"""TDD — shallow fetch: depth-limited BFS in wire_fetch_stream.

Test plan
---------
S1  WireFetchRequest accepts an optional depth field (schema contract).
S2  depth=None (default) returns full commit history — N commits in, N out.
S3  depth=1 returns only the tip commit, not its ancestors.
S4  depth=2 returns tip + one ancestor, stops before the root.
S5  H frame carries shallow_commits list when depth is active.
S6  Objects in the response are only those referenced by depth-limited commits.
S7  depth=1 on a single-commit repo behaves identically to depth=None.
S8  depth cap works correctly on a branching (multi-parent) DAG.
"""
from __future__ import annotations
from collections.abc import AsyncIterator

import asyncio
import zlib
from datetime import datetime, timezone

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

from muse.core.types import blob_id, fake_id
from musehub.db import musehub_models as db
from musehub.models.wire import (
    SFRAME_COMMIT_PACK,
    SFRAME_END,
    SFRAME_HEADER,
    SFRAME_OBJECT,
    WireFetchRequest,
)
from musehub.types.json_types import JSONObject
from tests.factories import create_repo


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

def _uid(seed: str) -> str:
    return fake_id(seed)


def _now() -> datetime:
    return datetime.now(tz=timezone.utc)


async def _make_snapshot(
    session: AsyncSession,
    repo_id: str,
    manifest: dict[str, str],
) -> db.MusehubSnapshot:
    sid = _uid(str(manifest))
    snap = db.MusehubSnapshot(
        snapshot_id=sid,
        repo_id=repo_id,
        directories=[],
        manifest_blob=msgpack.packb(manifest, use_bin_type=True),
        entry_count=len(manifest),
        created_at=_now(),
    )
    session.add(snap)
    await session.commit()
    return snap


async def _make_commit(
    session: AsyncSession,
    repo_id: str,
    *,
    parent_ids: list[str] | None = None,
    snapshot_id: str | None = None,
    seed: str = "",
) -> db.MusehubCommit:
    commit_id = _uid(f"commit-{seed}-{id(session)}")
    row = db.MusehubCommit(
        commit_id=commit_id,
        repo_id=repo_id,
        branch="main",
        parent_ids=parent_ids or [],
        message=f"commit {seed}",
        author="test",
        timestamp=_now(),
        snapshot_id=snapshot_id,
    )
    session.add(row)
    await session.commit()
    return row


async def _make_object(
    session: AsyncSession,
    repo_id: str,
    content: bytes,
) -> str:
    """Store an object in musehub_objects + musehub_object_refs. Returns object_id."""
    from musehub.services.musehub_wire import get_backend
    backend = get_backend()
    oid = blob_id(content)
    uri = await backend.put(oid, content)
    row = db.MusehubObject(
        object_id=oid,
        path="test.dat",
        size_bytes=len(content),
        disk_path=uri.replace("local://", ""),
        storage_uri=uri,
    )
    from sqlalchemy.dialects.postgresql import insert as pg_insert
    await session.execute(
        pg_insert(db.MusehubObject)
        .values(
            object_id=oid,
            path="test.dat",
            size_bytes=len(content),
            disk_path=uri.replace("local://", ""),
            storage_uri=uri,
        )
        .on_conflict_do_nothing(index_elements=["object_id"])
    )
    await session.execute(
        pg_insert(db.MusehubObjectRef)
        .values(repo_id=repo_id, object_id=oid)
        .on_conflict_do_nothing()
    )
    await session.commit()
    return oid


async def _collect_frames(gen: AsyncIterator[bytes]) -> list[JSONObject]:
    frames = []
    async for chunk in gen:
        for frame in _unpack_chunks([chunk]):
            frames.append(frame)
    return frames


def _unpack_chunks(chunks: list[bytes]) -> list[dict]:
    unpacker = msgpack.Unpacker(raw=False)
    for chunk in chunks:
        unpacker.feed(chunk)
    return list(unpacker)


# ---------------------------------------------------------------------------
# S1 — WireFetchRequest accepts depth
# ---------------------------------------------------------------------------

class TestWireFetchRequestSchema:
    def test_depth_defaults_to_none(self) -> None:
        req = WireFetchRequest(want=[], have=[])
        assert req.depth is None

    def test_depth_accepts_positive_int(self) -> None:
        req = WireFetchRequest(want=[], have=[], depth=1)
        assert req.depth == 1

    def test_depth_accepts_none_explicitly(self) -> None:
        req = WireFetchRequest(want=[], have=[], depth=None)
        assert req.depth is None


# ---------------------------------------------------------------------------
# S2 — depth=None returns full history
# ---------------------------------------------------------------------------

