#!/usr/bin/env python3
"""Generate immutable P4 split manifests and fail closed on leakage."""

from __future__ import annotations

import argparse
import hashlib
import json
import math
import re
from collections import Counter, defaultdict
from datetime import date
from pathlib import Path
from typing import Any, Iterable

from pyproj import CRS

SCHEMA_VERSION = 1
GENERATOR_VERSION = "1.1.0"
TRAIN_SPLITS = {"train"}
SELECTION_SPLITS = {"val", "calibration"}
DEVELOPMENT_SPLITS = TRAIN_SPLITS | SELECTION_SPLITS
PROTECTED_SPLITS = {"test", "background-test", "challenge"}
ALL_SPLITS = DEVELOPMENT_SPLITS | PROTECTED_SPLITS
NORMATIVE_SPLITS = {"train", "val", "calibration", "test", "background-test"}
SHA256_PATTERN = re.compile(r"^[0-9a-f]{64}$")
FINGERPRINT_PATTERN = re.compile(r"^[0-9a-f]{16}$")


class LeakageError(ValueError):
    """Raised when split isolation is not demonstrably safe."""


def canonical_bytes(value: Any) -> bytes:
    return json.dumps(
        value, ensure_ascii=False, sort_keys=True, separators=(",", ":")
    ).encode("utf-8")


def canonical_hash(value: Any) -> str:
    return hashlib.sha256(canonical_bytes(value)).hexdigest()


def write_json(path: Path, payload: Any) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    content = json.dumps(payload, ensure_ascii=False, indent=2, sort_keys=True) + "\n"
    temporary = path.with_name(f".{path.name}.tmp")
    temporary.write_text(content, encoding="utf-8")
    temporary.replace(path)


def validate_metric_crs(value: Any) -> str:
    try:
        crs = CRS.from_user_input(value)
    except Exception as exc:  # noqa: BLE001 - pyproj exposes multiple parser errors
        raise LeakageError(f"Invalid CRS: {value!r}") from exc
    axes = crs.axis_info
    if crs.is_geographic or not axes:
        raise LeakageError(f"CRS must be projected in metres: {value!r}")
    if any(
        abs(float(axis.unit_conversion_factor or 0.0) - 1.0) > 1e-12
        for axis in axes[:2]
    ):
        raise LeakageError(f"CRS axes must use metres: {value!r}")
    return crs.to_string()


def require_hex(
    sample_id: str, field: str, value: Any, pattern: re.Pattern[str]
) -> str:
    normalized = str(value or "").lower()
    if not pattern.fullmatch(normalized):
        raise LeakageError(f"{sample_id}: {field} has an invalid fingerprint")
    return normalized


def fingerprint_distance(left: str, right: str) -> int:
    return (int(left, 16) ^ int(right, 16)).bit_count()


def bbox_distance(left: list[float], right: list[float]) -> float:
    dx = max(left[0] - right[2], right[0] - left[2], 0.0)
    dy = max(left[1] - right[3], right[1] - left[3], 0.0)
    return math.hypot(dx, dy)


