geointel/scripts/manage_grb_refresh.py

#!/usr/bin/env python3
"""Plan, stage and explicitly apply a governed regional GRB refresh.

Planning is read-only. Staging downloads resumable municipality partitions but
does not persist application data. Applying requires the exact staged plan
SHA-256 and delegates persistence to the existing DatasetService-based regional
operators. Existing snapshots are immutable and remain available.
"""

from __future__ import annotations

import argparse
import hashlib
import json
import os
import subprocess
import sys
from datetime import datetime, timezone
from pathlib import Path
from typing import Any
from urllib.error import HTTPError, URLError
from urllib.parse import urlparse
from urllib.request import Request, urlopen

from geographic_scopes import GEOGRAPHIC_SCOPES


DEFAULT_API_URL = "http://127.0.0.1:8000/api/v1"
DEFAULT_SCOPE = "kempen-transport-region"
DEFAULT_OUTPUT_ROOT = Path("/app/storage/operator-data/regional-themes")
DEFAULT_EVIDENCE_ROOT = Path("/app/storage/operator-evidence/grb-refresh")
THEMES = ("buildings", "roads", "water", "parcels")


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Governed GRB refresh: plan, stage, then checksum-confirmed apply.")
    parser.add_argument("action", choices=("plan", "stage", "apply"))
    parser.add_argument("--project-id", required=True, help="GeoIntel project UUID for the governed scope")
    parser.add_argument("--scope", choices=(DEFAULT_SCOPE,), default=DEFAULT_SCOPE)
    parser.add_argument("--api-url", default=os.environ.get("GEOINTEL_API_URL", DEFAULT_API_URL))
    parser.add_argument("--layers", nargs="+", default=list(THEMES), help="Subset: buildings roads water parcels")
    parser.add_argument("--confirm-edition", help="Exact official ISO edition date required for stage/apply")
    parser.add_argument("--confirm-plan-sha256", help="Exact staged plan hash required for apply")
    parser.add_argument("--plan-path", type=Path, help="Staged plan path; defaults to the governed evidence location")
    parser.add_argument("--output-root", type=Path, default=Path(os.environ.get("GEOINTEL_REGIONAL_THEME_OUTPUT_ROOT", DEFAULT_OUTPUT_ROOT)))
    parser.add_argument("--evidence-root", type=Path, default=Path(os.environ.get("GEOINTEL_OPERATOR_EVIDENCE_ROOT", DEFAULT_EVIDENCE_ROOT)))
    parser.add_argument("--refresh-catalog", action="store_true", help="Bypass the short catalog probe cache")
    parser.add_argument("--request-timeout", type=int, default=180)
    parser.add_argument("--api-timeout", type=int, default=180)
    parser.add_argument("--batch-size", type=int, default=1000)
    parser.add_argument("--page-limit", type=int, default=1000)
    parser.add_argument("--max-features-per-member", type=int, default=100000)
    parser.add_argument("--max-total-features", type=int, default=1500000)
    return parser.parse_args()


def selected_themes(values: str | list[str]) -> list[str]:
    raw = [values] if isinstance(values, str) else values
    requested = {item.strip().lower() for value in raw for item in value.split(",") if item.strip()}
    unknown = requested - set(THEMES)
    if unknown or not requested:
        raise ValueError(f"Unsupported GRB layers: {sorted(unknown)}")
    return [theme for theme in THEMES if theme in requested]


def api_data(api_url: str, path: str, timeout: int) -> dict[str, Any]:
    endpoint = f"{api_url.rstrip('/')}/{path.lstrip('/')}"
    request = Request(endpoint, headers={"Accept": "application/json", "User-Agent": "GeoIntel-GRB-refresh/1.0"})
    try:
        with urlopen(request, timeout=timeout) as response:
            payload = json.load(response)
    except HTTPError as exc:
        body = exc.read().decode("utf-8", errors="replace")
        raise RuntimeError(f"GeoIntel API returned HTTP {exc.code}: {body[-1000:]}") from exc
    except URLError as exc:
        raise RuntimeError(f"GeoIntel API is unreachable: {exc.reason}") from exc
    if not isinstance(payload, dict) or not isinstance(payload.get("data"), dict):
        raise RuntimeError("GeoIntel API response is not a canonical data envelope")
    return payload["data"]


