"""TDD — R2 presigned push path.

Problem
-------
Cloudflare times out streaming POSTs after ~100 seconds.  Repos with
more than ~500 objects (e.g. the muse repo at 6905 objects / 191 MB raw)
fail mid-upload.  The fix: for large pushes the client calls
``POST /push/presign`` first, PUTs objects directly to R2 (bypassing CF),
then sends a compressed H+C+E stream with zero O frames.  The server's
existing referential integrity check (lines 1758-1786 of musehub_wire.py)
already handles objects that arrive pre-stored.

Test plan
---------
P1  wire_push_presign — LocalBackend: missing objects land in stream_these,
    already-stored objects land in already_stored, presigned_urls is empty.

P2  wire_push_presign — mocked S3Backend: missing objects get presigned PUT
    URLs; already-stored objects land in already_stored.

P3  wire_push_presign — empty object list returns all-empty response.

P4  push/stream with zero O frames and all objects pre-stored in the DB
    finalises the commit successfully (referential integrity passes).

P5  push/stream with zero O frames but objects NOT in storage returns 422
    (referential integrity rejects the push).

P6  Full round-trip: presign → store objects via backend.put → push stream
    with zero O frames → /refs reflects the new head.

P7  S3Backend.presign_batch generates correctly-keyed URLs for a batch of
    object IDs (mocked boto3; no real AWS call).

Confirm step (missing from initial implementation — discovered via staging test)
--------------------------------------------------------------------------------
The presigned PUT goes directly to R2, bypassing the server entirely.  The
server therefore has no DB record for those objects, so the referential
integrity check on the zero-O-frame push/stream rejects with 422.

Fix: the client calls ``POST /push/confirm`` after all R2 PUTs succeed.
The server inserts a ``musehub_objects`` row (using ``backend.uri_for(oid)``)
and a ``musehub_object_refs`` row for each confirmed object, then the zero-O
push/stream passes the integrity check.

P9   wire_push_confirm inserts DB rows for confirmed objects; subsequent
     push/stream with zero O frames succeeds.
P10  wire_push_confirm is idempotent — confirming the same object twice
     does not duplicate rows or raise.
P11  Full end-to-end: presign → R2 PUT → confirm → zero-O push/stream →
     commit is accessible via wire_refs.
"""
from __future__ import annotations
from collections.abc import AsyncIterator

import asyncio
import zlib
from datetime import datetime, timezone
from unittest.mock import AsyncMock, MagicMock, patch

import msgpack
import pytest
from sqlalchemy.dialects.postgresql import insert as pg_insert
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,
)
from muse.core.mpack import MuseWireFrameWriter
from musehub.types.json_types import JSONObject
from tests.factories import create_repo

_fw = MuseWireFrameWriter()


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

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


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


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


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


def _header_frame(
    n_objects: int = 0,
    n_commits: int = 1,
    branch: str = "main",
    force: bool = True,
) -> bytes:
    return _wrap(SFRAME_HEADER, {
        "t": SFRAME_HEADER,
        "branch": branch,
        "force": force,
        "n_objects": n_objects,
        "n_commits": n_commits,
    })


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


def _commit_pack_frame(commits: list[dict], snapshots: list[dict] | None = None) -> bytes:
    return _wrap(SFRAME_COMMIT_PACK, {
        "t": SFRAME_COMMIT_PACK,
        "commits": commits,
        "snapshots": snapshots or [],
    })


def _o_frame(oid: str, content: bytes, path: str = "file.py") -> bytes:
    return _wrap(SFRAME_OBJECT, {
        "t": SFRAME_OBJECT,
        "id": oid,
        "path": path,
        "content": zlib.compress(content, level=1),
        "enc": "zlib",
        "size": len(content),
    })