def normalized_sample(sample: dict[str, Any]) -> dict[str, Any]:
    required = {
        "sample_id",
        "task",
        "split",
        "group_id",
        "source_family",
        "temporal_family",
        "object_ids",
        "bbox",
        "raw_image_sha256",
        "processed_image_sha256",
        "label_sha256",
        "perceptual_image_hash",
        "label_geometry_hash",
        "label_geometry_fingerprint",
        "native_feature_ids",
        "parent_raster_id",
        "acquisition_id",
        "acquisition_date",
    }
    missing = sorted(required - set(sample))
    if missing:
        raise LeakageError(
            f"{sample.get('sample_id', '<unknown>')}: missing fields {missing}"
        )
    split = str(sample["split"])
    if split not in ALL_SPLITS:
        raise LeakageError(f"{sample['sample_id']}: unsupported split {split!r}")
    bbox = sample["bbox"]
    if not isinstance(bbox, list) or len(bbox) != 4:
        raise LeakageError(f"{sample['sample_id']}: bbox must contain four values")
    values = [float(value) for value in bbox]
    if (
        not all(math.isfinite(value) for value in values)
        or values[0] >= values[2]
        or values[1] >= values[3]
    ):
        raise LeakageError(f"{sample['sample_id']}: invalid bbox")
    result = dict(sample)
    result["bbox"] = values
    for field in (
        "raw_image_sha256",
        "processed_image_sha256",
        "label_sha256",
        "label_geometry_hash",
    ):
        result[field] = require_hex(
            str(sample["sample_id"]), field, sample[field], SHA256_PATTERN
        )
    for field in ("perceptual_image_hash", "label_geometry_fingerprint"):
        result[field] = require_hex(
            str(sample["sample_id"]), field, sample[field], FINGERPRINT_PATTERN
        )
    for field in ("object_ids", "native_feature_ids"):
        if not isinstance(sample[field], list) or any(
            not str(value) for value in sample[field]
        ):
            raise LeakageError(
                f"{sample['sample_id']}: {field} must be a list of non-empty identities"
            )
        result[field] = sorted({str(value) for value in sample[field]})
    for field in (
        "group_id",
        "source_family",
        "temporal_family",
        "parent_raster_id",
        "acquisition_id",
    ):
        if not str(sample[field]).strip():
            raise LeakageError(f"{sample['sample_id']}: {field} must be non-empty")
        result[field] = str(sample[field])
    try:
        result["acquisition_date"] = date.fromisoformat(
            str(sample["acquisition_date"])
        ).isoformat()
    except ValueError as exc:
        raise LeakageError(
            f"{sample['sample_id']}: acquisition_date must be ISO-8601"
        ) from exc
    record_without_split = {
        key: value for key, value in result.items() if key != "split"
    }
    result["record_sha256"] = canonical_hash(record_without_split)
    return result


