"""Artifact lock tracking and verification helpers for .osmoimg workflows."""

from __future__ import annotations

import json
import subprocess
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Dict, List, Optional, Sequence


ARTIFACT_LOCKFILE_DEFAULT = Path(".model_osmosis") / "artifacts.lock.json"
ARTIFACT_HOOK_MARKER = "# model-osmosis-artifact-hook-v1"


def _utc_now() -> str:
    return datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ")


def _artifact_lock_path(path_override: Optional[str], *, default_lock_file: Path) -> Path:
    if path_override:
        return Path(path_override)
    return Path(default_lock_file)


def _normalize_artifact_path(path: Path, *, root: Path) -> str:
    resolved = path.resolve()
    try:
        rel = resolved.relative_to(root.resolve())
        return rel.as_posix()
    except Exception:
        return str(resolved)


def _tracked_artifact_abs_path(tracked_path: str, *, root: Path) -> Path:
    path_obj = Path(tracked_path)
    if path_obj.is_absolute():
        return path_obj
    return (root / path_obj).resolve()


def _load_artifact_lock(lock_path: Path) -> Dict[str, Any]:
    if not lock_path.exists():
        return {
            "schema_version": 1,
            "updated_at": "",
            "artifacts": [],
        }
    try:
        payload = json.loads(lock_path.read_text(encoding="utf-8"))
    except Exception as exc:
        raise ValueError(f"Invalid artifact lock file JSON: {lock_path}") from exc
    if not isinstance(payload, dict):
        raise ValueError("Artifact lock file must contain a JSON object")
    artifacts_raw = payload.get("artifacts", [])
    if not isinstance(artifacts_raw, list):
        raise ValueError("Artifact lock file field `artifacts` must be a JSON array")
    artifacts: List[Dict[str, Any]] = []
    for item in artifacts_raw:
        if not isinstance(item, dict):
            continue
        path_value = str(item.get("path", "") or "").strip()
        if not path_value:
            continue
        normalized = dict(item)
        normalized["path"] = path_value.replace("\\", "/")
        artifacts.append(normalized)
    artifacts = sorted(artifacts, key=lambda x: str(x.get("path", "")))
    return {
        "schema_version": int(payload.get("schema_version", 1) or 1),
        "updated_at": str(payload.get("updated_at", "") or ""),
        "artifacts": artifacts,
    }


def _write_artifact_lock(lock_path: Path, payload: Dict[str, Any]) -> Path:
    lock_path.parent.mkdir(parents=True, exist_ok=True)
    normalized = dict(payload)
    normalized["schema_version"] = int(payload.get("schema_version", 1) or 1)
    normalized["updated_at"] = str(payload.get("updated_at", _utc_now()) or _utc_now())
    artifacts = payload.get("artifacts", [])
    if not isinstance(artifacts, list):
        artifacts = []
    normalized["artifacts"] = sorted(
        [
            {k: v for k, v in dict(item).items()}
            for item in artifacts
            if isinstance(item, dict) and str(item.get("path", "")).strip()
        ],
        key=lambda x: str(x.get("path", "")),
    )
    lock_path.write_text(
        json.dumps(normalized, indent=2, sort_keys=True) + "\n",
        encoding="utf-8",
    )
    return lock_path


def _collect_artifact_record(image_path: Path, *, root: Path) -> Dict[str, Any]:
    if not image_path.exists():
        raise FileNotFoundError(f"Artifact not found: {image_path}")

    from model_osmosis.format.image import OsmoImage

    verify = OsmoImage.verify(str(image_path), require_trust=False)
    inspect = OsmoImage.inspect(str(image_path))
    runtime = inspect.get("runtime", {}) or {}
    return {
        "path": _normalize_artifact_path(image_path, root=root),
        "file_size_bytes": int(image_path.stat().st_size),
        "header_version": int(verify.get("header_version", 1) or 1),
        "payload_sha256": str(verify.get("payload_sha256", "") or ""),
        "artifact_id": str(verify.get("artifact_id", "") or runtime.get("artifact_id", "") or ""),
        "manifest_sha256": str(
            verify.get("manifest_sha256", "") or runtime.get("manifest_sha256", "") or ""
        ),
        "tensor_index_sha256": str(
            verify.get("tensor_index_sha256", "") or runtime.get("tensor_index_sha256", "") or ""
        ),
        "payload_codec": str(verify.get("payload_codec", "") or runtime.get("payload_codec", "unknown")),
        "compression_profile": str(
            verify.get("compression_profile", "") or runtime.get("compression_profile", "legacy")
        ),
        "trust_status": str(verify.get("trust_status", "unsigned") or "unsigned"),
        "signer_id": str(verify.get("signer_id") or ""),
        "key_id": str(verify.get("key_id") or ""),
        "signature_alg": str(verify.get("signature_alg") or ""),
    }