def fetch_refresh_plan(args: argparse.Namespace, *, refresh: bool) -> dict[str, Any]:
    query = "true" if refresh else "false"
    return api_data(
        args.api_url,
        f"projects/{args.project_id}/datasets/grb-refresh-plan?scope={args.scope}&refresh_catalog={query}",
        args.api_timeout,
    )


def validate_project_scope(args: argparse.Namespace) -> None:
    project = api_data(args.api_url, f"projects/{args.project_id}", args.api_timeout)
    expected_name = GEOGRAPHIC_SCOPES[args.scope].project_name
    if project.get("name") != expected_name:
        raise RuntimeError(
            f"Project {args.project_id} is '{project.get('name')}', but scope {args.scope} requires '{expected_name}'"
        )


def require_confirmed_edition(plan: dict[str, Any], confirmation: str | None) -> str:
    edition = str(plan.get("remote_edition_date") or "")
    if not edition:
        raise RuntimeError("The official GRB catalog did not provide a safe ISO edition date")
    if confirmation != edition:
        raise RuntimeError(f"Explicit --confirm-edition {edition} is required")
    return edition


def internal_base_url(api_url: str) -> str:
    value = api_url.rstrip("/")
    if value.endswith("/api/v1"):
        value = value[:-7]
    parsed = urlparse(value)
    if parsed.hostname not in {"127.0.0.1", "localhost", "::1"}:
        raise RuntimeError("Stage/apply must run inside GeoIntel against the local backend API")
    return value


def build_operator_commands(
    args: argparse.Namespace,
    themes: list[str],
    edition: str,
    *,
    fetch_only: bool,
) -> list[tuple[str, list[str]]]:
    scripts_dir = Path(__file__).resolve().parent
    common = [
        "--scope", args.scope,
        "--observed-date", edition,
        "--base-url", internal_base_url(args.api_url),
        "--output-root", str(args.output_root),
        "--request-timeout", str(args.request_timeout),
        "--api-timeout", str(args.api_timeout),
        "--batch-size", str(args.batch_size),
        "--page-limit", str(args.page_limit),
        "--max-features-per-member", str(args.max_features_per_member),
        "--max-total-features", str(args.max_total_features),
    ]
    if fetch_only:
        common.append("--fetch-only")
    commands: list[tuple[str, list[str]]] = []
    if "buildings" in themes:
        commands.append(("buildings", [sys.executable, str(scripts_dir / "provision_regional_grb_buildings.py"), *common]))
    context = [theme for theme in themes if theme != "buildings"]
    if context:
        commands.append(
            (
                "context",
                [
                    sys.executable,
                    str(scripts_dir / "provision_regional_grb_context.py"),
                    "--layers",
                    *context,
                    *common,
                ],
            )
        )
    return commands


def run_operator(label: str, command: list[str]) -> dict[str, Any]:
    print(f"GRB {label}: {'staging' if '--fetch-only' in command else 'applying'}...", file=sys.stderr, flush=True)
    completed = subprocess.run(command, check=False, capture_output=True, text=True, encoding="utf-8")
    if completed.returncode != 0:
        detail = completed.stderr.strip() or completed.stdout.strip() or "operator returned no diagnostics"
        raise RuntimeError(f"GRB {label} failed with exit {completed.returncode}: {detail[-3000:]}")
    try:
        payload = json.loads(completed.stdout)
    except json.JSONDecodeError as exc:
        raise RuntimeError(f"GRB {label} returned invalid JSON: {completed.stdout[-1000:]}") from exc
    if payload.get("status") != "ok":
        raise RuntimeError(f"GRB {label} did not report success")
    return payload


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


def canonical_plan_sha256(payload: dict[str, Any]) -> str:
    content = {key: value for key, value in payload.items() if key != "plan_sha256"}
    encoded = json.dumps(content, ensure_ascii=False, sort_keys=True, separators=(",", ":")).encode("utf-8")
    return hashlib.sha256(encoded).hexdigest()


def manifest_path(output_root: Path, scope: str, theme: str, edition: str) -> Path:
    name = "regional_buildings_manifest.json" if theme == "buildings" else f"regional_{theme}_manifest.json"
    return output_root / scope / theme / edition / name


