"""TDD — Phase 14B server-side: OC-frame reassembly.

Rules encoded here:

  P14B-5  Server assembles a 1-chunk OC sequence → stored object matches original.

  P14B-6  Server assembles a 3-chunk OC sequence → stored object matches original.

  P14B-7  Server receives partial OC sequence (2 of 3 chunks) then END frame →
          returns an error frame; does not store a partial object.
"""
from __future__ import annotations
from collections.abc import AsyncIterator

from unittest.mock import AsyncMock, MagicMock

import msgpack
import pytest
from sqlalchemy.ext.asyncio import AsyncSession

from muse.core.types import blob_id, fake_id
from musehub.types.json_types import JSONObject, JSONValue
from muse.core.mpack import MuseWireFrameWriter
from musehub.models.wire import (
    SFRAME_HEADER, SFRAME_OBJECT_CHUNK, SFRAME_COMMIT_PACK, SFRAME_END,
    SFRAME_ERROR, SFRAME_RESULT,
)


# ---------------------------------------------------------------------------
# Shared helpers (same pattern as test_wire_push_stream.py)
# ---------------------------------------------------------------------------

_fw = MuseWireFrameWriter()


def _pack(data: JSONValue) -> bytes:
    return msgpack.packb(data, use_bin_type=True)


def _wrap(ft: str, data: JSONValue) -> bytes:
    return _fw.wrap(frame_type=ft, payload=_pack(data))


def _header_frame(n_objects: int = 1) -> bytes:
    return _wrap(SFRAME_HEADER, {
        "t": SFRAME_HEADER, "branch": "main", "force": False, "have": [],
        "head": fake_id("push-head"), "n_objects": n_objects, "n_commits": 0,
        "agent_id": "", "model_id": "", "sig": "",
    })


def _oc_frame(oid: str, ci: int, tc: int, chunk: bytes, *,
              enc: str | None = None, path: str | None = None,
              sz: int | None = None) -> bytes:
    payload: JSONObject = {"t": SFRAME_OBJECT_CHUNK, "id": oid,
                     "ci": ci, "tc": tc, "content": chunk}
    if enc is not None:
        payload["enc"] = enc
    if path is not None:
        payload["path"] = path
    if sz is not None:
        payload["sz"] = sz
    return _wrap(SFRAME_OBJECT_CHUNK, payload)


def _c_frame() -> bytes:
    return _wrap(SFRAME_COMMIT_PACK, {
        "t": SFRAME_COMMIT_PACK, "commits": [], "snapshots": [],
        "snapshot_deltas": [],
    })


def _e_frame(n_objects: int = 1, n_commits: int = 0) -> bytes:
    return _wrap(SFRAME_END, {"t": SFRAME_END, "n_objects": n_objects, "n_commits": n_commits})


async def _collect_frames(gen: AsyncIterator[bytes]) -> list[JSONObject]:
    results = []
    async for chunk in gen:
        unpacker = msgpack.Unpacker(raw=False)
        unpacker.feed(chunk)
        results.extend(list(unpacker))
    return results


# ---------------------------------------------------------------------------
# Shared fixtures
# ---------------------------------------------------------------------------

@pytest.fixture()
def stored_objects() -> None:
    return {}


@pytest.fixture()
def stub_backend(monkeypatch: pytest.MonkeyPatch, stored_objects: None) -> None:
    backend = AsyncMock()
    backend.exists = AsyncMock(side_effect=lambda oid, **_: oid in stored_objects)
    def _put(oid: str, data: bytes, **_: JSONValue) -> None:
        stored_objects[oid] = data
        return f"local://{oid}"
    backend.put = AsyncMock(side_effect=_put)
    backend.get = AsyncMock(side_effect=lambda oid, **_: stored_objects.get(oid))
    monkeypatch.setattr("musehub.services.musehub_wire.get_backend", lambda: backend)
    return backend