def assign_split_roles(source: dict[str, Any]) -> list[dict[str, Any]]:
    """Assign complete leakage groups deterministically, or validate a frozen assignment."""
    raw_samples = [dict(item) for item in source.get("samples", [])]
    mode = str(source.get("assignment_mode") or "preassigned")
    split_presence = [bool(item.get("split")) for item in raw_samples]
    if mode == "preassigned":
        if raw_samples and not all(split_presence):
            raise LeakageError("Preassigned mode requires a split on every sample")
        return raw_samples
    if mode != "deterministic_grouped":
        raise LeakageError(f"Unsupported assignment_mode: {mode!r}")
    if any(split_presence):
        raise LeakageError(
            "deterministic_grouped mode refuses partially or fully preassigned splits"
        )
    config = source.get("split_assignment") or {}
    roles = list(
        config.get("roles")
        or ["train", "val", "calibration", "test", "background-test"]
    )
    if not roles or len(set(roles)) != len(roles) or set(roles) - ALL_SPLITS:
        raise LeakageError("split_assignment.roles must be unique supported roles")
    weights = {
        role: float((config.get("weights") or {}).get(role, 1.0)) for role in roles
    }
    if any(not math.isfinite(value) or value <= 0 for value in weights.values()):
        raise LeakageError("split_assignment weights must be positive finite values")
    placeholder = [
        normalized_sample({**item, "split": "train"}) for item in raw_samples
    ]
    parent = list(range(len(placeholder)))

    def find(index: int) -> int:
        while parent[index] != index:
            parent[index] = parent[parent[index]]
            index = parent[index]
        return index

    def union(left: int, right: int) -> None:
        left_root = find(left)
        right_root = find(right)
        if left_root != right_root:
            parent[max(left_root, right_root)] = min(left_root, right_root)

    identity_fields = (
        "group_id",
        "source_family",
        "temporal_family",
        "raw_image_sha256",
        "processed_image_sha256",
        "label_sha256",
        "label_geometry_hash",
        "parent_raster_id",
        "acquisition_id",
    )
    seen: dict[tuple[str, str], int] = {}
    for index, sample in enumerate(placeholder):
        identities = [(field, str(sample[field])) for field in identity_fields]
        identities.extend(("object_id", str(value)) for value in sample["object_ids"])
        identities.extend(
            ("native_feature_id", str(value)) for value in sample["native_feature_ids"]
        )
        for identity in identities:
            if identity in seen:
                union(index, seen[identity])
            else:
                seen[identity] = index
    image_threshold = int(source.get("perceptual_hamming_threshold", 4))
    geometry_threshold = int(source.get("label_geometry_hamming_threshold", 2))
    buffer_m = float(source.get("independence_buffer_m") or 0)
    for index, left in enumerate(placeholder):
        for right_index, right in enumerate(placeholder[index + 1 :], start=index + 1):
            if (
                fingerprint_distance(
                    left["perceptual_image_hash"], right["perceptual_image_hash"]
                )
                <= image_threshold
                or fingerprint_distance(
                    left["label_geometry_fingerprint"],
                    right["label_geometry_fingerprint"],
                )
                <= geometry_threshold
                or bbox_distance(left["bbox"], right["bbox"]) < buffer_m
            ):
                union(index, right_index)
    components: dict[int, list[int]] = defaultdict(list)
    for index in range(len(placeholder)):
        components[find(index)].append(index)
    if len(components) < len(roles):
        raise LeakageError(
            f"Not enough independent groups for required roles: {len(components)} < {len(roles)}"
        )
    seed = str(config.get("seed") or "geointel-p4-group-split-v1")
    groups = sorted(
        components.values(),
        key=lambda indices: canonical_hash(
            {
                "seed": seed,
                "sample_ids": sorted(
                    placeholder[index]["sample_id"] for index in indices
                ),
            }
        ),
    )
    stratify_by = tuple(config.get("stratify_by") or ["task"])
    assigned_counts: Counter[str] = Counter()
    stratum_counts: dict[str, Counter[str]] = defaultdict(Counter)
    assignment: dict[int, str] = {}
    for group_index, indices in enumerate(groups):
        strata = {
            f"{field}={placeholder[index].get(field) or (placeholder[index].get('metadata') or {}).get(field)}"
            for index in indices
            for field in stratify_by
        }
        candidates = (
            roles[group_index : group_index + 1] if group_index < len(roles) else roles
        )
        role = min(
            candidates,
            key=lambda candidate: (
                assigned_counts[candidate] / weights[candidate]
                + sum(
                    stratum_counts[stratum][candidate] / weights[candidate]
                    for stratum in strata
                ),
                candidate,
            ),
        )
        for index in indices:
            assignment[index] = role
        assigned_counts[role] += len(indices)
        for stratum in strata:
            stratum_counts[stratum][role] += len(indices)
    return [
        {**sample, "split": assignment[index]}
        for index, sample in enumerate(raw_samples)
    ]


