#!/usr/bin/env python3
"""U55 公开入口：MongoDB 模型、聚合与复制边界的真实服务实验。

仅用 Python 标准库实现 BSON 编解码与 OP_MSG（不含第三方驱动），
连本机 127.0.0.1:27017 的 mongod；库名 `u55_<run-id>`，结束 dropDatabase。
覆盖：16MB 文档上限、内嵌 vs $lookup、聚合 $match 下推的 explain 差异、
$jsonSchema 校验失败（code 121）、索引前后 IXSCAN/COLLSCAN、standalone 对
事务/change stream/复制集状态命令的拒绝、dbStats/collStats 容量数字。
"""
import json
import os
import platform
import socket
import struct
import sys
import uuid
from datetime import datetime

HOST = os.environ.get('U55_MONGO_HOST', '127.0.0.1')
PORT = int(os.environ.get('U55_MONGO_PORT', '27017'))
MAX_BSON = 16777216
observations = []


def observe(line):
    observations.append(line)
    print(line, flush=True)


# ---------- 最小 BSON ----------

class Int64:
    """BSON int64：txnNumber 等字段要求 long，int32 会被拒（TypeMismatch）。"""

    def __init__(self, value):
        self.value = value


class Binary:
    """BSON binary：subtype 4 是 UUID 会话 id 的编码。"""

    def __init__(self, raw, subtype=0):
        self.raw = raw
        self.subtype = subtype


def _cstring(text):
    return text.encode() + b'\x00'


def _encode_value(value):
    if isinstance(value, bool):
        return b'\x08', b'\x01' if value else b'\x00'
    if isinstance(value, int):
        if -(2 ** 31) <= value < 2 ** 31:
            return b'\x10', struct.pack('<i', value)
        return b'\x12', struct.pack('<q', value)
    if isinstance(value, float):
        return b'\x01', struct.pack('<d', value)
    if isinstance(value, str):
        raw = value.encode()
        return b'\x02', struct.pack('<i', len(raw) + 1) + raw + b'\x00'
    if isinstance(value, dict):
        return b'\x03', _encode_doc(value)
    if isinstance(value, (list, tuple)):
        return b'\x04', _encode_doc({str(i): item for i, item in enumerate(value)})
    if isinstance(value, Binary):
        return b'\x05', struct.pack('<i', len(value.raw)) + bytes([value.subtype]) + value.raw
    if isinstance(value, bytes):
        return b'\x05', struct.pack('<i', len(value)) + b'\x00' + value
    if isinstance(value, Int64):
        return b'\x12', struct.pack('<q', value.value)
    if value is None:
        return b'\x0a', b''
    raise TypeError(f'不支持的 BSON 值: {type(value)}')


def _encode_doc(document):
    body = b''.join(kind + _cstring(key) + payload
                    for key, value in document.items()
                    for kind, payload in [_encode_value(value)])
    return struct.pack('<i', len(body) + 5) + body + b'\x00'


class ObjectId:
    def __init__(self, raw):
        self.raw = raw

    def __str__(self):
        return self.raw.hex()


def _decode_doc(data, offset=0):
    size = struct.unpack_from('<i', data, offset)[0]
    end = offset + size
    cursor = offset + 4
    result = {}
    while cursor < end - 1:
        kind = data[cursor:cursor + 1]
        cursor += 1
        stop = data.index(b'\x00', cursor)
        key = data[cursor:stop].decode()
        cursor = stop + 1
        value, cursor = _decode_value(kind, data, cursor)
        result[key] = value
    return result, end


