from __future__ import annotations

import argparse
import hashlib
import json
import tarfile
from pathlib import Path

import pytest

from pyro_mcp.runtime_build import (
    _build_paths,
    _load_lock,
    build_bundle,
    export_environment_oci_layout,
    generate_manifest,
    main,
    materialize_binaries,
    stage_agent,
    stage_binaries,
    stage_kernel,
    stage_rootfs,
    validate_sources,
)


def _write_text(path: Path, content: str) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(content, encoding="utf-8")


def _make_source_tree(tmp_path: Path) -> tuple[Path, Path, Path]:
    source_dir = tmp_path / "runtime_sources"
    platform_root = source_dir / "linux-x86_64"
    _write_text(source_dir / "NOTICE", "notice\n")
    _write_text(platform_root / "bin/firecracker", "firecracker\n")
    _write_text(platform_root / "bin/jailer", "jailer\n")
    _write_text(platform_root / "guest/pyro_guest_agent.py", "#!/usr/bin/env python3\n")
    _write_text(platform_root / "profiles/debian-base/vmlinux", "kernel-base\n")
    _write_text(platform_root / "profiles/debian-base/rootfs.ext4", "rootfs-base\n")
    _write_text(platform_root / "profiles/debian-git/vmlinux", "kernel-git\n")
    _write_text(platform_root / "profiles/debian-git/rootfs.ext4", "rootfs-git\n")
    _write_text(platform_root / "profiles/debian-build/vmlinux", "kernel-build\n")
    _write_text(platform_root / "profiles/debian-build/rootfs.ext4", "rootfs-build\n")
    lock = {
        "bundle_version": "9.9.9",
        "platform": "linux-x86_64",
        "component_versions": {
            "firecracker": "1.0.0",
            "jailer": "1.0.0",
            "kernel": "6.0.0",
            "guest_agent": "0.2.0",
            "base_distro": "debian-12",
        },
        "capabilities": {"vm_boot": True, "guest_exec": True, "guest_network": True},
        "binaries": {"firecracker": "bin/firecracker", "jailer": "bin/jailer"},
        "guest": {"agent": {"path": "guest/pyro_guest_agent.py"}},
        "profiles": {
            "debian-base": {
                "description": "base",
                "kernel": "profiles/debian-base/vmlinux",
                "rootfs": "profiles/debian-base/rootfs.ext4",
            },
            "debian-git": {
                "description": "git",
                "kernel": "profiles/debian-git/vmlinux",
                "rootfs": "profiles/debian-git/rootfs.ext4",
            },
            "debian-build": {
                "description": "build",
                "kernel": "profiles/debian-build/vmlinux",
                "rootfs": "profiles/debian-build/rootfs.ext4",
            },
        },
    }
    (platform_root / "runtime.lock.json").write_text(
        json.dumps(lock, indent=2) + "\n", encoding="utf-8"
    )
    build_dir = tmp_path / "build/runtime_bundle"
    bundle_dir = tmp_path / "bundle_out"
    return source_dir, build_dir, bundle_dir


def test_runtime_build_stages_and_manifest(tmp_path: Path) -> None:
    source_dir, build_dir, bundle_dir = _make_source_tree(tmp_path)
    paths = _build_paths(
        source_dir=source_dir,
        build_dir=build_dir,
        bundle_dir=bundle_dir,
        materialized_dir=tmp_path / "materialized_sources",
        platform="linux-x86_64",
    )
    lock = _load_lock(paths)

    paths.build_platform_root.mkdir(parents=True, exist_ok=True)
    stage_binaries(paths, lock)
    stage_kernel(paths, lock)
    stage_rootfs(paths, lock)
    stage_agent(paths, lock)
    manifest = generate_manifest(paths, lock)

    assert manifest["bundle_version"] == "9.9.9"
    assert manifest["capabilities"]["guest_exec"] is True
    assert manifest["component_versions"]["guest_agent"] == "0.2.0"
    assert (paths.build_platform_root / "guest/pyro_guest_agent.py").exists()


