#!/usr/bin/env python3
"""Validate or stamp reproducible non-AI Python release locks."""

from __future__ import annotations

import argparse
import hashlib
import json
import re
import sys
import tomllib
from pathlib import Path


ROOT = Path(__file__).resolve().parents[1]
PYPROJECT_PATH = ROOT / "backend" / "pyproject.toml"
STAMP_PREFIX = "# geointel-input-sha256: "
LOCKED_PACKAGE_RE = re.compile(r"^([A-Za-z0-9_.-]+)==", re.MULTILINE)
LOCK_GROUPS = {
    "runtime": ("gis",),
    "ci": ("gis", "dev"),
}


def normalize_name(requirement: str) -> str:
    name = re.split(r"[\[<>=!~;@\s]", requirement, maxsplit=1)[0]
    return name.lower().replace("_", "-").replace(".", "-")


def lock_input(
    pyproject: dict[str, object], optional_groups: tuple[str, ...]
) -> dict[str, object]:
    project = pyproject["project"]
    assert isinstance(project, dict)
    optional = project.get("optional-dependencies", {})
    assert isinstance(optional, dict)
    result = {
        "requires-python": project.get("requires-python"),
        "dependencies": project.get("dependencies", []),
    }
    for group in optional_groups:
        result[group] = optional.get(group, [])
    return result


def input_digest(
    pyproject: dict[str, object], optional_groups: tuple[str, ...]
) -> str:
    payload = json.dumps(
        lock_input(pyproject, optional_groups),
        ensure_ascii=True,
        separators=(",", ":"),
        sort_keys=True,
    )
    return hashlib.sha256(payload.encode("utf-8")).hexdigest()


def stamp_lock(lock_text: str, digest: str) -> str:
    lines = [
        line for line in lock_text.splitlines() if not line.startswith(STAMP_PREFIX)
    ]
    insert_at = next(
        (index for index, line in enumerate(lines) if line and not line.startswith("#")),
        len(lines),
    )
    lines.insert(insert_at, f"{STAMP_PREFIX}{digest}")
    return "\n".join(lines) + "\n"


def validate_lock(
    pyproject: dict[str, object],
    lock_text: str,
    lock_name: str,
    optional_groups: tuple[str, ...],
) -> list[str]:
    errors: list[str] = []
    digest = input_digest(pyproject, optional_groups)
    if f"{STAMP_PREFIX}{digest}" not in lock_text:
        errors.append(
            f"{lock_name} lock input stamp is missing or stale; "
            "run scripts/generate_python_lock.sh"
        )
    if "pip-compile with Python 3.11" not in lock_text:
        errors.append("lock must be generated with the release Python 3.11 runtime")
    if "--generate-hashes" not in lock_text:
        errors.append("lock must contain pip hashes")

    project = pyproject["project"]
    assert isinstance(project, dict)
    optional = project.get("optional-dependencies", {})
    assert isinstance(optional, dict)
    required = {
        normalize_name(requirement)
        for group in (
            project.get("dependencies", []),
            *(optional.get(name, []) for name in optional_groups),
        )
        for requirement in group
    }
    locked = {
        normalize_name(package)
        for package in LOCKED_PACKAGE_RE.findall(lock_text)
    }
    missing = sorted(required - locked)
    if missing:
        errors.append(f"direct release dependencies missing from lock: {', '.join(missing)}")

    ai_packages = {
        normalize_name(requirement)
        for requirement in optional.get("ai", [])
    }
    leaked_ai = sorted(ai_packages & locked)
    if leaked_ai:
        errors.append(
            f"optional AI dependencies leaked into the {lock_name} lock: "
            + ", ".join(leaked_ai)
        )
    if lock_name == "runtime":
        dev_packages = {
            normalize_name(requirement)
            for requirement in optional.get("dev", [])
        }
        leaked_dev = sorted(dev_packages & locked)
        if leaked_dev:
            errors.append(
                "developer-only dependencies leaked into the runtime lock: "
                + ", ".join(leaked_dev)
            )

    package_blocks = re.split(r"\n(?=[A-Za-z0-9_.-]+==)", lock_text)
    unhashed = []
    for block in package_blocks:
        match = re.match(r"([A-Za-z0-9_.-]+)==", block)
        if match and "--hash=sha256:" not in block:
            unhashed.append(match.group(1))
    if unhashed:
        errors.append(f"locked packages without hashes: {', '.join(sorted(unhashed))}")
    return errors


def main() -> int:
    parser = argparse.ArgumentParser(
        description="Verify the GeoIntel Python 3.11 runtime and CI locks."
    )
    parser.add_argument(
        "--stamp",
        action="store_true",
        help="stamp a freshly generated lock with its canonical input digest",
    )
    parser.add_argument(
        "--lock",
        choices=("all", *LOCK_GROUPS),
        default="all",
        help="limit validation/stamping to one lock",
    )
    args = parser.parse_args()

    pyproject = tomllib.loads(PYPROJECT_PATH.read_text(encoding="utf-8"))
    selected = LOCK_GROUPS if args.lock == "all" else {args.lock: LOCK_GROUPS[args.lock]}
    errors: list[str] = []
    for lock_name, optional_groups in selected.items():
        lock_path = ROOT / "backend" / f"requirements-{lock_name}.lock"
        if not lock_path.exists():
            errors.append(f"missing lock: {lock_path.relative_to(ROOT)}")
            continue
        lock_text = lock_path.read_text(encoding="utf-8")
        if args.stamp:
            lock_text = stamp_lock(
                lock_text,
                input_digest(pyproject, optional_groups),
            )
            lock_path.write_text(lock_text, encoding="utf-8", newline="\n")
        errors.extend(
            validate_lock(
                pyproject,
                lock_text,
                lock_name,
                optional_groups,
            )
        )
    if errors:
        for error in errors:
            print(f"ERROR: {error}", file=sys.stderr)
        return 1
    print(
        "Python lock policy passed: Python 3.11, hashed runtime/CI locks, "
        "AI optional."
    )
    return 0


if __name__ == "__main__":
    raise SystemExit(main())