def validate_manifest(path: Path, *, output_root: Path, scope: str, theme: str, edition: str) -> dict[str, Any]:
    root = output_root.resolve()
    resolved = path.resolve()
    if not resolved.is_relative_to(root):
        raise RuntimeError(f"Manifest is outside the governed output root: {path}")
    if not path.is_file():
        raise RuntimeError(f"Staged manifest is missing: {path}")
    payload = json.loads(path.read_text(encoding="utf-8"))
    if payload.get("status") != "complete" or payload.get("scope") != scope or payload.get("theme") != theme:
        raise RuntimeError(f"Staged manifest identity is invalid: {path}")
    if payload.get("observed_at") != edition or payload.get("reference_truncated") is not False:
        raise RuntimeError(f"Staged manifest edition or completeness is invalid: {path}")
    artifact = path.parent / str(payload.get("artifact_filename") or "")
    if not artifact.is_file() or sha256_file(artifact) != payload.get("artifact_sha256"):
        raise RuntimeError(f"Staged artifact checksum is invalid: {artifact}")
    partitions = payload.get("partitions")
    if not isinstance(partitions, list) or len(partitions) != int(payload.get("member_count") or 0):
        raise RuntimeError(f"Staged partition count is invalid: {path}")
    partition_root = path.parent / "partitions"
    for item in partitions:
        partition = partition_root / str(item.get("filename") or "")
        if not partition.is_file() or sha256_file(partition) != item.get("sha256"):
            raise RuntimeError(f"Staged partition checksum is invalid: {partition}")
    return payload


def default_plan_path(args: argparse.Namespace, edition: str) -> Path:
    return args.evidence_root / args.scope / edition / "staged-plan.json"


def build_staged_plan(
    args: argparse.Namespace,
    remote_plan: dict[str, Any],
    themes: list[str],
    edition: str,
) -> dict[str, Any]:
    local_by_theme = {item["theme"]: item for item in remote_plan.get("layers", [])}
    layers: list[dict[str, Any]] = []
    for theme in themes:
        path = manifest_path(args.output_root, args.scope, theme, edition)
        manifest = validate_manifest(path, output_root=args.output_root, scope=args.scope, theme=theme, edition=edition)
        local = local_by_theme.get(theme, {})
        staged_count = int(manifest["feature_count"])
        local_count = local.get("local_feature_count")
        layers.append(
            {
                "theme": theme,
                "collections": manifest.get("grb_collections") or (["GBG"] if theme == "buildings" else []),
                "manifest_path": str(path),
                "artifact_path": str(path.parent / manifest["artifact_filename"]),
                "artifact_sha256": manifest["artifact_sha256"],
                "artifact_size_bytes": int(manifest["artifact_size_bytes"]),
                "partition_count": len(manifest["partitions"]),
                "staged_feature_count": staged_count,
                "local_dataset_id": local.get("local_dataset_id"),
                "local_source_version": local.get("local_source_version"),
                "local_feature_count": local_count,
                "feature_count_delta": staged_count - int(local_count) if local_count is not None else None,
                "existing_snapshot_retained": True,
            }
        )
    payload: dict[str, Any] = {
        "schema_version": 1,
        "status": "staged",
        "created_at": datetime.now(timezone.utc).isoformat(),
        "project_id": args.project_id,
        "scope": args.scope,
        "official_remote_version": remote_plan.get("remote_version"),
        "edition": edition,
        "layers": layers,
        "total_staged_feature_count": sum(item["staged_feature_count"] for item in layers),
        "total_artifact_size_bytes": sum(item["artifact_size_bytes"] for item in layers),
        "automatic_import": False,
        "destructive_replacement": False,
        "apply_requires_plan_sha256": True,
    }
    payload["plan_sha256"] = canonical_plan_sha256(payload)
    return payload