def leakage_findings(
    samples: list[dict[str, Any]],
    independence_buffer_m: float,
    perceptual_hamming_threshold: int,
    label_geometry_hamming_threshold: int,
) -> list[dict[str, Any]]:
    findings: list[dict[str, Any]] = []
    identity_fields = (
        "sample_id",
        "group_id",
        "source_family",
        "temporal_family",
        "raw_image_sha256",
        "processed_image_sha256",
        "label_sha256",
        "label_geometry_hash",
        "parent_raster_id",
        "acquisition_id",
        "record_sha256",
    )
    identities: dict[str, dict[str, set[str]]] = {
        key: defaultdict(set) for key in identity_fields
    }
    object_splits: dict[str, set[str]] = defaultdict(set)
    native_feature_splits: dict[str, set[str]] = defaultdict(set)
    for sample in samples:
        for key, index in identities.items():
            index[str(sample[key])].add(sample["split"])
        for object_id in sample["object_ids"]:
            object_splits[str(object_id)].add(sample["split"])
        for feature_id in sample["native_feature_ids"]:
            native_feature_splits[str(feature_id)].add(sample["split"])
    code_by_key = {
        "sample_id": "S-SAMPLE-IDENTITY",
        "group_id": "S-SPATIAL-GROUP",
        "source_family": "S-SOURCE-FAMILY",
        "temporal_family": "S-TEMPORAL-FAMILY",
        "raw_image_sha256": "S-RAW-IMAGE-DUPLICATE",
        "processed_image_sha256": "S-PROCESSED-IMAGE-DUPLICATE",
        "label_sha256": "S-LABEL-DUPLICATE",
        "label_geometry_hash": "S-LABEL-GEOMETRY-DUPLICATE",
        "parent_raster_id": "S-PARENT-RASTER",
        "acquisition_id": "S-ACQUISITION",
        "record_sha256": "S-EXACT-DUPLICATE",
    }
    for key, index in identities.items():
        for value, splits in sorted(index.items()):
            if len(splits) > 1:
                findings.append(
                    {
                        "code": code_by_key[key],
                        "identity": value,
                        "splits": sorted(splits),
                    }
                )
    for object_id, splits in sorted(object_splits.items()):
        if len(splits) > 1:
            findings.append(
                {
                    "code": "S-OBJECT-INSTANCE",
                    "identity": object_id,
                    "splits": sorted(splits),
                }
            )
    for feature_id, splits in sorted(native_feature_splits.items()):
        if len(splits) > 1:
            findings.append(
                {
                    "code": "S-NATIVE-FEATURE",
                    "identity": feature_id,
                    "splits": sorted(splits),
                }
            )
    for index, left in enumerate(samples):
        for right in samples[index + 1 :]:
            if left["split"] == right["split"]:
                continue
            image_distance = fingerprint_distance(
                left["perceptual_image_hash"], right["perceptual_image_hash"]
            )
            if image_distance <= perceptual_hamming_threshold:
                findings.append(
                    {
                        "code": "S-PERCEPTUAL-IMAGE-NEAR-DUPLICATE",
                        "left": left["sample_id"],
                        "right": right["sample_id"],
                        "distance": image_distance,
                        "maximum_allowed_distance": perceptual_hamming_threshold,
                    }
                )
            geometry_distance = fingerprint_distance(
                left["label_geometry_fingerprint"], right["label_geometry_fingerprint"]
            )
            if geometry_distance <= label_geometry_hamming_threshold:
                findings.append(
                    {
                        "code": "S-LABEL-GEOMETRY-NEAR-DUPLICATE",
                        "left": left["sample_id"],
                        "right": right["sample_id"],
                        "distance": geometry_distance,
                        "maximum_allowed_distance": label_geometry_hamming_threshold,
                    }
                )
            distance = bbox_distance(left["bbox"], right["bbox"])
            if distance < independence_buffer_m:
                findings.append(
                    {
                        "code": "S-SPATIAL-OVERLAP",
                        "left": left["sample_id"],
                        "right": right["sample_id"],
                        "left_split": left["split"],
                        "right_split": right["split"],
                        "distance_m": distance,
                        "required_distance_m": independence_buffer_m,
                    }
                )
    return findings