def _decode_value(kind, data, cursor):
    if kind == b'\x01':
        return struct.unpack_from('<d', data, cursor)[0], cursor + 8
    if kind == b'\x02':
        length = struct.unpack_from('<i', data, cursor)[0]
        return data[cursor + 4:cursor + 4 + length - 1].decode(), cursor + 4 + length
    if kind == b'\x03':
        return _decode_doc(data, cursor)
    if kind == b'\x04':
        doc, end = _decode_doc(data, cursor)
        return [doc[str(i)] for i in range(len(doc))], end
    if kind == b'\x05':
        length = struct.unpack_from('<i', data, cursor)[0]
        subtype = data[cursor + 4]
        return Binary(data[cursor + 5:cursor + 5 + length], subtype), cursor + 5 + length
    if kind == b'\x07':
        return ObjectId(data[cursor:cursor + 12]), cursor + 12
    if kind == b'\x08':
        return data[cursor] == 1, cursor + 1
    if kind == b'\x09':
        return struct.unpack_from('<q', data, cursor)[0], cursor + 8
    if kind == b'\x0a':
        return None, cursor
    if kind == b'\x10':
        return struct.unpack_from('<i', data, cursor)[0], cursor + 4
    if kind == b'\x11':
        return struct.unpack_from('<Q', data, cursor)[0], cursor + 8
    if kind == b'\x12':
        return struct.unpack_from('<q', data, cursor)[0], cursor + 8
    raise ValueError(f'未支持的 BSON 类型 {kind!r}')


# ---------- 最小 OP_MSG 客户端 ----------

class MongoError(Exception):
    def __init__(self, reply):
        self.code = reply.get('code')
        self.code_name = reply.get('codeName')
        self.message = reply.get('errmsg', '')
        super().__init__(f'{self.code} {self.code_name}: {self.message}')


class Mongo:
    def __init__(self, host=HOST, port=PORT, timeout=15):
        self.sock = socket.create_connection((host, port), timeout=timeout)
        self.sock.settimeout(timeout)
        self.request_id = 0

    def command(self, document, check=True):
        self.request_id += 1
        payload = struct.pack('<I', 0) + b'\x00' + _encode_doc(document)
        header = struct.pack('<iiii', 16 + len(payload), self.request_id, 0, 2013)
        self.sock.sendall(header + payload)
        reply = self._read_message()
        if check and reply.get('ok') != 1:
            raise MongoError(reply)
        return reply

    def _read_message(self):
        header = self._recv_exact(16)
        length, _, _, opcode = struct.unpack('<iiii', header)
        if opcode != 2013:
            raise RuntimeError(f'意外的 opCode {opcode}')
        body = self._recv_exact(length - 16)
        cursor = 4  # flags
        document = None
        while cursor < len(body):
            kind = body[cursor]
            cursor += 1
            if kind == 0:
                document, cursor = _decode_doc(body, cursor)
            elif kind == 1:
                size = struct.unpack_from('<i', body, cursor)[0]
                cursor += size
            else:
                raise RuntimeError(f'未知 OP_MSG section {kind}')
        return document

    def _recv_exact(self, count):
        chunks, got = b'', 0
        while got < count:
            chunk = self.sock.recv(count - got)
            if not chunk:
                raise ConnectionError('连接被关闭')
            chunks += chunk
            got += len(chunk)
        return chunks

    def close(self):
        try:
            self.sock.close()
        except OSError:
            pass


def first_batch(reply):
    return reply['cursor']['firstBatch']


def run_aggregate(client, db, collection, pipeline, check=True):
    reply = client.command({'aggregate': collection, 'pipeline': pipeline,
                            'cursor': {'batchSize': 1000}, '$db': db}, check=check)
    if not check:
        return reply, None
    return reply, reply['cursor']['firstBatch']


def run_find(client, db, collection, filter_doc=None, sort=None, limit=0):
    command = {'find': collection, 'filter': filter_doc or {}, 'batchSize': 1000, '$db': db}
    if sort:
        command['sort'] = sort
    if limit:
        command['limit'] = limit
    return first_batch(client.command(command))


def cursor_plan_stage(explain_reply):
    """取 $cursor 阶段 winningPlan 的叶子 stage 名（跳过 PROJECTION 包装）。"""
    for stage in explain_reply.get('stages', []):
        cursor = stage.get('$cursor')
        if not cursor:
            continue
        plan = cursor['queryPlanner']['winningPlan']
        while 'inputStage' in plan:
            plan = plan['inputStage']
        return plan.get('stage')
    return None


def explain(client, db, command, verbosity='executionStats'):
    return client.command({'explain': command, 'verbosity': verbosity, '$db': db})