class TestFullHistory:
    @pytest.mark.asyncio
    async def test_no_depth_returns_all_commits(self, db_session: AsyncSession) -> None:
        repo = await create_repo(db_session, owner="test-user-wire")
        repo_id = str(repo.repo_id)

        c1 = await _make_commit(db_session, repo_id, seed="root")
        c2 = await _make_commit(db_session, repo_id, parent_ids=[c1.commit_id], seed="middle")
        c3 = await _make_commit(db_session, repo_id, parent_ids=[c2.commit_id], seed="tip")

        from musehub.services.musehub_wire import wire_fetch_stream
        req = WireFetchRequest(want=[c3.commit_id], have=[], depth=None)
        frames = await _collect_frames(wire_fetch_stream(db_session, repo_id, req))

        c_frames = [f for f in frames if f.get("t") == SFRAME_COMMIT_PACK]
        all_commit_ids = {c["commit_id"] for cf in c_frames for c in cf.get("commits", [])}
        assert c1.commit_id in all_commit_ids
        assert c2.commit_id in all_commit_ids
        assert c3.commit_id in all_commit_ids


# ---------------------------------------------------------------------------
# S3 — depth=1 returns only the tip
# ---------------------------------------------------------------------------

class TestDepthOne:
    @pytest.mark.asyncio
    async def test_depth_1_returns_only_tip(self, db_session: AsyncSession) -> None:
        repo = await create_repo(db_session, owner="test-user-wire")
        repo_id = str(repo.repo_id)

        c1 = await _make_commit(db_session, repo_id, seed="root")
        c2 = await _make_commit(db_session, repo_id, parent_ids=[c1.commit_id], seed="middle")
        c3 = await _make_commit(db_session, repo_id, parent_ids=[c2.commit_id], seed="tip")

        from musehub.services.musehub_wire import wire_fetch_stream
        req = WireFetchRequest(want=[c3.commit_id], have=[], depth=1)
        frames = await _collect_frames(wire_fetch_stream(db_session, repo_id, req))

        c_frames = [f for f in frames if f.get("t") == SFRAME_COMMIT_PACK]
        all_commit_ids = {c["commit_id"] for cf in c_frames for c in cf.get("commits", [])}
        assert c3.commit_id in all_commit_ids
        assert c2.commit_id not in all_commit_ids
        assert c1.commit_id not in all_commit_ids

    @pytest.mark.asyncio
    async def test_depth_1_stream_ends_cleanly(self, db_session: AsyncSession) -> None:
        repo = await create_repo(db_session, owner="test-user-wire")
        repo_id = str(repo.repo_id)

        c1 = await _make_commit(db_session, repo_id, seed="a")
        c2 = await _make_commit(db_session, repo_id, parent_ids=[c1.commit_id], seed="b")

        from musehub.services.musehub_wire import wire_fetch_stream
        req = WireFetchRequest(want=[c2.commit_id], have=[], depth=1)
        frames = await _collect_frames(wire_fetch_stream(db_session, repo_id, req))

        types = [f.get("t") for f in frames]
        assert SFRAME_END in types


# ---------------------------------------------------------------------------
# S4 — depth=2 returns tip + one ancestor
# ---------------------------------------------------------------------------

class TestDepthTwo:
    @pytest.mark.asyncio
    async def test_depth_2_includes_parent_excludes_grandparent(
        self, db_session: AsyncSession
    ) -> None:
        repo = await create_repo(db_session, owner="test-user-wire")
        repo_id = str(repo.repo_id)

        c1 = await _make_commit(db_session, repo_id, seed="root")
        c2 = await _make_commit(db_session, repo_id, parent_ids=[c1.commit_id], seed="parent")
        c3 = await _make_commit(db_session, repo_id, parent_ids=[c2.commit_id], seed="tip")

        from musehub.services.musehub_wire import wire_fetch_stream
        req = WireFetchRequest(want=[c3.commit_id], have=[], depth=2)
        frames = await _collect_frames(wire_fetch_stream(db_session, repo_id, req))

        c_frames = [f for f in frames if f.get("t") == SFRAME_COMMIT_PACK]
        all_commit_ids = {c["commit_id"] for cf in c_frames for c in cf.get("commits", [])}
        assert c3.commit_id in all_commit_ids
        assert c2.commit_id in all_commit_ids
        assert c1.commit_id not in all_commit_ids


# ---------------------------------------------------------------------------
# S5 — H frame carries shallow_commits when depth is active
# ---------------------------------------------------------------------------