def build_manifests(
    source: dict[str, Any],
) -> tuple[dict[str, Any], dict[str, Any], dict[str, Any]]:
    if source.get("schema_version") != 1:
        raise LeakageError("Unsupported source manifest schema_version")
    metric_crs = validate_metric_crs(source.get("crs"))
    assigned_source_samples = assign_split_roles(source)
    samples = [normalized_sample(item) for item in assigned_source_samples]
    if not samples:
        raise LeakageError("Source manifest has no samples")
    if len({item["sample_id"] for item in samples}) != len(samples):
        raise LeakageError("Duplicate sample_id in source manifest")
    buffer_m = float(source.get("independence_buffer_m") or 0)
    if not math.isfinite(buffer_m) or buffer_m <= 0:
        raise LeakageError("independence_buffer_m must be a positive finite number")
    split_counts = Counter(item["split"] for item in samples)
    required_splits = set(source.get("required_splits") or NORMATIVE_SPLITS)
    unsupported_required = sorted(required_splits - ALL_SPLITS)
    if unsupported_required:
        raise LeakageError(f"Unsupported required splits: {unsupported_required}")
    missing_splits = sorted(required_splits - set(split_counts))
    if missing_splits:
        raise LeakageError(f"Required splits are absent: {missing_splits}")
    perceptual_threshold = int(source.get("perceptual_hamming_threshold", 4))
    geometry_threshold = int(source.get("label_geometry_hamming_threshold", 2))
    if not 0 <= perceptual_threshold < 64 or not 0 <= geometry_threshold < 64:
        raise LeakageError("Near-duplicate Hamming thresholds must be between 0 and 63")
    findings = leakage_findings(
        samples, buffer_m, perceptual_threshold, geometry_threshold
    )
    canonical_source = dict(source)
    canonical_source["samples"] = sorted(
        source.get("samples", []), key=lambda item: str(item.get("sample_id", ""))
    )
    source_hash = canonical_hash(canonical_source)

    development_samples = sorted(
        (item for item in samples if item["split"] in DEVELOPMENT_SPLITS),
        key=lambda item: item["sample_id"],
    )
    protected_samples = sorted(
        (item for item in samples if item["split"] in PROTECTED_SPLITS),
        key=lambda item: item["sample_id"],
    )
    common = {
        "schema_version": SCHEMA_VERSION,
        "generator_version": GENERATOR_VERSION,
        "dataset_version": source["dataset_version"],
        "source_manifest_sha256": source_hash,
        "crs": metric_crs,
        "independence_buffer_m": buffer_m,
        "perceptual_hamming_threshold": perceptual_threshold,
        "label_geometry_hamming_threshold": geometry_threshold,
        "assignment_mode": str(source.get("assignment_mode") or "preassigned"),
        "assignment_algorithm": "group-before-split-deficit-balancer-v1",
        "assignment_seed": str(
            (source.get("split_assignment") or {}).get("seed")
            or "geointel-p4-group-split-v1"
        ),
        "claim_boundary": source.get("claim_boundary"),
    }
    development = {
        **common,
        "manifest_role": "development_and_calibration",
        "allowed_splits": sorted(DEVELOPMENT_SPLITS),
        "training_access_allowed_by_split": {
            "train": True,
            "val": False,
            "calibration": False,
        },
        "selection_access_allowed_by_split": {
            "train": False,
            "val": True,
            "calibration": True,
        },
        "samples": development_samples,
    }
    development["manifest_sha256"] = canonical_hash(development)
    protected = {
        **common,
        "manifest_role": "protected_release_only",
        "allowed_splits": sorted(PROTECTED_SPLITS),
        "training_access_allowed": False,
        "selection_use_allowed": False,
        "labels_available_by_split": {
            "test": "frozen_evaluator_only",
            "background-test": "frozen_evaluator_only",
            "challenge": "sealed_external",
        },
        "access_policy": (
            "Only the frozen release evaluator may consume test/background-test labels; "
            "challenge labels remain external/sealed."
        ),
        "samples": protected_samples,
    }
    protected["manifest_sha256"] = canonical_hash(protected)
    leakage = {
        "schema_version": SCHEMA_VERSION,
        "generator_version": GENERATOR_VERSION,
        "status": "pass" if not findings else "fail",
        "source_manifest_sha256": source_hash,
        "development_manifest_sha256": development["manifest_sha256"],
        "protected_manifest_sha256": protected["manifest_sha256"],
        "inventory_total": len(samples),
        "split_counts": dict(sorted(split_counts.items())),
        "checked_identity_fields": [
            "sample_id",
            "group_id",
            "source_family",
            "temporal_family",
            "object_ids",
            "native_feature_ids",
            "raw_image_sha256",
            "processed_image_sha256",
            "label_sha256",
            "record_sha256",
            "perceptual_image_hash",
            "label_geometry_hash",
            "label_geometry_fingerprint",
            "parent_raster_id",
            "acquisition_id",
            "bbox_distance",
        ],
        "crs_validation": {"status": "pass", "crs": metric_crs, "distance_units": "m"},
        "independence_buffer_m": buffer_m,
        "perceptual_hamming_threshold": perceptual_threshold,
        "label_geometry_hamming_threshold": geometry_threshold,
        "finding_count": len(findings),
        "findings": findings,
    }
    return development, protected, leakage


