"""TDD — fetch compression: zstd + path-sorted objects + delta frames.

Test plan
---------
Layer 1 — zstd
C1  Server uses zstd encoding (not zlib) for O frames when zstd is available.
C2  Client can decode the zstd frame back to original bytes.

Layer 2 — path-sorted objects
C3  O frames for a multi-object fetch arrive sorted by path within each batch.
    Sorted order is a prerequisite for profitable intra-batch compression and
    deterministic delta base selection.

Layer 3 — delta frames
C4  Two successive versions of the same file → server sends the second as a
    delta frame (enc="delta+zstd") with a non-empty base_id.
C5  Delta frame content reconstructs to the correct target bytes.
C6  Delta is not sent when unprofitable (random content → no shared runs →
    delta ≥ full object; server falls back to full zstd frame).
C7  Full clone with a mix of full + delta frames delivers the right object_id
    set to the client.
C8  Objects from a prior batch serve as delta bases for the next batch
    (cross-batch delta).
"""
from __future__ import annotations

import zlib
from datetime import datetime, timezone

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

from muse.core.compression import ZSTD_AVAILABLE, apply_delta, compute_delta
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 tests.factories import create_repo


# ---------------------------------------------------------------------------
# Helpers (mirror test_wire_batch_stream.py)
# ---------------------------------------------------------------------------

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


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