class TestShallowHeader:
    @pytest.mark.asyncio
    async def test_h_frame_has_shallow_commits_on_depth(
        self, db_session: AsyncSession
    ) -> None:
        repo = await create_repo(db_session, owner="test-user-wire")
        repo_id = str(repo.repo_id)

        c1 = await _make_commit(db_session, repo_id, seed="root")
        c2 = await _make_commit(db_session, repo_id, parent_ids=[c1.commit_id], seed="tip")

        from musehub.services.musehub_wire import wire_fetch_stream
        req = WireFetchRequest(want=[c2.commit_id], have=[], depth=1)
        frames = await _collect_frames(wire_fetch_stream(db_session, repo_id, req))

        h = next(f for f in frames if f.get("t") == SFRAME_HEADER)
        assert "shallow_commits" in h
        assert c2.commit_id in h["shallow_commits"]

    @pytest.mark.asyncio
    async def test_h_frame_no_shallow_commits_without_depth(
        self, db_session: AsyncSession
    ) -> None:
        repo = await create_repo(db_session, owner="test-user-wire")
        repo_id = str(repo.repo_id)

        c1 = await _make_commit(db_session, repo_id, seed="only")

        from musehub.services.musehub_wire import wire_fetch_stream
        req = WireFetchRequest(want=[c1.commit_id], have=[], depth=None)
        frames = await _collect_frames(wire_fetch_stream(db_session, repo_id, req))

        h = next(f for f in frames if f.get("t") == SFRAME_HEADER)
        assert not h.get("shallow_commits")


# ---------------------------------------------------------------------------
# S6 — Only objects from depth-limited commits are sent
# ---------------------------------------------------------------------------

class TestShallowObjects:
    @pytest.mark.asyncio
    async def test_depth_1_only_sends_tip_objects(self, db_session: AsyncSession) -> None:
        repo = await create_repo(db_session, owner="test-user-wire")
        repo_id = str(repo.repo_id)

        obj_root = await _make_object(db_session, repo_id, b"root-content")
        obj_tip  = await _make_object(db_session, repo_id, b"tip-content")

        snap_root = await _make_snapshot(db_session, repo_id, {"root.dat": obj_root})
        snap_tip  = await _make_snapshot(db_session, repo_id, {"root.dat": obj_root, "tip.dat": obj_tip})

        c1 = await _make_commit(db_session, repo_id, snapshot_id=snap_root.snapshot_id, seed="root")
        c2 = await _make_commit(
            db_session, repo_id, parent_ids=[c1.commit_id],
            snapshot_id=snap_tip.snapshot_id, seed="tip",
        )

        from musehub.services.musehub_wire import wire_fetch_stream
        req = WireFetchRequest(want=[c2.commit_id], have=[], depth=1)
        frames = await _collect_frames(wire_fetch_stream(db_session, repo_id, req))

        o_frames = [f for f in frames if f.get("t") == SFRAME_OBJECT]
        received_oids = {f["id"] for f in o_frames}

        # Both objects are in tip's snapshot — both should be sent
        assert obj_tip in received_oids
        assert obj_root in received_oids


# ---------------------------------------------------------------------------
# S7 — depth=1 on single-commit repo == full clone
# ---------------------------------------------------------------------------

class TestDepthOnSingleCommit:
    @pytest.mark.asyncio
    async def test_single_commit_depth_1_equals_full(
        self, db_session: AsyncSession
    ) -> None:
        repo = await create_repo(db_session, owner="test-user-wire")
        repo_id = str(repo.repo_id)

        c1 = await _make_commit(db_session, repo_id, seed="only")

        from musehub.services.musehub_wire import wire_fetch_stream

        full = await _collect_frames(
            wire_fetch_stream(db_session, repo_id, WireFetchRequest(want=[c1.commit_id], have=[], depth=None))
        )
        shallow = await _collect_frames(
            wire_fetch_stream(db_session, repo_id, WireFetchRequest(want=[c1.commit_id], have=[], depth=1))
        )

        def _commit_ids(frames: list[JSONObject]) -> set[str]:
            return {c["commit_id"] for f in frames if f.get("t") == SFRAME_COMMIT_PACK for c in f.get("commits", [])}

        assert _commit_ids(full) == _commit_ids(shallow) == {c1.commit_id}


# ---------------------------------------------------------------------------
# S8 — depth on a DAG with two parents (merge commit)
# ---------------------------------------------------------------------------

class TestShallowDAG:
    @pytest.mark.asyncio
    async def test_depth_1_on_merge_commit_returns_only_merge(
        self, db_session: AsyncSession
    ) -> None:
        repo = await create_repo(db_session, owner="test-user-wire")
        repo_id = str(repo.repo_id)

        base   = await _make_commit(db_session, repo_id, seed="base")
        feat   = await _make_commit(db_session, repo_id, parent_ids=[base.commit_id], seed="feat")
        merge  = await _make_commit(
            db_session, repo_id,
            parent_ids=[base.commit_id, feat.commit_id],
            seed="merge",
        )

        from musehub.services.musehub_wire import wire_fetch_stream
        req = WireFetchRequest(want=[merge.commit_id], have=[], depth=1)
        frames = await _collect_frames(wire_fetch_stream(db_session, repo_id, req))

        c_frames = [f for f in frames if f.get("t") == SFRAME_COMMIT_PACK]
        all_ids = {c["commit_id"] for cf in c_frames for c in cf.get("commits", [])}
        assert merge.commit_id in all_ids
        assert base.commit_id not in all_ids
        assert feat.commit_id not in all_ids