"""TDD — n_objects in H/E frames must equal distinct objects, not raw OC-entry count.

Bug: _frame_generator passed len(objects) to both H and E frame n_objects.
     When an object is split into N OC chunks, it contributes N entries to
     `objects` but only 1 to n_objects_received on the server side, which
     counts assembled OC sequences.  This produces a count-mismatch 400 on
     the first real push that contains any chunked (>OC_CHUNK_SIZE) object.

     Observed failure:
       "count mismatch: E frame claims 1029 object frames, server received 1028"
     Root cause: 1027 O-frames + 1 OC object × 2 chunks = 1029 entries in
     objects list, but server assembled 1027 O + 1 OC = 1028 distinct objects.

Rules:

  P-OC-N0  Empty objects list → H.n_objects=0, E.n_objects=0.

  P-OC-N1  Two plain O-objects → H.n_objects=2, E.n_objects=2.

  P-OC-N2  One OC object with 2 chunk entries in objects list →
           H.n_objects=1, E.n_objects=1.
           (Regression test: the broken code emits n_objects=2 here.)

  P-OC-N3  1 O-object + 1 OC-object (3 chunk entries) = 4 list entries →
           H.n_objects=2, E.n_objects=2.

  P-OC-N4  H.n_objects == E.n_objects for all inputs (consistency invariant).
"""
from __future__ import annotations

import struct

import msgpack
from muse.core._types import MsgpackDict, blob_id


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

def _sha256_oid(data: bytes) -> str:
    return blob_id(data)


def _make_o_object(content: bytes, path: str = "") -> MsgpackDict:
    """Plain O-frame object entry."""
    return {
        "object_id": _sha256_oid(content),
        "content": content,
        "encoding": "raw",
        "path": path,
    }


