from __future__ import annotations

import json
import shutil
import struct
import uuid
from datetime import UTC, datetime, timedelta
from pathlib import Path

import httpx
import pytest
from modelforge_api.domain.acquisition import AgentArtifactJobFile, AgentArtifactJobLease
from pydantic import ValidationError

from modelforge_node_agent.acquisition import (
    AcquisitionCancelled,
    AcquisitionFailure,
    ArtifactAcquirer,
    _confined,
    _inspect_file,
    _relative_path,
)
from modelforge_node_agent.settings import AgentSettings
from modelforge_node_agent.transport import AgentTransport


class FakeTransport:
    def __init__(self, *, cancel: bool = False) -> None:
        self.cancel = cancel
        self.progress: list[dict] = []
        self.completed: dict | None = None

    def artifact_progress(self, _job_id: str, payload: dict, _credential: str) -> dict:
        self.progress.append(payload)
        return {
            "accepted": True,
            "cancel_requested": self.cancel and len(self.progress) > 1,
            "lease_expires_at": (datetime.now(UTC) + timedelta(seconds=120)).isoformat(),
        }

    def artifact_complete(self, _job_id: str, payload: dict, _credential: str) -> dict:
        self.completed = payload
        return {"status": "completed"}


def lease(target_root: Path, *, size: int = 8) -> AgentArtifactJobLease:
    return AgentArtifactJobLease(
        job_id=uuid.uuid4(),
        lease_token="x" * 40,
        lease_expires_at=datetime.now(UTC) + timedelta(minutes=2),
        repository_id="org/model",
        resolved_commit_sha="a" * 40,
        storage_root_id=uuid.uuid4(),
        target_root=str(target_root),
        total_size_bytes=size,
        reserve_bytes=0,
        reserve_percent=0,
        files=[
            AgentArtifactJobFile(
                ordinal=0,
                path="config.json",
                size_bytes=size,
                upstream_sha256=None,
                file_format="json",
                role="configuration",
                risk_flags=[],
            )
        ],
    )


def test_path_confinement_rejects_escape(tmp_path: Path) -> None:
    with pytest.raises(AcquisitionFailure, match="confined"):
        _relative_path("../secret")
    with pytest.raises(AcquisitionFailure, match="escaped"):
        _confined(tmp_path, tmp_path / ".." / "outside")


def test_empty_job_poll_decodes_json_null_as_no_job() -> None:
    transport = AgentTransport("https://control.invalid", 5, True)
    transport.client = httpx.Client(
        base_url="https://control.invalid",
        transport=httpx.MockTransport(
            lambda request: httpx.Response(200, content=b"null", request=request)
        ),
    )
    try:
        assert transport.next_artifact_job("credential") is None
    finally:
        transport.close()


def test_typed_job_rejects_arbitrary_command_data(tmp_path: Path) -> None:
    payload = lease(tmp_path).model_dump(mode="json")
    payload["command"] = "sh -c anything"
    with pytest.raises(ValidationError, match="Extra inputs are not permitted"):
        AgentArtifactJobLease.model_validate(payload)


def test_static_inspection_reads_header_without_deserializing(tmp_path: Path) -> None:
    header = json.dumps({"weight": {"dtype": "F32", "shape": [1], "data_offsets": [0, 4]}}).encode()
    target = tmp_path / "model.safetensors"
    target.write_bytes(struct.pack("<Q", len(header)) + header + b"\0\0\0\0")
    plan_file = AgentArtifactJobFile(
        ordinal=0,
        path=target.name,
        size_bytes=target.stat().st_size,
        upstream_sha256=None,
        file_format="safetensors",
        role="weights",
        risk_flags=[],
    )
    findings = _inspect_file(target, plan_file)
    assert findings[0]["type"] == "safetensors_header"
    assert findings[0]["status"] == "passed"