async def _store_object(
    session: AsyncSession,
    repo_id: str,
    oid: str,
    content: bytes,
    *,
    owner: str,
    slug: str,
) -> None:
    """Insert object directly into DB + local storage (bypasses wire protocol)."""
    from musehub.services.musehub_wire import get_backend
    from musehub.storage.backends import repo_root_for
    backend = get_backend()
    repo_root = repo_root_for(owner, slug)
    uri = await backend.put(oid, content, repo_root=repo_root)
    await session.execute(
        pg_insert(db.MusehubObject)
        .values(
            object_id=oid,
            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()


async def _make_commit_and_snapshot(
    session: AsyncSession,
    repo_id: str,
    *,
    manifest: dict[str, str],
    commit_seed: str = "c1",
    parent_ids: list[str] | None = None,
) -> tuple[db.MusehubCommit, db.MusehubSnapshot]:
    snap_id = _uid(f"snap-{commit_seed}")
    snap = db.MusehubSnapshot(
        snapshot_id=snap_id,
        repo_id=repo_id,
        directories=[],
        manifest_blob=msgpack.packb(manifest, use_bin_type=True),
        entry_count=len(manifest),
        created_at=_now(),
    )
    session.add(snap)
    commit_id = _uid(f"commit-{commit_seed}")
    commit = db.MusehubCommit(
        commit_id=commit_id,
        repo_id=repo_id,
        branch="main",
        parent_ids=parent_ids or [],
        message=f"commit {commit_seed}",
        author="gabriel",
        timestamp=_now(),
        snapshot_id=snap_id,
    )
    session.add(commit)
    await session.commit()
    return commit, snap


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


async def _body_iter(raw: bytes) -> None:
    yield raw


# ---------------------------------------------------------------------------
# P1 — LocalBackend: missing → stream_these, present → already_stored
# ---------------------------------------------------------------------------

@pytest.mark.asyncio
async def test_p1_local_backend_missing_goes_to_stream_these(
    db_session: AsyncSession,
) -> None:
    """LocalBackend.presign_batch returns {} → missing oids land in stream_these."""
    from musehub.services.musehub_wire import wire_push_presign

    repo = await create_repo(db_session, owner="gabriel")
    repo_id = str(repo.repo_id)

    content_a = b"object-A content for presign test"
    content_b = b"object-B content for presign test"
    oid_a = blob_id(content_a)
    oid_b = blob_id(content_b)

    # Pre-store A so it lands in already_stored.
    await _store_object(db_session, repo_id, oid_a, content_a, owner=repo.owner, slug=repo.slug)

    result = await wire_push_presign(
        db_session,
        repo_id,
        [oid_a, oid_b],
    )

    assert oid_a in result["already_stored"]
    assert oid_b not in result["already_stored"]
    # LocalBackend returns no presigned URLs.
    assert result["presigned_urls"] == {}
    # Missing object with no presigned URL goes to stream_these.
    assert oid_b in result["stream_these"]
    assert oid_a not in result["stream_these"]


# ---------------------------------------------------------------------------
# P2 — S3Backend: missing → presigned_urls
# ---------------------------------------------------------------------------

@pytest.mark.asyncio
async def test_p2_s3_backend_missing_gets_presigned_url(
    db_session: AsyncSession,
) -> None:
    """S3Backend.presign_batch is called for missing objects; URLs are returned."""
    from musehub.services.musehub_wire import wire_push_presign

    repo = await create_repo(db_session, owner="gabriel")
    repo_id = str(repo.repo_id)

    content = b"object content for s3 presign test"
    oid = blob_id(content)
    fake_url = f"https://r2.example.com/objects/{oid}?X-Amz-Signature=abc"

    mock_backend = MagicMock()
    mock_backend.presign_batch = AsyncMock(return_value={oid: fake_url})

    with patch("musehub.services.musehub_wire.get_backend", return_value=mock_backend):
        result = await wire_push_presign(
            db_session,
            repo_id,
            [oid],
        )

    assert result["presigned_urls"] == {oid: fake_url}
    assert result["already_stored"] == []
    assert result["stream_these"] == []
    mock_backend.presign_batch.assert_called_once_with([oid], "put", 3600)


# ---------------------------------------------------------------------------
# P3 — empty object list → all-empty response
# ---------------------------------------------------------------------------

@pytest.mark.asyncio
async def test_p3_empty_object_list(db_session: AsyncSession) -> None:
    from musehub.services.musehub_wire import wire_push_presign

    repo = await create_repo(db_session, owner="gabriel")

    result = await wire_push_presign(db_session, str(repo.repo_id), [])

    assert result["presigned_urls"] == {}
    assert result["already_stored"] == []
    assert result["stream_these"] == []


# ---------------------------------------------------------------------------
# P4 — zero O frames, all objects pre-stored → commit finalises
# ---------------------------------------------------------------------------

@pytest.mark.asyncio
async def test_p4_zero_o_frames_all_prestored_commit_finalises(
    db_session: AsyncSession,
) -> None:
    """H + C + E with n_objects=0 succeeds when snapshot objects are in the DB."""
    from musehub.services.musehub_wire import wire_push_stream

    repo = await create_repo(db_session, owner="gabriel")
    repo_id = str(repo.repo_id)

    content = b"pre-stored object content p4"
    oid = blob_id(content)
    await _store_object(db_session, repo_id, oid, content, owner=repo.owner, slug=repo.slug)

    snap_id = _uid("snap-p4")
    commit_id = _uid("commit-p4")

    commit_wire = {
        "commit_id": commit_id,
        "parent_ids": [],
        "message": "presign round-trip",
        "author": "gabriel",
        "timestamp": _now().isoformat(),
        "snapshot_id": snap_id,
        "branch": "main",
    }
    snap_wire = {
        "snapshot_id": snap_id,
        "manifest": {"file.py": oid},
        "directories": [],
        "created_at": _now().isoformat(),
    }

    raw = (
        _header_frame(n_objects=0, n_commits=1)
        + _commit_pack_frame([commit_wire], [snap_wire])
        + _end_frame(n_objects=0, n_commits=1)
    )

    frames = await _collect_frames(
        wire_push_stream(db_session, repo_id, _body_iter(raw), "gabriel")
    )

    # No error frames.
    errors = [f for f in frames if f.get("t") == "X"]
    assert errors == [], f"unexpected error frames: {errors}"

    # Commit exists in DB.
    from sqlalchemy import select
    row = (await db_session.execute(
        select(db.MusehubCommit).where(db.MusehubCommit.commit_id == commit_id)
    )).scalar_one_or_none()
    assert row is not None


# ---------------------------------------------------------------------------
# P5 — zero O frames, objects NOT stored → 422
# ---------------------------------------------------------------------------

@pytest.mark.asyncio
async def test_p5_zero_o_frames_missing_objects_yields_422(
    db_session: AsyncSession,
) -> None:
    """H + C + E with n_objects=0 but snapshot references absent objects → 422."""
    from musehub.services.musehub_wire import wire_push_stream

    repo = await create_repo(db_session, owner="gabriel")
    repo_id = str(repo.repo_id)

    # This oid is NOT stored anywhere.
    oid = blob_id(b"ghost object never uploaded")

    snap_id = _uid("snap-p5")
    commit_id = _uid("commit-p5")

    commit_wire = {
        "commit_id": commit_id,
        "parent_ids": [],
        "message": "should fail",
        "author": "gabriel",
        "timestamp": _now().isoformat(),
        "snapshot_id": snap_id,
        "branch": "main",
    }
    snap_wire = {
        "snapshot_id": snap_id,
        "manifest": {"missing.py": oid},
        "directories": [],
        "created_at": _now().isoformat(),
    }

    raw = (
        _header_frame(n_objects=0, n_commits=1)
        + _commit_pack_frame([commit_wire], [snap_wire])
        + _end_frame(n_objects=0, n_commits=1)
    )

    frames = await _collect_frames(
        wire_push_stream(db_session, repo_id, _body_iter(raw), "gabriel")
    )

    error_frames = [f for f in frames if f.get("t") == "X"]
    assert error_frames, "expected an error frame for missing objects"
    assert error_frames[0].get("code") == 422


# ---------------------------------------------------------------------------
# P6 — full round-trip: presign → store → push stream → /refs reflects head
# ---------------------------------------------------------------------------

@pytest.mark.asyncio
async def test_p6_full_round_trip_presign_store_push(
    db_session: AsyncSession,
) -> None:
    """presign returns oids to upload; after manual storage, zero-O-frame push works."""
    from musehub.services.musehub_wire import wire_push_presign, wire_push_stream

    repo = await create_repo(db_session, owner="gabriel")
    repo_id = str(repo.repo_id)

    content = b"p6 round-trip object bytes"
    oid = blob_id(content)

    # Step 1: Ask presign endpoint what to upload.
    presign_result = await wire_push_presign(db_session, repo_id, [oid])
    assert oid in presign_result["stream_these"]  # LocalBackend → no presigned URL

    # Step 2: Client "uploads" directly to storage (simulates R2 PUT via presigned URL).
    await _store_object(db_session, repo_id, oid, content, owner=repo.owner, slug=repo.slug)

    # Step 3: Now presign endpoint sees it as already_stored.
    presign_result2 = await wire_push_presign(db_session, repo_id, [oid])
    assert oid in presign_result2["already_stored"]
    assert presign_result2["stream_these"] == []

    # Step 4: Push H + C + E with zero O frames.
    snap_id = _uid("snap-p6")
    commit_id = _uid("commit-p6")
    commit_wire = {
        "commit_id": commit_id,
        "parent_ids": [],
        "message": "p6 round-trip",
        "author": "gabriel",
        "timestamp": _now().isoformat(),
        "snapshot_id": snap_id,
        "branch": "main",
    }
    snap_wire = {
        "snapshot_id": snap_id,
        "manifest": {"round_trip.py": oid},
        "directories": [],
        "created_at": _now().isoformat(),
    }
    raw = (
        _header_frame(n_objects=0, n_commits=1)
        + _commit_pack_frame([commit_wire], [snap_wire])
        + _end_frame(n_objects=0, n_commits=1)
    )
    frames = await _collect_frames(
        wire_push_stream(db_session, repo_id, _body_iter(raw), "gabriel")
    )

    errors = [f for f in frames if f.get("t") == "X"]
    assert errors == [], f"push failed: {errors}"

    # Step 5: Verify commit is in the DB.
    from sqlalchemy import select
    row = (await db_session.execute(
        select(db.MusehubCommit).where(db.MusehubCommit.commit_id == commit_id)
    )).scalar_one_or_none()
    assert row is not None
    assert row.branch == "main"


# ---------------------------------------------------------------------------
# P7 — S3Backend.presign_batch produces correctly-keyed URLs
# ---------------------------------------------------------------------------

@pytest.mark.asyncio
async def test_p7_s3_backend_presign_batch_url_format() -> None:
    """S3Backend.presign_batch calls boto3 with the right key format per oid."""
    from musehub.storage.backends import S3Backend

    oid = blob_id(b"test object for presign key format")
    expected_key = f"objects/{oid}"
    fake_url = f"https://bucket.r2.cloudflarestorage.com/{expected_key}?sig=abc"

    mock_client = MagicMock()
    mock_client.generate_presigned_url.return_value = fake_url

    backend = S3Backend(
        bucket="test-bucket",
        region="auto",
        endpoint_url="https://r2.example.com",
        access_key_id="key",
        secret_access_key="secret",
    )
    backend._client = mock_client

    result = await backend.presign_batch([oid], "put", 3600)

    assert result == {oid: fake_url}
    mock_client.generate_presigned_url.assert_called_once_with(
        "put_object",
        Params={"Bucket": "test-bucket", "Key": expected_key},
        ExpiresIn=3600,
    )


# ---------------------------------------------------------------------------
# P9 — confirm inserts DB rows; zero-O push/stream then succeeds
# ---------------------------------------------------------------------------

@pytest.mark.asyncio
async def test_p9_confirm_registers_objects_and_push_succeeds(
    db_session: AsyncSession,
) -> None:
    """wire_push_confirm creates DB records; referential integrity check passes."""
    from musehub.services.musehub_wire import wire_push_confirm, wire_push_stream

    repo = await create_repo(db_session, owner="gabriel")
    repo_id = str(repo.repo_id)

    content = b"p9 object bytes confirm test"
    oid = blob_id(content)

    # Simulate: client PUT raw bytes to R2 (storage has bytes but no DB row).
    from musehub.services.musehub_wire import get_backend
    from musehub.storage.backends import repo_root_for
    backend = get_backend()
    await backend.put(oid, content, repo_root=repo_root_for(repo.owner, repo.slug))

    # Confirm step: server registers the object.
    await wire_push_confirm(
        db_session,
        repo_id,
        objects=[{"object_id": oid, "size_bytes": len(content), "path": "p9.py"}],
    )

    # Now zero-O push/stream should pass referential integrity.
    snap_id = _uid("snap-p9")
    commit_id = _uid("commit-p9")
    commit_wire = {
        "commit_id": commit_id,
        "parent_ids": [],
        "message": "p9 confirm test",
        "author": "gabriel",
        "timestamp": _now().isoformat(),
        "snapshot_id": snap_id,
        "branch": "main",
    }
    snap_wire = {
        "snapshot_id": snap_id,
        "manifest": {"p9.py": oid},
        "directories": [],
        "created_at": _now().isoformat(),
    }
    raw = (
        _header_frame(n_objects=0, n_commits=1)
        + _commit_pack_frame([commit_wire], [snap_wire])
        + _end_frame(n_objects=0, n_commits=1)
    )
    frames = await _collect_frames(
        wire_push_stream(db_session, repo_id, _body_iter(raw), "gabriel")
    )
    errors = [f for f in frames if f.get("t") == "X"]
    assert errors == [], f"push failed after confirm: {errors}"

    from sqlalchemy import select
    row = (await db_session.execute(
        select(db.MusehubCommit).where(db.MusehubCommit.commit_id == commit_id)
    )).scalar_one_or_none()
    assert row is not None


# ---------------------------------------------------------------------------
# P10 — confirm is idempotent
# ---------------------------------------------------------------------------

@pytest.mark.asyncio
async def test_p10_confirm_is_idempotent(db_session: AsyncSession) -> None:
    """Confirming the same object twice does not raise or duplicate rows."""
    from musehub.services.musehub_wire import wire_push_confirm
    from musehub.services.musehub_wire import get_backend
    from musehub.storage.backends import repo_root_for

    repo = await create_repo(db_session, owner="gabriel")
    repo_id = str(repo.repo_id)

    content = b"p10 idempotent confirm bytes"
    oid = blob_id(content)
    backend = get_backend()
    await backend.put(oid, content, repo_root=repo_root_for(repo.owner, repo.slug))

    obj_entry = [{"object_id": oid, "size_bytes": len(content), "path": "p10.py"}]

    await wire_push_confirm(db_session, repo_id, obj_entry)
    await wire_push_confirm(db_session, repo_id, obj_entry)  # second call must not raise

    from sqlalchemy import select, func
    count = (await db_session.execute(
        select(func.count()).where(db.MusehubObject.object_id == oid)
    )).scalar()
    assert count == 1


# ---------------------------------------------------------------------------
# P11 — full end-to-end: presign → R2 PUT → confirm → zero-O stream → refs
# ---------------------------------------------------------------------------

@pytest.mark.asyncio
async def test_p11_full_end_to_end_with_confirm(db_session: AsyncSession) -> None:
    """presign → backend.put (simulate R2) → confirm → push/stream → commit in refs."""
    from musehub.services.musehub_wire import (
        wire_push_presign,
        wire_push_confirm,
        wire_push_stream,
    )
    from musehub.services.musehub_wire import get_backend

    repo = await create_repo(db_session, owner="gabriel")
    repo_id = str(repo.repo_id)

    content = b"p11 end to end object bytes"
    oid = blob_id(content)

    # Step 1: presign → object is missing, goes to stream_these (LocalBackend)
    presign_result = await wire_push_presign(db_session, repo_id, [oid])
    assert oid in presign_result["stream_these"]

    # Step 2: client "PUT" to R2 (simulate by calling backend.put directly)
    from musehub.storage.backends import repo_root_for
    backend = get_backend()
    await backend.put(oid, content, repo_root=repo_root_for(repo.owner, repo.slug))

    # Step 3: confirm — registers DB row
    await wire_push_confirm(
        db_session,
        repo_id,
        objects=[{"object_id": oid, "size_bytes": len(content), "path": "p11.py"}],
    )

    # Step 4: presign now shows object as already_stored
    presign_result2 = await wire_push_presign(db_session, repo_id, [oid])
    assert oid in presign_result2["already_stored"]

    # Step 5: zero-O push/stream
    snap_id = _uid("snap-p11")
    commit_id = _uid("commit-p11")
    commit_wire = {
        "commit_id": commit_id,
        "parent_ids": [],
        "message": "p11 end to end",
        "author": "gabriel",
        "timestamp": _now().isoformat(),
        "snapshot_id": snap_id,
        "branch": "main",
    }
    snap_wire = {
        "snapshot_id": snap_id,
        "manifest": {"p11.py": oid},
        "directories": [],
        "created_at": _now().isoformat(),
    }
    raw = (
        _header_frame(n_objects=0, n_commits=1)
        + _commit_pack_frame([commit_wire], [snap_wire])
        + _end_frame(n_objects=0, n_commits=1)
    )
    frames = await _collect_frames(
        wire_push_stream(db_session, repo_id, _body_iter(raw), "gabriel")
    )
    errors = [f for f in frames if f.get("t") == "X"]
    assert errors == [], f"end-to-end push failed: {errors}"

    from sqlalchemy import select
    commit_row = (await db_session.execute(
        select(db.MusehubCommit).where(db.MusehubCommit.commit_id == commit_id)
    )).scalar_one_or_none()
    assert commit_row is not None