def test_runtime_build_bundle_syncs_output(tmp_path: Path) -> None:
    source_dir, build_dir, bundle_dir = _make_source_tree(tmp_path)
    paths = _build_paths(
        source_dir=source_dir,
        build_dir=build_dir,
        bundle_dir=bundle_dir,
        materialized_dir=tmp_path / "materialized_sources",
        platform="linux-x86_64",
    )

    manifest = build_bundle(paths, sync=True)

    assert manifest["profiles"]["debian-git"]["description"] == "git"
    assert (bundle_dir / "NOTICE").exists()
    assert (bundle_dir / "linux-x86_64/manifest.json").exists()
    assert (bundle_dir / "linux-x86_64/guest/pyro_guest_agent.py").exists()


def test_runtime_build_rejects_guest_capabilities_for_placeholder_sources(tmp_path: Path) -> None:
    source_dir, build_dir, bundle_dir = _make_source_tree(tmp_path)
    paths = _build_paths(
        source_dir=source_dir,
        build_dir=build_dir,
        bundle_dir=bundle_dir,
        materialized_dir=tmp_path / "materialized_sources",
        platform="linux-x86_64",
    )
    lock_path = source_dir / "linux-x86_64/runtime.lock.json"
    _write_text(
        source_dir / "linux-x86_64/bin/firecracker",
        "#!/usr/bin/env bash\n"
        "echo 'bundled firecracker shim'\n",
    )
    _write_text(
        source_dir / "linux-x86_64/profiles/debian-base/rootfs.ext4",
        "placeholder-rootfs\n",
    )
    lock = json.loads(lock_path.read_text(encoding="utf-8"))
    lock["capabilities"] = {"vm_boot": True, "guest_exec": True, "guest_network": True}
    lock_path.write_text(json.dumps(lock, indent=2) + "\n", encoding="utf-8")

    with pytest.raises(RuntimeError, match="guest-capable features"):
        validate_sources(paths, _load_lock(paths))


def test_runtime_build_materializes_firecracker_release(tmp_path: Path) -> None:
    source_dir, build_dir, bundle_dir = _make_source_tree(tmp_path)
    archive_path = tmp_path / "firecracker-v1.12.1-x86_64.tgz"
    release_dir = "release-v1.12.1-x86_64"
    with tarfile.open(archive_path, "w:gz") as archive:
        firecracker_path = tmp_path / "firecracker-bin"
        jailer_path = tmp_path / "jailer-bin"
        firecracker_path.write_text("real-firecracker\n", encoding="utf-8")
        jailer_path.write_text("real-jailer\n", encoding="utf-8")
        archive.add(firecracker_path, arcname=f"{release_dir}/firecracker-v1.12.1-x86_64")
        archive.add(jailer_path, arcname=f"{release_dir}/jailer-v1.12.1-x86_64")

    digest = hashlib.sha256(archive_path.read_bytes()).hexdigest()
    lock_path = source_dir / "linux-x86_64/runtime.lock.json"
    lock = json.loads(lock_path.read_text(encoding="utf-8"))
    lock["upstream"] = {
        "firecracker_release": {
            "archive_url": archive_path.resolve().as_uri(),
            "archive_sha256": digest,
            "firecracker_member": f"{release_dir}/firecracker-v1.12.1-x86_64",
            "jailer_member": f"{release_dir}/jailer-v1.12.1-x86_64",
        }
    }
    lock_path.write_text(json.dumps(lock, indent=2) + "\n", encoding="utf-8")

    paths = _build_paths(
        source_dir=source_dir,
        build_dir=build_dir,
        bundle_dir=bundle_dir,
        materialized_dir=tmp_path / "materialized_sources",
        platform="linux-x86_64",
    )
    materialize_binaries(paths, _load_lock(paths))

    assert (paths.materialized_platform_root / "bin/firecracker").read_text(encoding="utf-8") == (
        "real-firecracker\n"
    )
    assert (paths.materialized_platform_root / "bin/jailer").read_text(encoding="utf-8") == (
        "real-jailer\n"
    )