def test_pickle_and_remote_code_are_blocking(tmp_path: Path) -> None:
    target = tmp_path / "weights.bin"
    target.write_bytes(b"not executed")
    plan_file = AgentArtifactJobFile(
        ordinal=0,
        path=target.name,
        size_bytes=target.stat().st_size,
        upstream_sha256=None,
        file_format="bin",
        role="weights",
        risk_flags=["pickle_or_executable_serialization"],
    )
    assert _inspect_file(target, plan_file)[0]["severity"] == "block"


def test_auto_map_configuration_is_blocking_without_loading_code(tmp_path: Path) -> None:
    target = tmp_path / "config.json"
    target.write_text('{"auto_map":{"AutoModel":"modeling.Custom"}}', encoding="utf-8")
    plan_file = AgentArtifactJobFile(
        ordinal=0,
        path=target.name,
        size_bytes=target.stat().st_size,
        upstream_sha256=None,
        file_format="json",
        role="configuration",
        risk_flags=[],
    )
    assert _inspect_file(target, plan_file)[0]["type"] == "remote_code_configuration"
    assert _inspect_file(target, plan_file)[0]["severity"] == "block"


def test_mocked_download_resumes_and_promotes_atomically(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
    root = tmp_path / "artifacts" / "model-registry"
    job = lease(root)
    stage = root / ".quarantine" / str(job.job_id)
    stage.mkdir(parents=True)
    (stage / "config.json.part").write_bytes(b"1234")

    def handler(request: httpx.Request) -> httpx.Response:
        assert request.headers["range"] == "bytes=4-"
        return httpx.Response(206, content=b"5678", request=request)

    original_client = httpx.Client
    monkeypatch.setattr(
        "modelforge_node_agent.acquisition.httpx.Client",
        lambda **kwargs: original_client(transport=httpx.MockTransport(handler), **kwargs),
    )
    transport = FakeTransport()
    settings = AgentSettings(
        artifact_path=tmp_path / "artifacts",
        artifact_download_timeout_seconds=10,
    )
    ArtifactAcquirer(settings, transport, "credential").execute(job.model_dump(mode="json"))
    final = root / "repositories" / "org--model" / ("a" * 40) / "config.json"
    assert final.read_bytes() == b"12345678"
    assert not stage.exists()
    assert transport.completed is not None
    assert transport.completed["files"][0]["sha256"]
    verifying = next(item for item in transport.progress if item["status"] == "verifying")
    assert verifying["progress_bytes"] == 8


def test_download_rejects_bytes_beyond_declared_size_before_writing_them(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
    root = tmp_path / "artifacts" / "model-registry"
    job = lease(root)

    def handler(request: httpx.Request) -> httpx.Response:
        return httpx.Response(200, content=b"123456789", request=request)

    original_client = httpx.Client
    monkeypatch.setattr(
        "modelforge_node_agent.acquisition.httpx.Client",
        lambda **kwargs: original_client(transport=httpx.MockTransport(handler), **kwargs),
    )
    settings = AgentSettings(artifact_path=tmp_path / "artifacts")
    with pytest.raises(AcquisitionFailure, match="exceeded declared size") as failure:
        ArtifactAcquirer(settings, FakeTransport(), "credential").execute(
            job.model_dump(mode="json")
        )

    assert failure.value.code == "download_size_mismatch"
    assert failure.value.retryable is False
    partial = root / ".quarantine" / str(job.job_id) / "config.json.part"
    assert partial.stat().st_size == 0


def test_complete_partial_is_verified_without_another_request(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
    root = tmp_path / "artifacts" / "model-registry"
    job = lease(root)
    stage = root / ".quarantine" / str(job.job_id)
    stage.mkdir(parents=True)
    (stage / "config.json.part").write_bytes(b"12345678")

    def fail_if_called(_request: httpx.Request) -> httpx.Response:
        raise AssertionError("a complete partial must not trigger a network request")

    original_client = httpx.Client
    monkeypatch.setattr(
        "modelforge_node_agent.acquisition.httpx.Client",
        lambda **kwargs: original_client(
            transport=httpx.MockTransport(fail_if_called), **kwargs
        ),
    )
    transport = FakeTransport()
    settings = AgentSettings(artifact_path=tmp_path / "artifacts")
    ArtifactAcquirer(settings, transport, "credential").execute(job.model_dump(mode="json"))
    final = root / "repositories" / "org--model" / ("a" * 40) / "config.json"
    assert final.read_bytes() == b"12345678"
    assert transport.completed is not None


def test_promoted_manifest_reconciles_without_redownload(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
    root = tmp_path / "artifacts" / "model-registry"
    job = lease(root)
    final = root / "repositories" / "org--model" / ("a" * 40)
    final.mkdir(parents=True)
    (final / "config.json").write_bytes(b"12345678")
    (final / ".modelforge-manifest.json").write_text(
        json.dumps(
            {
                "repository_id": "org/model",
                "resolved_commit_sha": "a" * 40,
                "files": [],
            }
        ),
        encoding="utf-8",
    )

    def fail_if_called(_request: httpx.Request) -> httpx.Response:
        raise AssertionError("a promoted exact-revision set must not be downloaded again")

    original_client = httpx.Client
    monkeypatch.setattr(
        "modelforge_node_agent.acquisition.httpx.Client",
        lambda **kwargs: original_client(
            transport=httpx.MockTransport(fail_if_called), **kwargs
        ),
    )
    transport = FakeTransport()
    settings = AgentSettings(artifact_path=tmp_path / "artifacts")
    ArtifactAcquirer(settings, transport, "credential").execute(job.model_dump(mode="json"))
    assert transport.completed is not None
    assert transport.completed["promoted_relative_path"].endswith("a" * 40)


def test_cancel_is_observed_before_promotion(tmp_path: Path) -> None:
    root = tmp_path / "artifacts" / "model-registry"
    job = lease(root, size=0)
    job.files[0].size_bytes = 0
    stage = root / ".quarantine" / str(job.job_id)
    stage.mkdir(parents=True)
    (stage / "config.json").write_bytes(b"")
    transport = FakeTransport(cancel=True)
    settings = AgentSettings(artifact_path=tmp_path / "artifacts")
    with pytest.raises(AcquisitionCancelled):
        ArtifactAcquirer(settings, transport, "credential").execute(job.model_dump(mode="json"))
    assert not (root / "repositories" / "org--model" / ("a" * 40)).exists()


def test_execution_capacity_recheck_fails_before_network(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
    root = tmp_path / "artifacts" / "model-registry"
    job = lease(root, size=500)
    monkeypatch.setattr(
        "modelforge_node_agent.acquisition.shutil.disk_usage",
        lambda _path: shutil._ntuple_diskusage(1000, 950, 50),
    )
    settings = AgentSettings(artifact_path=tmp_path / "artifacts")
    with pytest.raises(AcquisitionFailure, match="capacity preflight"):
        ArtifactAcquirer(settings, FakeTransport(), "credential").execute(
            job.model_dump(mode="json")
        )


def test_size_mismatch_never_promotes_partial(
    tmp_path: Path, monkeypatch: pytest.MonkeyPatch
) -> None:
    root = tmp_path / "artifacts" / "model-registry"
    job = lease(root, size=8)

    def handler(request: httpx.Request) -> httpx.Response:
        return httpx.Response(200, content=b"short", request=request)

    original_client = httpx.Client
    monkeypatch.setattr(
        "modelforge_node_agent.acquisition.httpx.Client",
        lambda **kwargs: original_client(transport=httpx.MockTransport(handler), **kwargs),
    )
    settings = AgentSettings(artifact_path=tmp_path / "artifacts")
    with pytest.raises(AcquisitionFailure, match="size mismatch"):
        ArtifactAcquirer(settings, FakeTransport(), "credential").execute(
            job.model_dump(mode="json")
        )
    assert not (root / "repositories" / "org--model" / ("a" * 40)).exists()