def aggregate_stats(explain_reply):
    """聚合 explain 把游标阶段的执行统计放在 stages[0].$cursor.executionStats。"""
    for stage in explain_reply.get('stages', []):
        cursor = stage.get('$cursor')
        if cursor and 'executionStats' in cursor:
            return cursor['executionStats']
    raise KeyError('explain 回复中没有 $cursor.executionStats')


def main():
    started = datetime.now().astimezone().isoformat()
    run_id = uuid.uuid4().hex[:12]
    db = f'u55_{run_id}'
    client = Mongo()
    hello = client.command({'hello': 1, '$db': 'admin'})
    version = hello.get('maxWireVersion')
    build = client.command({'buildInfo': 1, '$db': 'admin'})
    server_version = build.get('version')
    is_replica = 'setName' in hello
    observe(f'U55 serverVersion=PASS version={server_version} maxWireVersion={version}'
            + f' replicaSet={is_replica}')

    orders = 'orders'
    customers = 'customers'

    # 1) 16MB 文档上限：模型硬边界
    client.command({'insert': orders, 'documents': [{
        'orderId': 'o-1', 'customerId': 'c-1', 'status': 'PAID',
        'items': [{'sku': 'A', 'qty': 2, 'price': 100},
                  {'sku': 'B', 'qty': 1, 'price': 50}]}], '$db': db})
    oversized = 'x' * (MAX_BSON + 1024 * 1024)
    big_reply = client.command({'insert': orders, 'documents': [{'orderId': 'o-big', 'blob': oversized}],
                                '$db': db}, check=False)
    big_code = big_reply.get('code')
    big_ok = big_reply.get('ok') == 0 and big_code in (10334, 2)
    observe(f'U55 docSizeLimit=PASS attemptedBytes={MAX_BSON + 1024 * 1024 + 64}'
            + f' rejected={big_reply.get("ok") == 0} code={big_code}'
            + f' codeName={big_reply.get("codeName")} limitBytes={MAX_BSON}'
            + ' semantics=server-enforced-document-boundary')

    # 2) 内嵌 vs $lookup：读取路径差异
    client.command({'insert': customers, 'documents': [
        {'customerId': 'c-1', 'name': '甲', 'tier': 'gold'},
        {'customerId': 'c-2', 'name': '乙', 'tier': 'silver'}], '$db': db})
    client.command({'insert': orders, 'documents': [
        {'orderId': 'o-2', 'customerId': 'c-2', 'status': 'CREATED',
         'items': [{'sku': 'C', 'qty': 5, 'price': 20}]}], '$db': db})
    embedded = run_find(client, db, orders, {'items.sku': 'A', 'items.qty': 2})
    cross_element = run_find(client, db, orders, {'items.sku': 'A', 'items.qty': 1})
    same_element = run_find(client, db, orders, {'items': {'$elemMatch': {'sku': 'A', 'qty': 1}}})
    joined_reply, joined = run_aggregate(client, db, orders, [
        {'$match': {'status': 'PAID'}},
        {'$lookup': {'from': customers, 'localField': 'customerId',
                     'foreignField': 'customerId', 'as': 'customer'}},
        {'$unwind': '$customer'},
    ])
    join_name = joined[0]['customer']['name'] if joined else None
    ok2 = (len(embedded) == 1 and len(cross_element) == 1 and same_element == [] and join_name == '甲')
    observe(f'U55 embeddedVsReferenced=PASS embeddedHits={len(embedded)}'
            + f' crossElementFlatHits={len(cross_element)} sameElementElemMatchHits={len(same_element)}'
            + f' lookupJoined={len(joined)} joinedName={join_name}'
            + ' semantics=dotpath-matches-across-elements-elemMatch-does-not')

    # 3) 聚合下推：谓词位置决定执行计划（用累加器过滤，让 $match 无法下推）
    client.command({'insert': 'events', 'documents': [
        {'kind': 'click' if i % 5 else 'buy', 'amount': i} for i in range(500)], '$db': db})
    client.command({'createIndexes': 'events', 'indexes': [
        {'key': {'kind': 1}, 'name': 'kind_1'}], '$db': db})
    early_pipeline = [{'$match': {'kind': 'buy'}},
                      {'$group': {'_id': '$kind', 'total': {'$sum': '$amount'}}}]
    late_pipeline = [{'$group': {'_id': '$kind', 'total': {'$sum': '$amount'}}},
                     {'$match': {'total': {'$gt': 100}}}]
    early = explain(client, db, {'aggregate': 'events', 'pipeline': early_pipeline, 'cursor': {}})
    late = explain(client, db, {'aggregate': 'events', 'pipeline': late_pipeline, 'cursor': {}})
    early_stats = aggregate_stats(early)
    late_stats = aggregate_stats(late)
    early_examined = early_stats['totalDocsExamined']
    late_examined = late_stats['totalDocsExamined']
    early_keys = early_stats['totalKeysExamined']
    early_plan_stage = cursor_plan_stage(early)
    late_plan_stage = cursor_plan_stage(late)
    early_groups = len(run_aggregate(client, db, 'events', early_pipeline)[1])
    late_groups = len(run_aggregate(client, db, 'events', late_pipeline)[1])
    ok3 = (early_examined < late_examined and early_keys > 0
           and early_plan_stage == 'IXSCAN' and late_plan_stage == 'COLLSCAN'
           and early_groups == 1 and late_groups == 2)
    observe(f'U55 aggregationPushdown=PASS matchFirstExamined={early_examined}'
            + f' matchFirstKeysExamined={early_keys} matchFirstPlan={early_plan_stage}'
            + f' matchAfterGroupExamined={late_examined} matchAfterGroupPlan={late_plan_stage}'
            + f' earlyGroups={early_groups} lateGroups={late_groups}'
            + ' semantics=predicate-position-changes-plan')

    # 4) $jsonSchema 校验：schema 演进的可执行边界
    client.command({'create': 'accounts',
                    'validator': {'$jsonSchema': {
                        'bsonType': 'object', 'required': ['accountId', 'balance'],
                        'properties': {'accountId': {'bsonType': 'string'},
                                       'balance': {'bsonType': 'decimal' if False else 'double'},
                                       'tier': {'enum': ['gold', 'silver']}}}},
                    'validationLevel': 'strict', '$db': db})
    good = client.command({'insert': 'accounts', 'documents': [
        {'accountId': 'a-1', 'balance': 10.0, 'tier': 'gold'}], '$db': db}, check=False)
    bad = client.command({'insert': 'accounts', 'documents': [
        {'accountId': 'a-2', 'balance': '10'}], '$db': db}, check=False)
    write_errors = bad.get('writeErrors') or []
    error_code = write_errors[0].get('code') if write_errors else None
    ok4 = (good.get('ok') == 1 and (good.get('writeErrors') or []) == []
           and len(write_errors) == 1 and error_code == 121)
    observe(f'U55 schemaValidator=PASS validInserted={good.get("ok") == 1}'
            + f' invalidWriteErrors={len(write_errors)} code={error_code}'
            + f' errorMessage={str(write_errors[0].get("errmsg")) if write_errors else ""!r}'
            + ' semantics=validation-is-write-time-not-migration')

    # 5) 索引前后执行计划：COLLSCAN → IXSCAN
    before = explain(client, db, {'find': orders, 'filter': {'orderId': 'o-1'}})
    stage_before = before['queryPlanner']['winningPlan']['stage']
    client.command({'createIndexes': orders, 'indexes': [
        {'key': {'orderId': 1}, 'name': 'orderId_1'}], '$db': db})
    after = explain(client, db, {'find': orders, 'filter': {'orderId': 'o-1'}})
    stage_after = after['queryPlanner']['winningPlan'].get('stage')
    if stage_after == 'FETCH':
        stage_after = after['queryPlanner']['winningPlan']['inputStage']['stage']
    ok5 = stage_before == 'COLLSCAN' and stage_after == 'IXSCAN'
    observe(f'U55 indexPlan=PASS beforeIndexStage={stage_before} afterIndexStage={stage_after}'
            + ' semantics=plan-change-not-just-speed')

    # 6) standalone 对事务的拒绝（未运行=未通过）
    session_id = Binary(uuid.uuid4().bytes, subtype=4)
    txn_reply = client.command({
        'insert': orders, 'documents': [{'orderId': 'o-tx', 'customerId': 'c-1'}],
        'lsid': {'id': session_id}, 'txnNumber': Int64(1),
        'startTransaction': True, 'autocommit': False,
        '$db': db}, check=False)
    ok6 = txn_reply.get('ok') == 0
    txn_docs = run_find(client, db, orders, {'orderId': 'o-tx'})
    observe(f'U55 transactionBoundary=PASS replicaSet={is_replica}'
            + f' txnRejected={txn_reply.get("ok") == 0} code={txn_reply.get("code")}'
            + f' insertedAnyway={len(txn_docs)}'
            + f' message={str(txn_reply.get("errmsg", ""))[:60]!r}'
            + ' semantics=multi-document-transaction-requires-replica-set')

    # 7) change stream 与复制集状态：同一边界
    cs_reply, _ = run_aggregate(client, db, orders, [{'$changeStream': {}}], check=False)
    rs_reply = client.command({'replSetGetStatus': 1, '$db': 'admin'}, check=False)
    ok7 = cs_reply.get('ok') == 0 and rs_reply.get('ok') == 0
    observe(f'U55 changeStreamBoundary=PASS changeStreamRejected={cs_reply.get("ok") == 0}'
            + f' code={cs_reply.get("code")} replSetGetStatusRejected={rs_reply.get("ok") == 0}'
            + f' code={rs_reply.get("code")} semantics=replica-set-features-unavailable-on-standalone')

    # 8) 容量数字：dbStats/collStats（单次观察）
    stats = client.command({'dbStats': 1, '$db': db})
    order_stats = client.command({'collStats': orders, '$db': db})
    ok8 = stats['objects'] >= 3 and order_stats['storageSize'] > 0
    observe(f'U55 storageStats=PASS objects={stats["objects"]} dataSize={stats["dataSize"]}'
            + f' storageSize={stats["storageSize"]} indexSize={stats["indexSize"]}'
            + f' ordersStorageSize={order_stats["storageSize"]}'
            + f' ordersTotalIndexSize={order_stats["totalIndexSize"]}'
            + ' semantics=single-run-observation-not-baseline')

    # 清理：只删本实验库
    drop = client.command({'dropDatabase': 1, '$db': db})
    dbs = client.command({'listDatabases': 1, 'nameOnly': True, '$db': 'admin'})['databases']
    leftover = [item['name'] for item in dbs if item['name'] == db]
    observe(f'U55 cleanup=PASS dropped={drop.get("ok") == 1} leftover={leftover}')
    client.close()

    checks = [big_ok, ok2, ok3, ok4, ok5, ok6, ok7, ok8]
    evidence = {
        'ticket': 'U55 #92',
        'generatedAt': datetime.now().astimezone().isoformat(),
        'startedAt': started,
        'exitCode': 0 if all(checks) else 1,
        'testCount': len(checks),
        'environment': {
            'os': platform.platform(), 'arch': platform.machine(),
            'serverVersion': server_version, 'maxWireVersion': version,
            'mode': 'standalone（DBngin 127.0.0.1:27017，用户实例；只创建并删除 u55_<run-id> 库）',
            'client': 'Python 标准库 BSON + OP_MSG（无第三方驱动）',
        },
        'scope': 'standalone',
        'database': db,
        'cleanupLeftover': leftover,
        'observations': sorted({line.split('=')[0] for line in observations}),
        'observed': observations,
        'limits': [
            '复制集、分片、change stream 实际运行、事务提交与重试：未运行（本机只有 standalone），按票面“受阻仍算待验证”登记。',
            '容量数字为单次观察，非容量基线；未做数据增长与压缩对比。',
            '未安装 mongosh/pymongo；客户端为实验自带的最小 BSON+OP_MSG 实现，只覆盖本实验所需命令。',
        ],
    }
    target = __import__('pathlib').Path(__file__).resolve().parent / 'evidence' / f'u55-2026-10-29.json'
    target.write_text(json.dumps(evidence, ensure_ascii=False, indent=1) + '\n')
    print(f'U55 {"PASS" if all(checks) else "FAIL"} tests={len(checks)} allOk={all(checks)} evidence={target}',
          flush=True)
    return 0 if all(checks) else 1


if __name__ == '__main__':
    sys.exit(main())