"""Tests for snapshot delta encoding in _frame_generator.

Verifies that _frame_generator:
  A. Sends the first snapshot as a full manifest
  B. Sends subsequent snapshots as deltas
  C. Deltas chain correctly (each references the previous snapshot)
  D. A stream with only one commit sends a full snapshot (no delta)
  E. write_commit_pack includes snapshot_deltas in msgpack output
  F. Applying every delta in stream order reconstructs all manifests
"""
from __future__ import annotations

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


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

def _parse_mwp_frames(data: bytes) -> list[dict]:
    """Parse concatenated MWP-framed bytes into 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}"
        pos += 4
        pos += 1  # version byte
        header_len = struct.unpack_from(">I", data, pos)[0]
        pos += 4 + header_len
        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 _parse_mwp_payload(mwp_frame: bytes) -> MsgpackDict:
    """Extract msgpack payload from a raw MWP frame (past the binary envelope)."""
    # MWP: b"muse" | 0x01 | uint32 header_len | msgpack(envelope) | uint64 payload_len | payload
    assert mwp_frame[:4] == b"muse"
    assert mwp_frame[4] == 0x01
    header_len = struct.unpack(">I", mwp_frame[5:9])[0]
    payload_offset = 9 + header_len + 8  # skip envelope + 8-byte payload_len
    payload = mwp_frame[payload_offset:]
    return msgpack.unpackb(payload, raw=False)


def _build_manifest(n_files: int) -> Manifest:
    return {f"src/module_{i}.py": f"sha256:{'a' * 60}{i:04d}" for i in range(n_files)}


def _make_commit(snapshot_id: str, i: int) -> MsgpackDict:
    return {
        "commit_id": f"sha256:{'c' * 60}{i:04d}",
        "snapshot_id": snapshot_id,
        "message": f"commit {i}",
        "branch": "dev",
        "author": "gabriel",
        "committed_at": "2026-04-23T00:00:00+00:00",
        "parent_commit_id": None,
        "agent_id": "",
        "model_id": "",
        "toolchain_id": "",
        "signer_key": "",
        "signature": "",
    }


def _make_snapshot(manifest: Manifest, i: int) -> MsgpackDict:
    snap_bytes = msgpack.packb(manifest, use_bin_type=True)
    snap_id = blob_id(snap_bytes)
    return {
        "snapshot_id": snap_id,
        "manifest": manifest,
        "directories": [],
        "created_at": "2026-04-23T00:00:00+00:00",
        "note": "",
    }


def _build_linear_history(n_commits: int, n_files: int) -> tuple[list, list, dict]:
    """Build a sequence of commits each changing 3 files in a manifest.

    Returns: (commits, snapshots, snap_by_id)
    """
    commits = []
    snapshots = []
    snap_by_id = {}

    manifest = _build_manifest(n_files)
    for i in range(n_commits):
        # Mutate 3 files per commit
        new_manifest = dict(manifest)
        for j in range(3):
            path = f"src/module_{(i * 3 + j) % n_files}.py"
            new_manifest[path] = f"sha256:{'b' * 60}{i * 3 + j:04d}"
        snap = _make_snapshot(new_manifest, i)
        commit = _make_commit(snap["snapshot_id"], i)
        commits.append(commit)
        snapshots.append(snap)
        snap_by_id[snap["snapshot_id"]] = snap
        manifest = new_manifest

    return commits, snapshots, snap_by_id


# ---------------------------------------------------------------------------
# A. First snapshot in stream is always a full manifest
# ---------------------------------------------------------------------------

def test_first_snapshot_is_full():
    from muse.core.transport import _frame_generator

    commits, snapshots, snap_by_id = _build_linear_history(n_commits=5, n_files=100)

    raw = b"".join(_frame_generator([], commits, snapshots, local_head=None))
    c_frames = [f for f in _parse_mwp_frames(raw) if f.get("t") == "C"]
    assert c_frames, "expected at least one C frame"
    assert len(c_frames[0].get("snapshots", [])) >= 1, "first C frame must contain at least one full snapshot"


# ---------------------------------------------------------------------------
# B. Subsequent snapshots sent as deltas
# ---------------------------------------------------------------------------

def test_subsequent_snapshots_are_deltas():
    from muse.core.transport import _frame_generator

    commits, snapshots, _ = _build_linear_history(n_commits=10, n_files=100)

    raw = b"".join(_frame_generator([], commits, snapshots, local_head=None))
    c_frames = [f for f in _parse_mwp_frames(raw) if f.get("t") == "C"]

    all_deltas = []
    for cf in c_frames:
        all_deltas.extend(cf.get("snapshot_deltas") or [])

    assert len(all_deltas) > 0, "expected delta snapshots for commits 2+"


# ---------------------------------------------------------------------------
# C. Delta chain integrity — each delta references a previously sent snapshot
# ---------------------------------------------------------------------------

def test_delta_chain_references_valid_base():
    from muse.core.transport import _frame_generator

    commits, snapshots, _ = _build_linear_history(n_commits=20, n_files=100)

    raw = b"".join(_frame_generator([], commits, snapshots, local_head=None))
    c_frames = [f for f in _parse_mwp_frames(raw) if f.get("t") == "C"]

    seen_snapshot_ids: set[str] = set()
    for cf in c_frames:
        for snap in cf.get("snapshots") or []:
            seen_snapshot_ids.add(snap["snapshot_id"])
        for delta in cf.get("snapshot_deltas") or []:
            base_id = delta.get("base_id")
            assert base_id in seen_snapshot_ids, (
                f"delta references base {base_id!r} not yet seen in stream"
            )
            seen_snapshot_ids.add(delta["snapshot_id"])


# ---------------------------------------------------------------------------
# D. Single commit — full snapshot, no deltas
# ---------------------------------------------------------------------------

def test_single_commit_sends_full_snapshot_no_deltas():
    from muse.core.transport import _frame_generator

    commits, snapshots, _ = _build_linear_history(n_commits=1, n_files=50)

    raw = b"".join(_frame_generator([], commits, snapshots, local_head=None))
    c_frames = [f for f in _parse_mwp_frames(raw) if f.get("t") == "C"]

    assert len(c_frames) == 1
    assert len(c_frames[0].get("snapshots", [])) == 1
    assert c_frames[0].get("snapshot_deltas") in (None, [])


# ---------------------------------------------------------------------------
# E. write_commit_pack includes snapshot_deltas key when provided
# ---------------------------------------------------------------------------

def test_write_commit_pack_includes_snapshot_deltas():
    from muse.core.mpack import MPackStreamWriter
    writer = MPackStreamWriter()

    delta = {
        "snapshot_id": "sha256:" + "a" * 64,
        "base_id": "sha256:" + "b" * 64,
        "added": {"src/foo.py": "sha256:" + "c" * 64},
        "removed": [],
        "directories": [],
        "created_at": "2026-04-23T00:00:00+00:00",
        "note": "",
    }
    raw = writer.write_commit_pack(commits=[], snapshots=[], snapshot_deltas=[delta])
    frame = msgpack.unpackb(raw, raw=False)
    assert "snapshot_deltas" in frame
    assert len(frame["snapshot_deltas"]) == 1
    assert frame["snapshot_deltas"][0]["base_id"] == "sha256:" + "b" * 64


def test_write_commit_pack_no_deltas_key_when_empty():
    from muse.core.mpack import MPackStreamWriter
    writer = MPackStreamWriter()
    raw = writer.write_commit_pack(commits=[], snapshots=[])
    frame = msgpack.unpackb(raw, raw=False)
    assert "snapshot_deltas" not in frame


# ---------------------------------------------------------------------------
# F. Applying every delta in stream order reconstructs all manifests
# ---------------------------------------------------------------------------

def test_delta_chain_full_reconstruction():
    """All manifests reconstructed by walking the delta chain must equal
    the original snapshots passed to _frame_generator."""
    from muse.core.transport import _frame_generator
    from muse.core.mpack import apply_snapshot_delta

    n_commits = 30
    commits, snapshots, snap_by_id = _build_linear_history(n_commits=n_commits, n_files=200)

    raw = b"".join(_frame_generator([], commits, snapshots, local_head=None))
    c_frames = [f for f in _parse_mwp_frames(raw) if f.get("t") == "C"]

    # Walk stream, reconstructing every snapshot
    reconstructed: dict[str, dict] = {}  # snapshot_id → manifest

    for cf in c_frames:
        for snap in cf.get("snapshots") or []:
            reconstructed[snap["snapshot_id"]] = snap["manifest"]
        for delta in cf.get("snapshot_deltas") or []:
            base_manifest = reconstructed[delta["base_id"]]
            full = apply_snapshot_delta(base_manifest, delta["added"], delta["removed"])
            reconstructed[delta["snapshot_id"]] = full

    for snap in snapshots:
        sid = snap["snapshot_id"]
        assert sid in reconstructed, f"snapshot {sid[:16]} not in reconstructed"
        assert reconstructed[sid] == snap["manifest"], (
            f"reconstructed manifest for {sid[:16]} does not match original"
        )