def protected_identities(protected_manifest: dict[str, Any]) -> set[str]:
    identities: set[str] = set()
    scalar_fields = (
        "sample_id",
        "group_id",
        "record_sha256",
        "source_family",
        "temporal_family",
        "raw_image_sha256",
        "processed_image_sha256",
        "label_sha256",
        "label_geometry_hash",
        "perceptual_image_hash",
        "label_geometry_fingerprint",
        "parent_raster_id",
        "acquisition_id",
    )
    for item in protected_manifest.get("samples", []):
        identities.update(str(item[field]) for field in scalar_fields)
        identities.update(str(value) for value in item.get("object_ids", []))
        identities.update(str(value) for value in item.get("native_feature_ids", []))
    return identities


def assert_training_inputs_safe(
    input_paths: Iterable[Path],
    input_records: Iterable[dict[str, Any]],
    protected_manifest: dict[str, Any],
) -> None:
    """Refuse non-train roles and protected identities at a fitting boundary."""
    forbidden = protected_identities(protected_manifest)
    violations: list[str] = []
    for path in input_paths:
        lowered = path.as_posix().lower()
        if any(
            token in lowered
            for token in ("protected", "holdout", "challenge", "background-test")
        ):
            violations.append(f"protected_path:{path}")
    scalar_fields = (
        "sample_id",
        "group_id",
        "record_sha256",
        "source_family",
        "temporal_family",
        "raw_image_sha256",
        "processed_image_sha256",
        "label_sha256",
        "label_geometry_hash",
        "perceptual_image_hash",
        "label_geometry_fingerprint",
        "parent_raster_id",
        "acquisition_id",
    )
    for record in input_records:
        if record.get("split") not in TRAIN_SPLITS:
            violations.append(
                f"non_train_role:{record.get('sample_id')}:{record.get('split')}"
            )
        values = {str(record.get(field, "")) for field in scalar_fields}
        values.update(str(value) for value in record.get("object_ids", []))
        values.update(str(value) for value in record.get("native_feature_ids", []))
        overlap = sorted((values - {""}) & forbidden)
        if overlap:
            violations.append(f"protected_identity:{','.join(overlap)}")
    if violations:
        raise LeakageError(
            "Training input firewall blocked: " + "; ".join(sorted(violations))
        )


def generate(source_path: Path, output_dir: Path) -> dict[str, Any]:
    source = json.loads(source_path.read_text(encoding="utf-8"))
    development, protected, leakage = build_manifests(source)
    output_dir.mkdir(parents=True, exist_ok=True)
    generation_status = {
        "schema_version": SCHEMA_VERSION,
        "generator_version": GENERATOR_VERSION,
        "status": leakage["status"],
        "source_manifest_sha256": leakage["source_manifest_sha256"],
        "development_manifest_sha256": (
            development["manifest_sha256"] if leakage["status"] == "pass" else None
        ),
        "protected_manifest_sha256": (
            protected["manifest_sha256"] if leakage["status"] == "pass" else None
        ),
    }
    write_json(output_dir / "generation-status.json", generation_status)
    write_json(output_dir / "leakage-gate-report.json", leakage)
    if leakage["status"] != "pass":
        raise LeakageError(
            f"Leakage gate failed with {leakage['finding_count']} findings"
        )
    train_samples = [
        item for item in development["samples"] if item["split"] == "train"
    ]
    assert_training_inputs_safe([], train_samples, protected)
    write_json(output_dir / "development-split-manifest.json", development)
    write_json(output_dir / "protected-split-manifest.json", protected)
    return {
        "development": development,
        "protected": protected,
        "leakage": leakage,
        "generation_status": generation_status,
    }


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--source", type=Path, required=True)
    parser.add_argument("--output-dir", type=Path, required=True)
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    try:
        result = generate(args.source, args.output_dir)
    except (OSError, json.JSONDecodeError, LeakageError) as exc:
        print(json.dumps({"status": "fail", "error": str(exc)}, indent=2))
        return 2
    print(
        json.dumps(
            {
                "status": "pass",
                "development_manifest_sha256": result["development"]["manifest_sha256"],
                "protected_manifest_sha256": result["protected"]["manifest_sha256"],
                "split_counts": result["leakage"]["split_counts"],
            },
            indent=2,
        )
    )
    return 0


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