"""Artifact registry: content-addressed hashing, atomic JSON persistence."""

from __future__ import annotations

import hashlib
import json
import logging
import os
import re
import subprocess
from datetime import UTC, datetime
from pathlib import Path

from vmcluster_mcp.models import ArtifactDiff, BuildArtifact, ErrorDetail, ToolResult

logger = logging.getLogger("vmcluster_mcp.artifact")

_EXCLUDE_DIRS = {"__pycache__", ".git"}


def hash_tree(root: str) -> str:
    """Compute a deterministic sha256 hash of all files in a directory tree.

    Excludes __pycache__ and .git directories.
    Each file contributes "relative_path:file_sha256" to the combined hash.
    Files are processed in sorted order for determinism.
    """
    parts: list[str] = []
    for dirpath, dirnames, filenames in os.walk(root):
        dirnames[:] = sorted(d for d in dirnames if d not in _EXCLUDE_DIRS)
        for filename in sorted(filenames):
            filepath = os.path.join(dirpath, filename)
            rel_path = os.path.relpath(filepath, root)
            file_hash = hashlib.sha256()
            with open(filepath, "rb") as f:
                while chunk := f.read(8192):
                    file_hash.update(chunk)
            parts.append(f"{rel_path}:{file_hash.hexdigest()}")
    return hashlib.sha256("\n".join(parts).encode()).hexdigest()


def _detect_kernel_version(root: str) -> str | None:
    """Try to detect kernel version from include/generated/utsrelease.h."""
    utsrelease_path = os.path.join(root, "include", "generated", "utsrelease.h")
    if not os.path.isfile(utsrelease_path):
        return None
    with open(utsrelease_path) as f:
        content = f.read()
    match = re.search(r'#define\s+UTS_RELEASE\s+"([^"]+)"', content)
    return match.group(1) if match else None


def _scan_modules(root: str) -> list[str]:
    """Find all .ko files in the tree, return sorted relative paths."""
    root_path = Path(root)
    return sorted(str(p.relative_to(root_path)) for p in root_path.rglob("*.ko"))


def _scan_bins(root: str) -> list[str]:
    """Find ELF executable binaries (non-.ko) in the tree, return sorted relative paths."""
    root_path = Path(root)
    result: list[str] = []
    for p in sorted(root_path.rglob("*")):
        if not p.is_file():
            continue
        if p.suffix == ".ko":
            continue
        try:
            with p.open("rb") as f:
                magic = f.read(4)
            if magic == b"\x7fELF":
                result.append(str(p.relative_to(root_path)))
        except OSError:
            pass
    return result


def _hash_files(root: str) -> dict[str, str]:
    """Return {relative_path: sha256_hex} for all files in root (excluding hidden dirs)."""
    result: dict[str, str] = {}
    for dirpath, dirnames, filenames in os.walk(root):
        dirnames[:] = sorted(d for d in dirnames if d not in _EXCLUDE_DIRS)
        for filename in sorted(filenames):
            filepath = os.path.join(dirpath, filename)
            rel_path = os.path.relpath(filepath, root)
            fh = hashlib.sha256()
            with open(filepath, "rb") as f:
                while chunk := f.read(8192):
                    fh.update(chunk)
            result[rel_path] = fh.hexdigest()
    return result


class ArtifactRegistry:
    """Content-addressed artifact registry with atomic JSON persistence."""

    def __init__(self, registry_path: str) -> None:
        self._path = registry_path
        self._registry: dict[str, dict] = {}

    def load(self) -> None:
        """Load registry from disk. If file missing, start with empty registry."""
        if not os.path.isfile(self._path):
            self._registry = {}
            return
        with open(self._path) as f:
            self._registry = json.load(f)

    def save(self) -> None:
        """Atomically write registry to disk using tmp + os.replace."""
        parent = os.path.dirname(self._path)
        if parent:
            os.makedirs(parent, exist_ok=True)
        tmp_path = f"{self._path}.tmp"
        with open(tmp_path, "w") as f:
            json.dump(self._registry, f, indent=2)
        os.replace(tmp_path, self._path)

    def register(
        self,
        source_path: str,
        build_type: str = "raw_tree",
        kernel_version: str | None = None,
        metadata: dict[str, str] | None = None,
    ) -> BuildArtifact:
        """Register a build tree. Returns existing artifact if already registered (dedup)."""
        artifact_id = hash_tree(source_path)

        if artifact_id in self._registry:
            return BuildArtifact(**self._registry[artifact_id])

        if kernel_version is None:
            kernel_version = _detect_kernel_version(source_path) or "unknown"

        modules = _scan_modules(source_path)
        userspace_bins = _scan_bins(source_path)

        artifact = BuildArtifact(
            artifact_id=artifact_id,
            source_path=source_path,
            build_type=build_type,  # type: ignore[arg-type]
            kernel_version=kernel_version,
            modules=modules,
            userspace_bins=userspace_bins,
            build_timestamp=datetime.now(tz=UTC).isoformat(),
            metadata=metadata or {},
        )

        self._registry[artifact_id] = artifact.model_dump()
        self.save()
        return artifact

    def list_artifacts(self) -> list[BuildArtifact]:
        """Return all registered artifacts."""
        return [BuildArtifact(**v) for v in self._registry.values()]

    def get(self, artifact_id: str) -> BuildArtifact | None:
        """Look up an artifact by ID. Returns None if not found."""
        data = self._registry.get(artifact_id)
        return BuildArtifact(**data) if data is not None else None

    def diff(self, id_a: str, id_b: str) -> ArtifactDiff | ToolResult:
        """Compare two artifacts by re-hashing their source trees.

        Returns ArtifactDiff on success, ToolResult(success=False) if an artifact
        ID is not found, or raises FileNotFoundError if a source_path is missing.
        """
        art_a = self.get(id_a)
        if art_a is None:
            return ToolResult(
                success=False,
                error=ErrorDetail(
                    code="ARTIFACT_NOT_FOUND",
                    message=f"Artifact '{id_a}' not found in registry",
                    recoverable=False,
                ),
            )
        art_b = self.get(id_b)
        if art_b is None:
            return ToolResult(
                success=False,
                error=ErrorDetail(
                    code="ARTIFACT_NOT_FOUND",
                    message=f"Artifact '{id_b}' not found in registry",
                    recoverable=False,
                ),
            )

        if not os.path.isdir(art_a.source_path):
            raise FileNotFoundError(f"Source path not found: {art_a.source_path}")
        if not os.path.isdir(art_b.source_path):
            raise FileNotFoundError(f"Source path not found: {art_b.source_path}")

        files_a = _hash_files(art_a.source_path)
        files_b = _hash_files(art_b.source_path)

        all_files = set(files_a.keys()) | set(files_b.keys())
        changed_modules: list[str] = []
        changed_bins: list[str] = []

        for f in sorted(all_files):
            if files_a.get(f) != files_b.get(f):
                if f.endswith(".ko"):
                    changed_modules.append(f)
                else:
                    changed_bins.append(f)

        git_log: str | None = None
        if art_a.git_sha and art_b.git_sha:
            try:
                proc = subprocess.run(
                    ["git", "-C", art_a.source_path, "log",
                     f"{art_a.git_sha}..{art_b.git_sha}", "--oneline"],
                    capture_output=True,
                    text=True,
                    timeout=10,
                )
                git_log = proc.stdout.strip() if proc.returncode == 0 else None
            except Exception:
                git_log = None

        return ArtifactDiff(
            artifact_id_a=id_a,
            artifact_id_b=id_b,
            changed_modules=changed_modules,
            changed_bins=changed_bins,
            git_log=git_log,
        )
