"""Batch manifest creation and provenance capture."""

from __future__ import annotations

import importlib
import importlib.metadata
import platform
import subprocess
import sys
from pathlib import Path
from typing import Any, Mapping

import numpy as np

from .hashing import file_sha256, stable_hash
from .io import atomic_write_json, utc_now
from .plan import validate_trial_plan
from .schema import SCHEMA_VERSION, STORAGE_FORMATS, ContractError, require_schema


_SOURCE_FILES = (
    Path("pyproject.toml"),
    Path("uv.lock"),
    Path("code/config/master_7dof.urdf"),
    Path("code/config/real_slave_7dof.urdf"),
)


def _git_provenance(cwd: Path) -> dict[str, Any]:
    try:
        commit = subprocess.run(
            ["git", "rev-parse", "HEAD"],
            cwd=cwd,
            check=True,
            capture_output=True,
            text=True,
            timeout=5,
        ).stdout.strip()
        status = subprocess.run(
            ["git", "status", "--porcelain"],
            cwd=cwd,
            check=True,
            capture_output=True,
            text=True,
            timeout=5,
        ).stdout
        return {"available": True, "commit": commit, "dirty": bool(status.strip())}
    except (OSError, subprocess.SubprocessError):
        return {"available": False, "commit": None, "dirty": None}


def _source_file_hashes(source_root: Path) -> dict[str, dict[str, str]]:
    """Hash reproducibility-critical files without requiring every file."""
    hashes: dict[str, dict[str, str]] = {}
    for relative_path in _SOURCE_FILES:
        path = source_root / relative_path
        if path.is_file():
            hashes[relative_path.as_posix()] = {"sha256": file_sha256(path)}
    return hashes


def _optional_module_version(module_name: str) -> str | None:
    """Return an importable module's version, including distribution fallback."""
    try:
        module = importlib.import_module(module_name)
    except Exception:
        return None

    version = getattr(module, "__version__", None)
    if version is not None:
        return str(version)

    try:
        distributions = importlib.metadata.packages_distributions().get(
            module_name, ()
        )
        for distribution in distributions:
            try:
                return importlib.metadata.version(distribution)
            except importlib.metadata.PackageNotFoundError:
                continue
    except (ImportError, OSError):
        pass
    return "unknown"


def _runtime_provenance() -> dict[str, str]:
    runtime = {
        "python": sys.version.split()[0],
        "numpy": np.__version__,
        "platform": platform.platform(),
        "processor": platform.processor(),
    }
    for module_name in ("scipy", "pinocchio"):
        version = _optional_module_version(module_name)
        if version is not None:
            runtime[module_name] = version
    return runtime


def _require_locked_git(
    plan: Mapping[str, Any],
    provenance: Mapping[str, Any],
    source_root: Path,
) -> None:
    """Reject evidence-locked runs that cannot identify immutable source."""
    if plan.get("split") != "locked":
        return
    if not provenance.get("available"):
        raise ContractError(
            "locked split requires available Git metadata; "
            f"source_root={source_root} is not a usable Git worktree"
        )
    if provenance.get("dirty"):
        raise ContractError(
            "locked split requires a clean Git worktree; "
            "commit or stash all source changes before running"
        )


def _require_locked_resume_match(
    stored_manifest: Mapping[str, Any],
    current_provenance: Mapping[str, Any],
    source_root: Path,
) -> None:
    """Require a resumed locked batch to use its original immutable source."""
    stored_source = stored_manifest.get("source")
    if not isinstance(stored_source, Mapping):
        raise ContractError("locked batch manifest is missing source provenance")

    stored_commit = stored_source.get("commit")
    current_commit = current_provenance.get("commit")
    if not stored_commit or current_commit != stored_commit:
        raise ContractError(
            "locked batch Git commit mismatch; "
            f"stored={stored_commit!r}, current={current_commit!r}"
        )

    stored_files = stored_source.get("files")
    if not isinstance(stored_files, Mapping):
        raise ContractError(
            "locked batch manifest is missing reproducibility file hashes"
        )
    current_files = _source_file_hashes(source_root)
    if dict(stored_files) != current_files:
        changed_paths = sorted(
            relative_path
            for relative_path in set(stored_files) | set(current_files)
            if stored_files.get(relative_path) != current_files.get(relative_path)
        )
        raise ContractError(
            "locked batch source file hash mismatch for: "
            + ", ".join(changed_paths)
        )


def build_batch_manifest(
    plan: Mapping[str, Any],
    *,
    source_root: Path | None = None,
) -> dict[str, Any]:
    validate_trial_plan(plan)
    source_root = Path.cwd() if source_root is None else Path(source_root)
    source = _git_provenance(source_root)
    _require_locked_git(plan, source, source_root)
    source["files"] = _source_file_hashes(source_root)
    batch_id = f"batch-{str(plan['plan_hash'])[:16]}"
    manifest = {
        "kind": "batch_manifest",
        "schema_version": SCHEMA_VERSION,
        "batch_id": batch_id,
        "study_id": plan["study_id"],
        "split": plan["split"],
        "plan_hash": plan["plan_hash"],
        "created_utc": utc_now(),
        "expected_pair_count": plan["pair_count"],
        "expected_trial_count": plan["trial_count"],
        "storage_formats": dict(STORAGE_FORMATS),
        "runtime": _runtime_provenance(),
        "source": source,
    }
    hash_basis = dict(manifest)
    hash_basis.pop("created_utc")
    manifest["manifest_hash"] = stable_hash(hash_basis, prefix="batch-manifest")
    return manifest


def ensure_batch(
    batch_dir: Path,
    plan: Mapping[str, Any],
    *,
    source_root: Path | None = None,
) -> dict[str, Any]:
    """Create or verify the immutable plan and batch manifest."""
    validate_trial_plan(plan)
    batch_dir = Path(batch_dir)
    source_root = Path.cwd() if source_root is None else Path(source_root)
    manifest_path = batch_dir / "batch_manifest.json"
    plan_path = batch_dir / "plan.json"
    if manifest_path.exists() or plan_path.exists():
        if not manifest_path.exists() or not plan_path.exists():
            raise ContractError("batch has only one of plan.json/batch_manifest.json")
        current_source = _git_provenance(source_root)
        _require_locked_git(plan, current_source, source_root)
        from .plan import load_document

        stored_plan = load_document(plan_path)
        stored_manifest = load_document(manifest_path)
        validate_trial_plan(stored_plan)
        require_schema(stored_manifest, "batch_manifest")
        if plan.get("split") == "locked":
            _require_locked_resume_match(
                stored_manifest,
                current_source,
                source_root,
            )
        if stored_plan["plan_hash"] != plan["plan_hash"]:
            raise ContractError("existing batch was created from a different plan")
        if stored_manifest.get("plan_hash") != plan["plan_hash"]:
            raise ContractError("batch manifest plan_hash mismatch")
        return stored_manifest

    batch_dir.mkdir(parents=True, exist_ok=True)
    manifest = build_batch_manifest(plan, source_root=source_root)
    atomic_write_json(plan_path, plan)
    atomic_write_json(manifest_path, manifest)
    return manifest