geointel/scripts/activate_promoted_yolo_candidate.py

from __future__ import annotations

import argparse
import hashlib
import json
import os
import re
import sys
from pathlib import Path
from typing import Any


SUPPORTED_MODEL_SUFFIXES = {".pt", ".onnx", ".engine"}
ENV_KEYS = ("GEOINTEL_INSTALL_AI", "YOLO_ENABLED", "YOLO_MODELS_DIR", "YOLO_MODEL_PATH")


def _sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as stream:
        for chunk in iter(lambda: stream.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def _asset_id(path: Path) -> str:
    raw = f"{path.stem}-{path.suffix.lower().lstrip('.')}"
    normalized = re.sub(r"[^a-z0-9]+", "-", raw.lower()).strip("-")
    return normalized or "model-asset"


def _candidate_paths(models_dir: Path) -> list[Path]:
    if not models_dir.exists() or not models_dir.is_dir():
        return []
    return sorted(
        path.resolve()
        for path in models_dir.rglob("*")
        if path.is_file() and path.suffix.lower() in SUPPORTED_MODEL_SUFFIXES
    )


def _container_path(host_model_path: Path, models_dir: Path, container_model_dir: str) -> str:
    relative = host_model_path.resolve().relative_to(models_dir.resolve())
    base = container_model_dir.rstrip("/")
    return f"{base}/{relative.as_posix()}" if relative.as_posix() else base


def _read_env_lines(env_file: Path) -> list[str]:
    if not env_file.exists():
        return []
    return env_file.read_text(encoding="utf-8").splitlines()


def _update_env_file(env_file: Path, updates: dict[str, str]) -> None:
    existing_lines = _read_env_lines(env_file)
    seen: set[str] = set()
    next_lines: list[str] = []

    for line in existing_lines:
        stripped = line.strip()
        if not stripped or stripped.startswith("#") or "=" not in line:
            next_lines.append(line)
            continue

        key = line.split("=", 1)[0].strip()
        if key in updates:
            next_lines.append(f"{key}={updates[key]}")
            seen.add(key)
        else:
            next_lines.append(line)

    for key, value in updates.items():
        if key not in seen:
            next_lines.append(f"{key}={value}")

    env_file.parent.mkdir(parents=True, exist_ok=True)
    env_file.write_text("\n".join(next_lines).rstrip() + "\n", encoding="utf-8")


def _load_report(path: Path) -> dict[str, Any]:
    if not path.exists() or not path.is_file():
        raise ValueError("Promotion report file does not exist")
    payload = json.loads(path.read_text(encoding="utf-8-sig"))
    if not isinstance(payload, dict):
        raise ValueError("Promotion report must be a JSON object")
    return payload


def _candidate_from_report(report: dict[str, Any], candidate_key: str) -> dict[str, Any] | None:
    recommended = report.get("recommended_candidate")
    if isinstance(recommended, dict):
        if recommended.get("candidate_key") == candidate_key:
            return recommended
        return None
    if isinstance(recommended, str) and recommended == candidate_key:
        for item in report.get("candidate_decisions") or []:
            if isinstance(item, dict) and item.get("candidate_key") == candidate_key:
                return item
    for item in report.get("candidate_decisions") or []:
        if isinstance(item, dict) and item.get("candidate_key") == candidate_key:
            return item
    return None


def _gate_failures(report: dict[str, Any], candidate: dict[str, Any]) -> list[str]:
    gates = report.get("gates") if isinstance(report.get("gates"), dict) else {}
    failures: list[str] = []
    if candidate.get("promotion_status") != "promote_candidate":
        failures.append("candidate_not_promoted")
    if candidate.get("rejection_reasons"):
        failures.append("candidate_has_rejection_reasons")

    min_positive_samples = int(gates.get("min_positive_samples") or 0)
    min_background_samples = int(gates.get("min_background_samples") or 0)
    min_mean_f1 = float(gates.get("min_mean_f1") or 0)
    max_background_detections = int(gates.get("max_background_detections_per_sample") or 0)

    if int(candidate.get("positive_sample_count") or 0) < min_positive_samples:
        failures.append("insufficient_positive_samples")
    if int(candidate.get("background_sample_count") or 0) < min_background_samples:
        failures.append("insufficient_background_samples")
    if float(candidate.get("mean_f1") or 0) < min_mean_f1:
        failures.append("positive_mean_f1_below_gate")
    if int(candidate.get("max_background_detections") or 0) > max_background_detections:
        failures.append("background_false_positive_pressure")
    return failures


def _governed_release_failures(
    release_gate: dict[str, Any],
    *,
    candidate_key: str,
    model_sha256: str,
) -> list[str]:
    """Require the Phase-4/5 product decision that the legacy gate never had."""

    failures: list[str] = []
    if release_gate.get("status") != "pass":
        failures.append("phase4_release_gate_not_passed")
    if release_gate.get("product_benchmark_status") != "pass":
        failures.append("product_benchmark_not_passed")
    if release_gate.get("promotion_allowed") is not True:
        failures.append("governed_promotion_not_allowed")
    if release_gate.get("phase_decision") not in {"done", "ready", "complete"}:
        failures.append("phase5_not_ready")
    if release_gate.get("candidate_key") != candidate_key:
        failures.append("governed_candidate_key_mismatch")
    recorded_sha256 = str(release_gate.get("candidate_model_sha256") or "").lower()
    if not re.fullmatch(r"[0-9a-f]{64}", recorded_sha256):
        failures.append("governed_candidate_model_sha256_missing")
    elif recorded_sha256 != model_sha256:
        failures.append("governed_candidate_model_sha256_mismatch")
    benchmark_sha256 = str(release_gate.get("benchmark_manifest_sha256") or "").lower()
    if not re.fullmatch(r"[0-9a-f]{64}", benchmark_sha256):
        failures.append("governed_benchmark_manifest_sha256_missing")
    return failures


def _resolve_model_asset(models_dir: Path, model_asset_id: str) -> Path | None:
    for path in _candidate_paths(models_dir):
        if _asset_id(path) == model_asset_id:
            return path
    return None


def _base_payload(args: argparse.Namespace) -> dict[str, Any]:
    return {
        "promotion_report": str(Path(args.promotion_report).resolve()),
        "phase4_release_gate_report": str(Path(args.phase4_release_gate_report).resolve()),
        "candidate_key": args.candidate_key,
        "models_dir": str(Path(args.models_dir).resolve()),
        "env_file": str(Path(args.env_file).resolve()),
        "candidate": None,
        "selected_host_model_path": None,
        "selected_container_model_path": None,
        "env_updates": {},
        "apply": args.apply,
        "applied": False,
        "docker_restart_required": False,
        "will_download_models": False,
        "will_run_inference": False,
    }


def activate(args: argparse.Namespace) -> tuple[int, dict[str, Any]]:
    payload = _base_payload(args)
    try:
        report = _load_report(Path(args.promotion_report))
    except (OSError, ValueError, json.JSONDecodeError) as exc:
        payload.update({"status": "invalid_promotion_report", "message": str(exc)})
        return 2, payload

    candidate = _candidate_from_report(report, args.candidate_key)
    if not candidate:
        payload.update(
            {
                "status": "candidate_not_recommended",
                "message": "Candidate key does not match the report recommended candidate.",
            }
        )
        return 3, payload

    payload["candidate"] = candidate
    failures = _gate_failures(report, candidate)
    if failures:
        payload.update(
            {
                "status": "candidate_not_promoted",
                "message": "Candidate did not pass the promotion gates.",
                "rejection_reasons": failures,
            }
        )
        return 3, payload

    model_asset_id = str(candidate.get("model_asset_id") or "").strip()
    if not model_asset_id:
        payload.update({"status": "model_asset_missing", "message": "Candidate does not record model_asset_id."})
        return 2, payload

    models_dir = Path(args.models_dir).resolve()
    selected_model = _resolve_model_asset(models_dir, model_asset_id)
    if selected_model is None:
        payload.update(
            {
                "status": "model_asset_not_found",
                "message": "Promoted candidate model asset was not found in the local models directory.",
                "model_asset_id": model_asset_id,
            }
        )
        return 2, payload

    model_sha256 = _sha256(selected_model)
    try:
        release_gate_path = Path(args.phase4_release_gate_report)
        release_gate = _load_report(release_gate_path)
    except (OSError, ValueError, json.JSONDecodeError) as exc:
        payload.update({"status": "invalid_phase4_release_gate", "message": str(exc)})
        return 2, payload
    governance_failures = _governed_release_failures(
        release_gate,
        candidate_key=args.candidate_key,
        model_sha256=model_sha256,
    )
    if governance_failures:
        payload.update(
            {
                "status": "governed_release_not_approved",
                "message": (
                    "Candidate passed the legacy diagnostic gates but lacks a matching "
                    "successful governed product benchmark and Phase-4/5 release decision."
                ),
                "rejection_reasons": governance_failures,
                "selected_host_model_path": str(selected_model),
                "selected_model_sha256": model_sha256,
            }
        )
        return 3, payload

    selected_container_path = _container_path(selected_model, models_dir, args.container_model_dir)
    updates = {
        "GEOINTEL_INSTALL_AI": "true",
        "YOLO_ENABLED": "true",
        "YOLO_MODELS_DIR": args.container_model_dir.rstrip("/"),
        "YOLO_MODEL_PATH": selected_container_path,
    }
    payload.update(
        {
            "status": "ready_to_apply",
            "message": "Promoted YOLO candidate validated. Re-run with --apply to update the environment file.",
            "selected_host_model_path": str(selected_model),
            "selected_model_sha256": model_sha256,
            "selected_container_model_path": selected_container_path,
            "phase4_release_gate_sha256": _sha256(release_gate_path),
            "env_updates": updates,
            "docker_restart_required": True,
        }
    )

    if args.apply:
        _update_env_file(Path(args.env_file).resolve(), updates)
        payload.update(
            {
                "status": "applied",
                "message": "Environment file updated. Rebuild or restart the GeoIntel container to use the promoted model path.",
                "applied": True,
            }
        )

    return 0, payload


def _emit(payload: dict[str, Any], *, as_json: bool) -> None:
    if as_json:
        print(json.dumps(payload, indent=2, sort_keys=True))
        return

    print(f"status: {payload['status']}")
    print(f"message: {payload['message']}")
    if payload.get("selected_host_model_path"):
        print(f"host model: {payload['selected_host_model_path']}")
        print(f"container model: {payload['selected_container_model_path']}")
    if payload.get("env_updates"):
        print("env updates:")
        for key in ENV_KEYS:
            if key in payload["env_updates"]:
                print(f"  {key}={payload['env_updates'][key]}")


def parse_args(argv: list[str]) -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Activate an existing local YOLO model only after a promotion report recommends the exact candidate key."
    )
    parser.add_argument("--promotion-report", required=True, help="Path to detection_model_promotion_report.json.")
    parser.add_argument(
        "--phase4-release-gate-report",
        required=True,
        help=(
            "Path to the governed Phase-4/5 release-gate report. It must bind the exact "
            "candidate key and model SHA-256 and explicitly allow promotion."
        ),
    )
    parser.add_argument("--candidate-key", required=True, help="Exact promoted candidate key from the report.")
    parser.add_argument(
        "--models-dir",
        default=os.environ.get("GEOINTEL_MODELS_PATH", "models"),
        help="Host directory containing mounted local model files.",
    )
    parser.add_argument("--container-model-dir", default="/app/models", help="Container path where --models-dir is mounted.")
    parser.add_argument("--env-file", default=".env", help="Environment file to update when --apply is supplied.")
    parser.add_argument("--apply", action="store_true", help="Write YOLO env updates after all promotion gates pass.")
    parser.add_argument("--json", action="store_true", help="Emit machine-readable JSON.")
    return parser.parse_args(argv)


def main(argv: list[str] | None = None) -> int:
    args = parse_args(argv or sys.argv[1:])
    exit_code, payload = activate(args)
    _emit(payload, as_json=args.json)
    return exit_code


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