def write_json(path: Path, payload: dict[str, Any]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    temporary = path.with_suffix(f"{path.suffix}.partial")
    temporary.write_text(json.dumps(payload, ensure_ascii=False, indent=2, sort_keys=True) + "\n", encoding="utf-8")
    temporary.replace(path)


def load_and_validate_staged_plan(args: argparse.Namespace, edition: str) -> tuple[Path, dict[str, Any]]:
    path = args.plan_path or default_plan_path(args, edition)
    if not path.is_file():
        raise RuntimeError(f"Staged refresh plan is missing: {path}")
    payload = json.loads(path.read_text(encoding="utf-8"))
    actual_sha = canonical_plan_sha256(payload)
    if payload.get("plan_sha256") != actual_sha:
        raise RuntimeError("Staged refresh plan checksum is invalid")
    if args.confirm_plan_sha256 != actual_sha:
        raise RuntimeError(f"Explicit --confirm-plan-sha256 {actual_sha} is required")
    if payload.get("status") != "staged" or payload.get("project_id") != args.project_id:
        raise RuntimeError("Staged refresh plan project or status is invalid")
    if payload.get("scope") != args.scope or payload.get("edition") != edition:
        raise RuntimeError("Staged refresh plan scope or edition is invalid")
    for layer in payload.get("layers") or []:
        manifest = validate_manifest(
            Path(layer["manifest_path"]),
            output_root=args.output_root,
            scope=args.scope,
            theme=layer["theme"],
            edition=edition,
        )
        if manifest.get("artifact_sha256") != layer.get("artifact_sha256"):
            raise RuntimeError(f"Staged plan no longer matches {layer['theme']} artifact")
    return path, payload


def main() -> int:
    args = parse_args()
    try:
        themes = selected_themes(args.layers)
        if min(args.request_timeout, args.api_timeout, args.batch_size, args.page_limit, args.max_features_per_member, args.max_total_features) <= 0:
            raise ValueError("All timeout, paging, batching and feature safety limits must be positive")
        validate_project_scope(args)
        remote_plan = fetch_refresh_plan(args, refresh=args.refresh_catalog or args.action == "stage")
        if args.action == "plan":
            print(json.dumps(remote_plan, ensure_ascii=False, indent=2))
            return 0

        edition = require_confirmed_edition(remote_plan, args.confirm_edition)
        if args.action == "stage":
            actionable = {
                item["theme"]
                for item in remote_plan.get("layers", [])
                if item.get("status") in {"update_available", "not_loaded"}
            }
            blocked = set(themes) - actionable
            if blocked:
                raise RuntimeError(f"Layers are not safely stageable according to the refresh plan: {sorted(blocked)}")
            commands = build_operator_commands(args, themes, edition, fetch_only=True)
            operator_results = {label: run_operator(label, command) for label, command in commands}
            staged = build_staged_plan(args, remote_plan, themes, edition)
            staged["operator_results"] = operator_results
            staged["plan_sha256"] = canonical_plan_sha256(staged)
            path = args.plan_path or default_plan_path(args, edition)
            write_json(path, staged)
            print(json.dumps({"status": "staged", "plan_path": str(path), **staged}, ensure_ascii=False, indent=2))
            return 0

        path, staged = load_and_validate_staged_plan(args, edition)
        staged_themes = [item["theme"] for item in staged["layers"]]
        commands = build_operator_commands(args, staged_themes, edition, fetch_only=False)
        operator_results = {label: run_operator(label, command) for label, command in commands}
        final_plan = fetch_refresh_plan(args, refresh=False)
        by_theme = {item["theme"]: item for item in final_plan.get("layers", [])}
        incomplete = [theme for theme in staged_themes if by_theme.get(theme, {}).get("status") != "current"]
        if incomplete:
            raise RuntimeError(f"Applied datasets did not become current: {incomplete}")
        evidence = {
            "schema_version": 1,
            "status": "applied",
            "applied_at": datetime.now(timezone.utc).isoformat(),
            "project_id": args.project_id,
            "scope": args.scope,
            "edition": edition,
            "staged_plan_path": str(path),
            "staged_plan_sha256": staged["plan_sha256"],
            "operator_results": operator_results,
            "resulting_layers": [by_theme[theme] for theme in staged_themes],
            "existing_snapshots_retained": True,
        }
        evidence_path = path.with_name("applied-evidence.json")
        write_json(evidence_path, evidence)
        print(json.dumps({"status": "applied", "evidence_path": str(evidence_path), **evidence}, ensure_ascii=False, indent=2))
        return 0
    except (OSError, RuntimeError, ValueError, KeyError, json.JSONDecodeError) as exc:
        print(json.dumps({"status": "error", "message": str(exc)}, ensure_ascii=False), file=sys.stderr)
        return 1


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