def _compare_artifact_records(expected: Dict[str, Any], current: Dict[str, Any]) -> Dict[str, Dict[str, Any]]:
    fields = [
        "header_version",
        "payload_sha256",
        "artifact_id",
        "manifest_sha256",
        "tensor_index_sha256",
        "payload_codec",
        "compression_profile",
        "signer_id",
        "key_id",
        "signature_alg",
    ]
    mismatches: Dict[str, Dict[str, Any]] = {}
    for field in fields:
        exp = expected.get(field)
        cur = current.get(field)
        if exp != cur:
            mismatches[field] = {
                "expected": exp,
                "actual": cur,
            }
    return mismatches


def _resolve_tracked_artifact_key(
    requested_path: str,
    artifacts: Sequence[Dict[str, Any]],
    *,
    root: Path,
) -> Optional[str]:
    requested = Path(requested_path)
    requested_abs = requested if requested.is_absolute() else (root / requested)
    requested_abs = requested_abs.resolve()
    requested_norm = _normalize_artifact_path(requested_abs, root=root)

    known = {str(item.get("path", "")): item for item in artifacts}
    if requested_norm in known:
        return requested_norm
    if str(requested).replace("\\", "/") in known:
        return str(requested).replace("\\", "/")

    for key in sorted(known.keys()):
        candidate_abs = _tracked_artifact_abs_path(key, root=root).resolve()
        if candidate_abs == requested_abs:
            return key
    return None


def _resolve_git_hooks_dir() -> Path:
    proc = subprocess.run(
        ["git", "rev-parse", "--git-dir"],
        capture_output=True,
        text=True,
        timeout=8,
    )
    if proc.returncode != 0:
        raise RuntimeError((proc.stderr or proc.stdout or "").strip() or "Not a git repository")
    git_dir = Path((proc.stdout or "").strip())
    if not git_dir.is_absolute():
        git_dir = (Path.cwd() / git_dir).resolve()
    hooks_dir = git_dir / "hooks"
    hooks_dir.mkdir(parents=True, exist_ok=True)
    return hooks_dir


def track_artifacts(
    *,
    images: Sequence[str],
    lock_file: Optional[str] = None,
    root: Optional[Path] = None,
    default_lock_file: Path = ARTIFACT_LOCKFILE_DEFAULT,
) -> Dict[str, Any]:
    root_path = (root or Path.cwd()).resolve()
    lock_path = _artifact_lock_path(lock_file, default_lock_file=default_lock_file)
    lock_data = _load_artifact_lock(lock_path)
    by_path = {
        str(item.get("path", "")): dict(item)
        for item in (lock_data.get("artifacts", []) or [])
        if isinstance(item, dict) and str(item.get("path", "")).strip()
    }

    tracked_paths: List[str] = []
    for image in list(images or []):
        record = _collect_artifact_record(Path(image), root=root_path)
        by_path[str(record["path"])] = record
        tracked_paths.append(str(record["path"]))

    if not tracked_paths:
        raise RuntimeError("No artifacts were provided to track.")

    lock_payload = {
        "schema_version": 1,
        "updated_at": _utc_now(),
        "artifacts": sorted(by_path.values(), key=lambda x: str(x.get("path", ""))),
    }
    saved = _write_artifact_lock(lock_path, lock_payload)
    return {
        "lock_file": str(saved),
        "tracked_artifacts": sorted(tracked_paths),
        "tracked_count": len(tracked_paths),
    }