async def _make_object(
    session: AsyncSession,
    repo_id: str,
    content: bytes,
    path: str = "file.dat",
    *,
    owner: str = "",
    slug: str = "",
) -> str:
    from musehub.services.musehub_wire import get_backend
    from musehub.storage.backends import repo_root_for
    from sqlalchemy.dialects.postgresql import insert as pg_insert

    backend = get_backend(owner or None, slug or None)
    oid = blob_id(content)
    # Use per-repo root so wire_fetch_stream (which passes repo_root=repo_root_for(...))
    # reads from the same location we write to.
    per_repo_root = repo_root_for(owner, slug) if (owner and slug) else None
    uri = await backend.put(oid, content, repo_root=per_repo_root)
    await session.execute(
        pg_insert(db.MusehubObject)
        .values(
            object_id=oid,
            path=path,
            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 _make_snapshot(
    session: AsyncSession,
    repo_id: str,
    manifest: dict[str, str],
) -> db.MusehubSnapshot:
    sid = _uid(str(sorted(manifest.items())))
    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:
    row = db.MusehubCommit(
        commit_id=_uid(f"commit-{seed}"),
        repo_id=repo_id,
        branch="main",
        parent_ids=parent_ids or [],
        message=f"commit {seed}",
        author="gabriel",
        timestamp=_now(),
        snapshot_id=snapshot_id,
    )
    session.add(row)
    await session.commit()
    return row


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


def _o_frames(frames: list[dict]) -> list[dict]:
    return [f for f in frames if f.get("t") == SFRAME_OBJECT]


def _realistic_source(version: str = "v1") -> bytes:
    """~4 KB of varied Python source — compresses 2:1 with zstd, not to nothing.

    Generates 60 unique functions with distinct hex identifiers so the content
    is varied enough that zstd does not win with an extreme ratio.  A tiny
    one-character version change creates a delta that is much smaller than the
    compressed full object — ideal for testing delta profitability.
    """
    lines: list[str] = [
        "# Auto-generated module\n",
        f"__version__ = '{version}'\n",
        "import hashlib, struct, os, sys, json, logging\n",
        "from typing import Dict, List, Optional, Any\n",
        "\n",
        "logger = logging.getLogger(__name__)\n",
        "\n",
    ]
    for i in range(60):
        h = hex(i * 0x9E3779B9 & 0xFFFFFFFF)[2:].zfill(8)
        lines += [
            f"def compute_{h}(x: int, ctx: Dict[str, Any]) -> int:\n",
            f"    base = ctx.get('{h}', {i * 7 + 1})\n",
            f"    return (x ^ base) % {0xFFFF - i * 3}\n",
            "\n",
        ]
    return "".join(lines).encode()


# ---------------------------------------------------------------------------
# C1 — Server uses zstd encoding on O frames
# ---------------------------------------------------------------------------

@pytest.mark.asyncio
@pytest.mark.skipif(not ZSTD_AVAILABLE, reason="zstd not installed")
async def test_c1_server_uses_zstd_encoding(db_session: AsyncSession) -> None:
    repo = await create_repo(db_session, owner="test-comp-c1")
    repo_id = str(repo.repo_id)

    content = b"hello world " * 100
    oid = await _make_object(db_session, repo_id, content, path="c1/file.py", owner=repo.owner, slug=repo.slug)
    snap = await _make_snapshot(db_session, repo_id, {"c1/file.py": oid})
    commit = await _make_commit(db_session, repo_id, snapshot_id=snap.snapshot_id, seed="c1")

    from musehub.services.musehub_wire import wire_fetch_stream
    frames = await _collect_frames(
        wire_fetch_stream(db_session, repo_id, WireFetchRequest(want=[commit.commit_id], have=[]))
    )
    o = _o_frames(frames)
    assert o, "Expected at least one O frame"
    assert all(f.get("enc") == "zstd" for f in o), (
        f"Expected enc=zstd on all O frames, got: {[f.get('enc') for f in o]}"
    )


# ---------------------------------------------------------------------------
# C2 — zstd frame decompresses back to original bytes
# ---------------------------------------------------------------------------

@pytest.mark.asyncio
@pytest.mark.skipif(not ZSTD_AVAILABLE, reason="zstd not installed")
async def test_c2_zstd_frame_decompresses_correctly(db_session: AsyncSession) -> None:
    import zstandard

    repo = await create_repo(db_session, owner="test-comp-c2")
    repo_id = str(repo.repo_id)

    content = b"def foo(): pass\n" * 50
    oid = await _make_object(db_session, repo_id, content, path="c2/mod.py", owner=repo.owner, slug=repo.slug)
    snap = await _make_snapshot(db_session, repo_id, {"c2/mod.py": oid})
    commit = await _make_commit(db_session, repo_id, snapshot_id=snap.snapshot_id, seed="c2")

    from musehub.services.musehub_wire import wire_fetch_stream
    frames = await _collect_frames(
        wire_fetch_stream(db_session, repo_id, WireFetchRequest(want=[commit.commit_id], have=[]))
    )
    o = next((f for f in frames if f.get("t") == SFRAME_OBJECT and f.get("id") == oid), None)
    assert o is not None
    assert o["enc"] == "zstd"

    reconstructed = zstandard.ZstdDecompressor().decompress(bytes(o["content"]))
    assert reconstructed == content


# ---------------------------------------------------------------------------
# C3 — O frames arrive sorted by path within each batch
# ---------------------------------------------------------------------------

@pytest.mark.asyncio
async def test_c3_o_frames_sorted_by_path(db_session: AsyncSession) -> None:
    repo = await create_repo(db_session, owner="test-comp-c3")
    repo_id = str(repo.repo_id)

    # Create objects at paths that would be out of alphabetical order if unsorted
    paths = ["z/last.py", "a/first.py", "m/middle.py", "b/second.py"]
    manifest: dict[str, str] = {}
    for p in paths:
        oid = await _make_object(db_session, repo_id, f"content of {p}".encode(), path=p, owner=repo.owner, slug=repo.slug)
        manifest[p] = oid

    snap = await _make_snapshot(db_session, repo_id, manifest)
    commit = await _make_commit(db_session, repo_id, snapshot_id=snap.snapshot_id, seed="c3")

    from musehub.services.musehub_wire import wire_fetch_stream
    frames = await _collect_frames(
        wire_fetch_stream(db_session, repo_id, WireFetchRequest(want=[commit.commit_id], have=[]))
    )
    o_paths = [f.get("path", "") for f in frames if f.get("t") == SFRAME_OBJECT]
    assert o_paths == sorted(o_paths), (
        f"O frames not sorted by path.\nGot:      {o_paths}\nExpected: {sorted(o_paths)}"
    )


# ---------------------------------------------------------------------------
# C4 — Successive versions of same file → delta frame with base_id
# ---------------------------------------------------------------------------

@pytest.mark.asyncio
async def test_c4_successive_versions_use_delta_frame(db_session: AsyncSession) -> None:
    repo = await create_repo(db_session, owner="test-comp-c4")
    repo_id = str(repo.repo_id)

    # Realistic Python source: large varied content, tiny version change.
    # Delta must win against zstd full frame for this test to be meaningful.
    base_content = _realistic_source("v1")
    new_content  = _realistic_source("v2")

    base_oid = await _make_object(db_session, repo_id, base_content, path="src/greet.py", owner=repo.owner, slug=repo.slug)
    new_oid  = await _make_object(db_session, repo_id, new_content,  path="src/greet.py", owner=repo.owner, slug=repo.slug)

    snap1 = await _make_snapshot(db_session, repo_id, {"src/greet.py": base_oid})
    c1 = await _make_commit(db_session, repo_id, snapshot_id=snap1.snapshot_id, seed="c4-1")

    snap2 = await _make_snapshot(db_session, repo_id, {"src/greet.py": new_oid})
    c2 = await _make_commit(
        db_session, repo_id,
        parent_ids=[c1.commit_id],
        snapshot_id=snap2.snapshot_id,
        seed="c4-2",
    )

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

    # The first version (base) must be a full frame; the second a delta
    frame_for_new = next((f for f in o_frames if f.get("id") == new_oid), None)
    assert frame_for_new is not None, f"No O frame for {new_oid[:16]}"
    assert "delta" in frame_for_new.get("enc", ""), (
        f"Expected delta encoding for successive version, got enc={frame_for_new.get('enc')!r}"
    )
    assert frame_for_new.get("base_id") == base_oid, (
        f"base_id should be {base_oid[:16]}, got {str(frame_for_new.get('base_id'))[:16]}"
    )


# ---------------------------------------------------------------------------
# C5 — Delta frame reconstructs to correct target bytes
# ---------------------------------------------------------------------------

@pytest.mark.asyncio
async def test_c5_delta_frame_reconstructs_correctly(db_session: AsyncSession) -> None:
    repo = await create_repo(db_session, owner="test-comp-c5")
    repo_id = str(repo.repo_id)

    base_content = _realistic_source("a1")
    new_content  = _realistic_source("a2")

    base_oid = await _make_object(db_session, repo_id, base_content, path="lib/foo.py", owner=repo.owner, slug=repo.slug)
    new_oid  = await _make_object(db_session, repo_id, new_content,  path="lib/foo.py", owner=repo.owner, slug=repo.slug)

    snap1 = await _make_snapshot(db_session, repo_id, {"lib/foo.py": base_oid})
    c1 = await _make_commit(db_session, repo_id, snapshot_id=snap1.snapshot_id, seed="c5-1")
    snap2 = await _make_snapshot(db_session, repo_id, {"lib/foo.py": new_oid})
    c2 = await _make_commit(
        db_session, repo_id,
        parent_ids=[c1.commit_id],
        snapshot_id=snap2.snapshot_id,
        seed="c5-2",
    )

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

    frame_base = next((f for f in o_frames if f.get("id") == base_oid), None)
    frame_new  = next((f for f in o_frames if f.get("id") == new_oid),  None)
    assert frame_base is not None and frame_new is not None

    # Decompress base
    from muse.core.compression import decompress_frame
    base_enc = frame_base.get("enc", "raw")
    decoded_base = decompress_frame(bytes(frame_base["content"]), base_enc)

    # Apply delta
    enc = frame_new.get("enc", "")
    assert "delta" in enc
    delta_bytes = bytes(frame_new["content"])
    # Delta stream is zlib-compressed inside (compute_delta always zlib-wraps)
    reconstructed = apply_delta(decoded_base, delta_bytes)

    assert reconstructed == new_content
    assert blob_id(reconstructed) == new_oid


# ---------------------------------------------------------------------------
# C6 — Unprofitable delta → falls back to full frame
# ---------------------------------------------------------------------------

@pytest.mark.asyncio
async def test_c6_unprofitable_delta_uses_full_frame(db_session: AsyncSession) -> None:
    import os

    repo = await create_repo(db_session, owner="test-comp-c6")
    repo_id = str(repo.repo_id)

    # Two completely random (incompressible) objects at the same path
    rng = __import__("random").Random(42)
    base_content = bytes(rng.randint(0, 255) for _ in range(4096))
    new_content  = bytes(rng.randint(0, 255) for _ in range(4096))

    base_oid = await _make_object(db_session, repo_id, base_content, path="bin/rand.dat", owner=repo.owner, slug=repo.slug)
    new_oid  = await _make_object(db_session, repo_id, new_content,  path="bin/rand.dat", owner=repo.owner, slug=repo.slug)

    snap1 = await _make_snapshot(db_session, repo_id, {"bin/rand.dat": base_oid})
    c1 = await _make_commit(db_session, repo_id, snapshot_id=snap1.snapshot_id, seed="c6-1")
    snap2 = await _make_snapshot(db_session, repo_id, {"bin/rand.dat": new_oid})
    c2 = await _make_commit(
        db_session, repo_id,
        parent_ids=[c1.commit_id],
        snapshot_id=snap2.snapshot_id,
        seed="c6-2",
    )

    from musehub.services.musehub_wire import wire_fetch_stream
    frames = await _collect_frames(
        wire_fetch_stream(db_session, repo_id, WireFetchRequest(want=[c2.commit_id], have=[]))
    )
    frame_new = next((f for f in frames if f.get("t") == SFRAME_OBJECT and f.get("id") == new_oid), None)
    assert frame_new is not None
    assert "delta" not in frame_new.get("enc", ""), (
        "Expected full frame for incompressible random content, got delta"
    )


# ---------------------------------------------------------------------------
# C7 — Full clone with delta mix delivers correct object_id set
# ---------------------------------------------------------------------------

@pytest.mark.asyncio
async def test_c7_full_clone_correct_object_ids(db_session: AsyncSession) -> None:
    repo = await create_repo(db_session, owner="test-comp-c7")
    repo_id = str(repo.repo_id)

    # Three commits: add file, modify it, add another file
    v1 = b"version one\n" * 50
    v2 = b"version two\n" * 50
    other = b"other file\n" * 30

    oid_v1    = await _make_object(db_session, repo_id, v1,    path="src/main.py", owner=repo.owner, slug=repo.slug)
    oid_v2    = await _make_object(db_session, repo_id, v2,    path="src/main.py", owner=repo.owner, slug=repo.slug)
    oid_other = await _make_object(db_session, repo_id, other, path="src/util.py", owner=repo.owner, slug=repo.slug)

    snap1 = await _make_snapshot(db_session, repo_id, {"src/main.py": oid_v1})
    c1 = await _make_commit(db_session, repo_id, snapshot_id=snap1.snapshot_id, seed="c7-1")

    snap2 = await _make_snapshot(db_session, repo_id, {"src/main.py": oid_v2})
    c2 = await _make_commit(db_session, repo_id, parent_ids=[c1.commit_id], snapshot_id=snap2.snapshot_id, seed="c7-2")

    snap3 = await _make_snapshot(db_session, repo_id, {"src/main.py": oid_v2, "src/util.py": oid_other})
    c3 = await _make_commit(db_session, repo_id, parent_ids=[c2.commit_id], snapshot_id=snap3.snapshot_id, seed="c7-3")

    from musehub.services.musehub_wire import wire_fetch_stream
    frames = await _collect_frames(
        wire_fetch_stream(db_session, repo_id, WireFetchRequest(want=[c3.commit_id], have=[]))
    )
    received_oids = {f["id"] for f in frames if f.get("t") == SFRAME_OBJECT}
    expected_oids = {oid_v1, oid_v2, oid_other}
    assert received_oids == expected_oids, (
        f"Missing: {expected_oids - received_oids}, Extra: {received_oids - expected_oids}"
    )


# ---------------------------------------------------------------------------
# C8 — Cross-batch delta: prior-batch object serves as base for next batch
# ---------------------------------------------------------------------------

@pytest.mark.asyncio
async def test_c8_cross_batch_delta(db_session: AsyncSession) -> None:
    """An object in batch N+1 that shares a path with an object in batch N
    should be delta-encoded against the batch-N object."""
    from musehub.services.musehub_wire import _COMMIT_BATCH

    repo = await create_repo(db_session, owner="test-comp-c8")
    repo_id = str(repo.repo_id)

    base_content = _realistic_source("r1")
    new_content  = _realistic_source("r2")

    base_oid = await _make_object(db_session, repo_id, base_content, path="worker.py", owner=repo.owner, slug=repo.slug)
    new_oid  = await _make_object(db_session, repo_id, new_content,  path="worker.py", owner=repo.owner, slug=repo.slug)

    # Put base in batch 1 and new version in batch 2 (requires > _COMMIT_BATCH commits between them)
    snap_base = await _make_snapshot(db_session, repo_id, {"worker.py": base_oid})
    c_base = await _make_commit(db_session, repo_id, snapshot_id=snap_base.snapshot_id, seed="c8-base")

    # Pad with _COMMIT_BATCH commits so base ends up in batch 1 (reuse same snap)
    prev = c_base.commit_id
    for i in range(_COMMIT_BATCH):
        c = await _make_commit(
            db_session, repo_id,
            parent_ids=[prev],
            snapshot_id=snap_base.snapshot_id,
            seed=f"c8-fill-{i}",
        )
        prev = c.commit_id

    snap_new = await _make_snapshot(db_session, repo_id, {"worker.py": new_oid})
    c_new = await _make_commit(
        db_session, repo_id,
        parent_ids=[prev],
        snapshot_id=snap_new.snapshot_id,
        seed="c8-new",
    )

    from musehub.services.musehub_wire import wire_fetch_stream
    frames = await _collect_frames(
        wire_fetch_stream(db_session, repo_id, WireFetchRequest(want=[c_new.commit_id], have=[]))
    )
    frame_new = next((f for f in frames if f.get("t") == SFRAME_OBJECT and f.get("id") == new_oid), None)
    assert frame_new is not None
    assert "delta" in frame_new.get("enc", ""), (
        f"Expected delta frame for cross-batch successive version, got enc={frame_new.get('enc')!r}"
    )
    assert frame_new.get("base_id") == base_oid