def _make_oc_chunks(content: bytes, path: str = "") -> list[dict]:
    """Two OC-frame entries for a single logical object (simulates Phase-14B chunking)."""
    oid = _sha256_oid(content)
    # Split into 2 halves — sizes don't matter, just that t="OC" with same oid.
    mid = max(1, len(content) // 2)
    chunk0 = content[:mid]
    chunk1 = content[mid:]
    return [
        {
            "t": "OC",
            "object_id": oid,
            "chunk_index": 0,
            "total_chunks": 2,
            "content": chunk0,
            "encoding": "raw",
            "path": path,
            "sz": len(content),
        },
        {
            "t": "OC",
            "object_id": oid,
            "chunk_index": 1,
            "total_chunks": 2,
            "content": chunk1,
        },
    ]


def _parse_mwp_frames(data: bytes) -> list[dict]:
    """Parse concatenated MWP-framed bytes into a list of decoded payload dicts."""
    frames: list[dict] = []
    pos = 0
    magic = b"muse"
    while pos < len(data):
        assert data[pos : pos + 4] == magic, f"bad magic at {pos}: {data[pos:pos+4]!r}"
        pos += 4
        pos += 1  # version byte
        header_len = struct.unpack_from(">I", data, pos)[0]
        pos += 4
        pos += header_len  # skip envelope header (ft/id/sz)
        payload_len = struct.unpack_from(">Q", data, pos)[0]
        pos += 8
        payload_bytes = data[pos : pos + payload_len]
        pos += payload_len
        frames.append(msgpack.unpackb(payload_bytes, raw=False))
    return frames


def _h_and_e(objects: list[MsgpackDict]) -> tuple[MsgpackDict, MsgpackDict]:
    """Run _frame_generator and return the decoded H and E frame payloads."""
    from muse.core.transport import _frame_generator

    raw = b"".join(_frame_generator(objects, [], [], local_head=None))
    frames = _parse_mwp_frames(raw)

    h_frame = next((f for f in frames if f.get("t") == "H"), None)
    e_frame = next((f for f in frames if f.get("t") == "E"), None)

    assert h_frame is not None, "no H frame in output"
    assert e_frame is not None, "no E frame in output"
    return h_frame, e_frame


# ---------------------------------------------------------------------------
# P-OC-N0 — empty objects list
# ---------------------------------------------------------------------------

class TestPOCN0Empty:
    def test_h_n_objects_is_zero(self) -> None:
        h, _ = _h_and_e([])
        assert h["n_objects"] == 0

    def test_e_n_objects_is_zero(self) -> None:
        _, e = _h_and_e([])
        assert e["n_objects"] == 0


# ---------------------------------------------------------------------------
# P-OC-N1 — two plain O-objects
# ---------------------------------------------------------------------------

class TestPOCN1TwoOObjects:
    def _objects(self) -> list[MsgpackDict]:
        return [
            _make_o_object(b"alpha" * 100),
            _make_o_object(b"beta" * 100),
        ]

    def test_h_n_objects_is_two(self) -> None:
        h, _ = _h_and_e(self._objects())
        assert h["n_objects"] == 2

    def test_e_n_objects_is_two(self) -> None:
        _, e = _h_and_e(self._objects())
        assert e["n_objects"] == 2


# ---------------------------------------------------------------------------
# P-OC-N2 — one OC object split into 2 chunks (regression for the bug)
# ---------------------------------------------------------------------------

class TestPOCN2OneOCObjectTwoChunks:
    """Core regression: 2 OC chunk entries represent 1 distinct object.

    The broken code emits n_objects=2 (len(objects)); the correct code emits 1.
    """

    def _objects(self) -> list[MsgpackDict]:
        return _make_oc_chunks(b"x" * 600_000)  # 2 entries, same oid

    def test_h_n_objects_is_one(self) -> None:
        h, _ = _h_and_e(self._objects())
        assert h["n_objects"] == 1, (
            f"H frame says n_objects={h['n_objects']}; expected 1. "
            "Two OC chunk entries for the same oid represent one distinct object."
        )

    def test_e_n_objects_is_one(self) -> None:
        _, e = _h_and_e(self._objects())
        assert e["n_objects"] == 1, (
            f"E frame says n_objects={e['n_objects']}; expected 1. "
            "Two OC chunk entries for the same oid represent one distinct object."
        )


# ---------------------------------------------------------------------------
# P-OC-N3 — 1 O + 1 OC(3 chunks) = 4 entries → 2 distinct objects
# ---------------------------------------------------------------------------

class TestPOCN3MixedObjects:
    def _objects(self) -> list[MsgpackDict]:
        o_obj = _make_o_object(b"plain" * 50)
        # Build a 3-chunk OC object manually (simpler than invoking split_object_into_oc_frames)
        oid = _sha256_oid(b"large" * 200)
        oc_chunks = [
            {
                "t": "OC",
                "object_id": oid,
                "chunk_index": i,
                "total_chunks": 3,
                "content": b"large" * 67,
                **({"encoding": "raw", "path": "big.bin", "sz": 1000} if i == 0 else {}),
            }
            for i in range(3)
        ]
        return [o_obj] + oc_chunks  # 4 entries, 2 distinct objects

    def test_h_n_objects_is_two(self) -> None:
        h, _ = _h_and_e(self._objects())
        assert h["n_objects"] == 2

    def test_e_n_objects_is_two(self) -> None:
        _, e = _h_and_e(self._objects())
        assert e["n_objects"] == 2


# ---------------------------------------------------------------------------
# P-OC-N4 — H.n_objects always equals E.n_objects
# ---------------------------------------------------------------------------

class TestPOCN4HEConsistency:
    def _cases(self) -> list[list]:
        return [
            [],
            [_make_o_object(b"a" * 100)],
            _make_oc_chunks(b"b" * 600_000),
            [_make_o_object(b"c" * 100)] + _make_oc_chunks(b"d" * 600_000),
        ]

    def test_h_equals_e_for_all_cases(self) -> None:
        for objects in self._cases():
            h, e = _h_and_e(objects)
            assert h["n_objects"] == e["n_objects"], (
                f"H.n_objects={h['n_objects']} != E.n_objects={e['n_objects']} "
                f"for objects list of length {len(objects)}"
            )