def test_export_environment_oci_layout_writes_boot_artifacts(tmp_path: Path) -> None:
    source_dir, build_dir, bundle_dir = _make_source_tree(tmp_path)
    paths = _build_paths(
        source_dir=source_dir,
        build_dir=build_dir,
        bundle_dir=bundle_dir,
        materialized_dir=tmp_path / "materialized_sources",
        platform="linux-x86_64",
    )

    result = export_environment_oci_layout(
        paths,
        environment="debian:12-base",
        output_dir=tmp_path / "oci_layouts",
    )

    layout_dir = Path(str(result["layout_dir"]))
    assert layout_dir == tmp_path / "oci_layouts" / "debian_12-base"
    assert json.loads((layout_dir / "oci-layout").read_text(encoding="utf-8")) == {
        "imageLayoutVersion": "1.0.0"
    }
    index = json.loads((layout_dir / "index.json").read_text(encoding="utf-8"))
    manifest_descriptor = index["manifests"][0]
    assert manifest_descriptor["annotations"]["org.opencontainers.image.ref.name"] == "1.0.0"
    assert manifest_descriptor["platform"] == {"os": "linux", "architecture": "amd64"}

    manifest_digest = str(result["manifest_digest"]).split(":", 1)[1]
    manifest = json.loads(
        (layout_dir / "blobs" / "sha256" / manifest_digest).read_text(encoding="utf-8")
    )
    assert manifest["config"]["mediaType"] == "application/vnd.oci.image.config.v1+json"
    assert len(manifest["layers"]) == 3

    extracted_files: dict[str, str] = {}
    for layer in manifest["layers"]:
        layer_digest = str(layer["digest"]).split(":", 1)[1]
        layer_path = layout_dir / "blobs" / "sha256" / layer_digest
        with tarfile.open(layer_path, "r:") as archive:
            members = archive.getmembers()
            assert len(members) == 1
            member = members[0]
            extracted = archive.extractfile(member)
            if extracted is None:
                raise AssertionError("expected layer member content")
            extracted_files[member.name] = extracted.read().decode("utf-8")

    assert extracted_files["vmlinux"] == "kernel-base\n"
    assert extracted_files["rootfs.ext4"] == "rootfs-base\n"
    metadata = json.loads(extracted_files["environment.json"])
    assert metadata["environment"] == "debian:12-base"
    assert metadata["source_profile"] == "debian-base"


def test_export_environment_oci_layout_rejects_missing_profile(tmp_path: Path) -> None:
    source_dir, build_dir, bundle_dir = _make_source_tree(tmp_path)
    paths = _build_paths(
        source_dir=source_dir,
        build_dir=build_dir,
        bundle_dir=bundle_dir,
        materialized_dir=tmp_path / "materialized_sources",
        platform="linux-x86_64",
    )
    lock_path = source_dir / "linux-x86_64/runtime.lock.json"
    lock = json.loads(lock_path.read_text(encoding="utf-8"))
    del lock["profiles"]["debian-base"]
    lock_path.write_text(json.dumps(lock, indent=2) + "\n", encoding="utf-8")

    with pytest.raises(RuntimeError, match="does not define source profile"):
        export_environment_oci_layout(
            paths,
            environment="debian:12-base",
            output_dir=tmp_path / "oci_layouts",
        )


def test_runtime_build_main_dispatches_export(
    monkeypatch: pytest.MonkeyPatch,
    capsys: pytest.CaptureFixture[str],
    tmp_path: Path,
) -> None:
    source_dir, build_dir, bundle_dir = _make_source_tree(tmp_path)
    runtime_paths = _build_paths(
        source_dir=source_dir,
        build_dir=build_dir,
        bundle_dir=bundle_dir,
        materialized_dir=tmp_path / "materialized",
        platform="linux-x86_64",
    )

    class StubParser:
        def parse_args(self) -> argparse.Namespace:
            return argparse.Namespace(
                command="export-environment-oci",
                platform="linux-x86_64",
                source_dir=str(source_dir),
                build_dir=str(build_dir),
                bundle_dir=str(bundle_dir),
                materialized_dir=str(tmp_path / "materialized"),
                environment="debian:12-base",
                output_dir=str(tmp_path / "oci_layouts"),
                reference="1.0.0",
            )

    monkeypatch.setattr("pyro_mcp.runtime_build._build_parser", lambda: StubParser())
    monkeypatch.setattr(
        "pyro_mcp.runtime_build._load_lock",
        lambda paths: _load_lock(runtime_paths),
    )
    monkeypatch.setattr(
        "pyro_mcp.runtime_build.export_environment_oci_layout",
        lambda paths, **kwargs: {
            "environment": kwargs["environment"],
            "layout_dir": str(Path(kwargs["output_dir"]) / "debian_12-base"),
            "manifest_digest": "sha256:test",
        },
    )

    main()

    payload = json.loads(capsys.readouterr().out)
    assert payload["environment"] == "debian:12-base"
    assert payload["manifest_digest"] == "sha256:test"