@pytest.fixture()
def stub_session() -> None:
    session = AsyncMock(spec=AsyncSession)
    session.execute = AsyncMock(
        return_value=MagicMock(scalar=lambda: None, fetchall=lambda: [])
    )
    session.commit = AsyncMock()
    session.add = MagicMock()
    return session


# ---------------------------------------------------------------------------
# P14B-5 — Server assembles 1-chunk OC sequence → stored correctly
# ---------------------------------------------------------------------------

class TestP14B5OneChunkOC:
    @pytest.mark.asyncio
    async def test_single_oc_chunk_stored_correctly(
        self, stub_backend: None, stub_session: None, stored_objects: None
    ) -> None:
        from musehub.services.musehub_wire import wire_push_stream

        content = b"hello world from a single OC chunk"
        oid = blob_id(content)

        async def body() -> None:
            yield (
                _header_frame(n_objects=1)
                + _oc_frame(oid, 0, 1, content, enc="raw", path="file.txt", sz=len(content))
                + _c_frame()
                + _e_frame()
            )

        await _collect_frames(wire_push_stream(stub_session, "repo-id", body(), "gabriel"))
        assert oid in stored_objects, f"Object {oid[:20]} not stored"
        assert stored_objects[oid] == content

    @pytest.mark.asyncio
    async def test_single_oc_no_error_frame(
        self, stub_backend: None, stub_session: None
    ) -> None:
        from musehub.services.musehub_wire import wire_push_stream

        content = b"no error expected"
        oid = blob_id(content)

        async def body() -> None:
            yield (
                _header_frame(n_objects=1)
                + _oc_frame(oid, 0, 1, content, enc="raw", path="", sz=len(content))
                + _c_frame()
                + _e_frame()
            )

        frames = await _collect_frames(wire_push_stream(stub_session, "repo-id", body(), "gabriel"))
        assert not [f for f in frames if f.get("t") == "X"]

    @pytest.mark.asyncio
    async def test_result_ok_after_single_oc(
        self, stub_backend: None, stub_session: None
    ) -> None:
        from musehub.services.musehub_wire import wire_push_stream

        content = b"result test"
        oid = blob_id(content)

        async def body() -> None:
            yield (
                _header_frame(n_objects=1)
                + _oc_frame(oid, 0, 1, content, enc="raw", path="", sz=len(content))
                + _c_frame()
                + _e_frame()
            )

        frames = await _collect_frames(wire_push_stream(stub_session, "repo-id", body(), "gabriel"))
        result_frames = [f for f in frames if f.get("t") == "R"]
        assert result_frames, "No RESULT frame"
        assert result_frames[0]["ok"] is True


# ---------------------------------------------------------------------------
# P14B-6 — Server assembles 3-chunk OC sequence → stored correctly
# ---------------------------------------------------------------------------