def verify_artifacts(
    *,
    image: Optional[str] = None,
    verify_all: bool = False,
    lock_file: Optional[str] = None,
    root: Optional[Path] = None,
    default_lock_file: Path = ARTIFACT_LOCKFILE_DEFAULT,
) -> Dict[str, Any]:
    root_path = (root or Path.cwd()).resolve()
    lock_path = _artifact_lock_path(lock_file, default_lock_file=default_lock_file)
    lock_data = _load_artifact_lock(lock_path)
    artifacts = [
        dict(item)
        for item in (lock_data.get("artifacts", []) or [])
        if isinstance(item, dict) and str(item.get("path", "")).strip()
    ]
    if not artifacts:
        raise RuntimeError(f"No tracked artifacts found in lock file: {lock_path}")

    requested_image = str(image or "").strip()
    if not verify_all and not requested_image:
        raise RuntimeError("Specify a tracked IMAGE or use --all.")

    targets: List[Dict[str, Any]] = []
    if verify_all:
        targets = sorted(artifacts, key=lambda x: str(x.get("path", "")))
    else:
        key = _resolve_tracked_artifact_key(requested_image, artifacts, root=root_path)
        if key is None:
            raise RuntimeError(f"Artifact is not tracked: {requested_image}")
        targets = [next(item for item in artifacts if str(item.get("path", "")) == key)]

    results: List[Dict[str, Any]] = []
    for expected in targets:
        tracked_path = str(expected.get("path", ""))
        abs_path = _tracked_artifact_abs_path(tracked_path, root=root_path)
        if not abs_path.exists():
            results.append(
                {
                    "path": tracked_path,
                    "passed": False,
                    "reason": "file_missing",
                    "mismatches": {},
                }
            )
            continue
        try:
            current = _collect_artifact_record(abs_path, root=root_path)
            mismatches = _compare_artifact_records(expected, current)
            results.append(
                {
                    "path": tracked_path,
                    "passed": len(mismatches) == 0,
                    "reason": "" if len(mismatches) == 0 else "metadata_mismatch",
                    "mismatches": mismatches,
                }
            )
        except Exception as exc:
            results.append(
                {
                    "path": tracked_path,
                    "passed": False,
                    "reason": f"inspect_failed:{type(exc).__name__}",
                    "mismatches": {},
                }
            )

    failed = [item for item in results if not bool(item.get("passed", False))]
    return {
        "schema_version": 1,
        "status": "PASS" if not failed else "FAIL",
        "lock_file": str(lock_path),
        "checked_artifacts": len(results),
        "failed_artifacts": len(failed),
        "results": sorted(results, key=lambda x: str(x.get("path", ""))),
    }


def install_artifact_hooks(
    *,
    lock_file: Optional[str] = None,
    force: bool = False,
    default_lock_file: Path = ARTIFACT_LOCKFILE_DEFAULT,
) -> Dict[str, Any]:
    hooks_dir = _resolve_git_hooks_dir()
    resolved_lock = _artifact_lock_path(lock_file, default_lock_file=default_lock_file)
    verify_cmd = "osmosis artifact verify --all"
    if lock_file:
        verify_cmd = f"{verify_cmd} --lock-file \"{lock_file}\""

    script = (
        "#!/bin/sh\n"
        f"{ARTIFACT_HOOK_MARKER}\n"
        "set -e\n"
        f"{verify_cmd}\n"
    )

    installed: List[str] = []
    for hook_name in ("pre-commit", "pre-push"):
        hook_path = hooks_dir / hook_name
        if hook_path.exists() and not force:
            existing = hook_path.read_text(encoding="utf-8", errors="ignore")
            if ARTIFACT_HOOK_MARKER not in existing:
                raise RuntimeError(
                    f"Hook already exists and is not managed by ModelOsmosis: {hook_path}"
                )
        hook_path.write_text(script, encoding="utf-8")
        try:
            hook_path.chmod(0o755)
        except Exception:
            pass
        installed.append(str(hook_path))

    return {
        "hooks_dir": str(hooks_dir),
        "verify_command": verify_cmd,
        "installed_hooks": installed,
        "lock_file": str(resolved_lock),
    }