class TestP14B6ThreeChunkOC:
    def _content_and_chunks(self, n: int = 900) -> None:
        content = bytes([i % 256 for i in range(n)])
        oid = blob_id(content)
        size = n // 3
        chunks = [content[i * size:(i + 1) * size if i < 2 else n] for i in range(3)]
        return content, oid, chunks

    @pytest.mark.asyncio
    async def test_three_chunks_stored_correctly(
        self, stub_backend: None, stub_session: None, stored_objects: None
    ) -> None:
        from musehub.services.musehub_wire import wire_push_stream

        content, oid, chunks = self._content_and_chunks()

        async def body() -> None:
            b = _header_frame(n_objects=1)
            for i, chunk in enumerate(chunks):
                kw = {"enc": "raw", "path": "big.bin", "sz": len(content)} if i == 0 else {}
                b += _oc_frame(oid, i, 3, chunk, **kw)
            b += _c_frame() + _e_frame()
            yield b

        await _collect_frames(wire_push_stream(stub_session, "repo-id", body(), "gabriel"))
        assert oid in stored_objects
        assert stored_objects[oid] == content

    @pytest.mark.asyncio
    async def test_three_chunks_no_error_frame(
        self, stub_backend: None, stub_session: None
    ) -> None:
        from musehub.services.musehub_wire import wire_push_stream

        content, oid, chunks = self._content_and_chunks()

        async def body() -> None:
            b = _header_frame(n_objects=1)
            for i, chunk in enumerate(chunks):
                kw = {"enc": "raw", "path": "", "sz": len(content)} if i == 0 else {}
                b += _oc_frame(oid, i, 3, chunk, **kw)
            b += _c_frame() + _e_frame()
            yield b

        frames = await _collect_frames(wire_push_stream(stub_session, "repo-id", body(), "gabriel"))
        assert not [f for f in frames if f.get("t") == "X"]

    @pytest.mark.asyncio
    async def test_chunks_out_of_order_assembles_correctly(
        self, stub_backend: None, stub_session: None, stored_objects: None
    ) -> None:
        from musehub.services.musehub_wire import wire_push_stream

        content, oid, chunks = self._content_and_chunks()

        async def body() -> None:
            b = _header_frame(n_objects=1)
            # Send in reverse order
            for i in reversed(range(3)):
                kw = {"enc": "raw", "path": "", "sz": len(content)} if i == 0 else {}
                b += _oc_frame(oid, i, 3, chunks[i], **kw)
            b += _c_frame() + _e_frame()
            yield b

        await _collect_frames(wire_push_stream(stub_session, "repo-id", body(), "gabriel"))
        assert oid in stored_objects
        assert stored_objects[oid] == content


# ---------------------------------------------------------------------------
# P14B-7 — Partial OC sequence + END → error frame, no partial object stored
# ---------------------------------------------------------------------------

class TestP14B7PartialOCSequenceErrors:
    @pytest.mark.asyncio
    async def test_one_of_three_chunks_returns_error(
        self, stub_backend: None, stub_session: None, stored_objects: None
    ) -> None:
        from musehub.services.musehub_wire import wire_push_stream

        content = bytes(300)
        oid = blob_id(content)

        async def body() -> None:
            yield (
                _header_frame(n_objects=1)
                + _oc_frame(oid, 0, 3, content[:100], enc="raw", path="", sz=300)
                # chunks 1 and 2 missing
                + _c_frame()
                + _e_frame()
            )

        frames = await _collect_frames(wire_push_stream(stub_session, "repo-id", body(), "gabriel"))
        assert [f for f in frames if f.get("t") == "X"], "Expected error for incomplete OC"

    @pytest.mark.asyncio
    async def test_partial_oc_object_not_stored(
        self, stub_backend: None, stub_session: None, stored_objects: None
    ) -> None:
        from musehub.services.musehub_wire import wire_push_stream

        content = bytes(300)
        oid = blob_id(content)

        async def body() -> None:
            yield (
                _header_frame(n_objects=1)
                + _oc_frame(oid, 0, 3, content[:100], enc="raw", path="", sz=300)
                + _c_frame()
                + _e_frame()
            )

        await _collect_frames(wire_push_stream(stub_session, "repo-id", body(), "gabriel"))
        assert oid not in stored_objects, "Partial object must not be stored"

    @pytest.mark.asyncio
    async def test_two_of_three_chunks_returns_error(
        self, stub_backend: None, stub_session: None, stored_objects: None
    ) -> None:
        from musehub.services.musehub_wire import wire_push_stream

        content = bytes(300)
        oid = blob_id(content)

        async def body() -> None:
            yield (
                _header_frame(n_objects=1)
                + _oc_frame(oid, 0, 3, content[:100], enc="raw", path="", sz=300)
                + _oc_frame(oid, 1, 3, content[100:200])
                # chunk 2 missing
                + _c_frame()
                + _e_frame()
            )

        frames = await _collect_frames(wire_push_stream(stub_session, "repo-id", body(), "gabriel"))
        assert [f for f in frames if f.get("t") == "X"], "Expected error for missing last chunk"
        assert oid not in stored_objects