#!/usr/bin/env python3
"""Compare canonical Postgres source and restored-target JSONL manifests."""

from __future__ import annotations

import argparse
import hashlib
import ipaddress
import json
from datetime import UTC, datetime
from pathlib import Path
from typing import Any

SINGLETON_KINDS = {
    "identity",
    "schemas",
    "extensions",
    "application_roles",
    "columns",
    "constraints",
    "indexes",
    "sequences",
    "views",
    "functions",
    "triggers",
    "types",
    "policies",
}


def canonical_hash(value: Any) -> str:
    payload = json.dumps(value, sort_keys=True, separators=(",", ":"), ensure_ascii=True)
    return hashlib.sha256(payload.encode()).hexdigest()


def load_manifest(path: Path) -> dict[str, Any]:
    singleton: dict[str, dict[str, Any]] = {}
    tables: dict[str, dict[str, Any]] = {}
    performance: dict[str, dict[str, Any]] = {}
    with path.open(encoding="utf-8") as handle:
        for line_number, raw_line in enumerate(handle, 1):
            line = raw_line.strip()
            if not line or line in {"BEGIN", "SET", "ROLLBACK"}:
                continue
            try:
                row = json.loads(line)
            except json.JSONDecodeError as exc:
                raise ValueError(f"{path}:{line_number}: invalid JSONL manifest row") from exc
            kind = row.get("kind")
            if kind in SINGLETON_KINDS:
                if kind in singleton:
                    raise ValueError(f"{path}: duplicate {kind} row")
                singleton[str(kind)] = row
            elif kind == "table":
                key = f"{row.get('schema')}.{row.get('table')}"
                if key in tables:
                    raise ValueError(f"{path}: duplicate table row {key}")
                tables[key] = row
            elif kind == "performance":
                key = str(row.get("query"))
                if key in performance:
                    raise ValueError(f"{path}: duplicate performance row {key}")
                performance[key] = row
            else:
                raise ValueError(f"{path}:{line_number}: unknown manifest kind {kind!r}")
    missing = SINGLETON_KINDS - set(singleton)
    if missing:
        raise ValueError(f"{path}: missing singleton rows {sorted(missing)}")
    if not tables:
        raise ValueError(f"{path}: no table rows")
    return {"singleton": singleton, "tables": tables, "performance": performance}


def item_map(row: dict[str, Any], key: str) -> dict[str, dict[str, Any]]:
    return {str(item[key]): item for item in row.get("items", [])}


def compare_manifests(
    source: dict[str, Any],
    target: dict[str, Any],
    *,
    max_target_query_ms: float,
    max_target_source_ratio: float,
) -> tuple[list[str], dict[str, Any]]:
    problems: list[str] = []
    source_singleton = source["singleton"]
    target_singleton = target["singleton"]

    source_identity = source_singleton["identity"]
    target_identity = target_singleton["identity"]
    source_major = int(source_identity["server_version_num"]) // 10000
    target_major = int(target_identity["server_version_num"]) // 10000
    if source_major != target_major:
        problems.append(f"postgres_major source={source_major} target={target_major}")
    if source_identity.get("transaction_read_only") != "on":
        problems.append("source_manifest_not_captured_read_only")
    if target_identity.get("transaction_read_only") != "on":
        problems.append("target_manifest_not_captured_read_only")

    source_tables = source["tables"]
    target_tables = target["tables"]
    if set(source_tables) != set(target_tables):
        problems.append(
            "table_set mismatch source_only="
            + json.dumps(sorted(set(source_tables) - set(target_tables)))
            + " target_only="
            + json.dumps(sorted(set(target_tables) - set(source_tables)))
        )
    table_mismatches = []
    for table in sorted(set(source_tables) & set(target_tables)):
        source_row = source_tables[table]
        target_row = target_tables[table]
        mismatch = {
            field: {"source": source_row.get(field), "target": target_row.get(field)}
            for field in ("row_count", "rowset_md5")
            if source_row.get(field) != target_row.get(field)
        }
        if mismatch:
            table_mismatches.append({"table": table, "mismatch": mismatch})
    if table_mismatches:
        problems.append(f"table_content_mismatches={len(table_mismatches)}")

    structural_hashes: dict[str, dict[str, str]] = {}
    for kind in (
        "schemas",
        "columns",
        "constraints",
        "indexes",
        "sequences",
        "views",
        "functions",
        "triggers",
        "types",
        "policies",
    ):
        source_hash = canonical_hash(source_singleton[kind].get("items", []))
        target_hash = canonical_hash(target_singleton[kind].get("items", []))
        structural_hashes[kind] = {"source": source_hash, "target": target_hash}
        if source_hash != target_hash:
            problems.append(f"{kind}_hash_mismatch")

    source_extensions = item_map(source_singleton["extensions"], "name")
    target_extensions = item_map(target_singleton["extensions"], "name")
    extension_mismatches = {
        name: {"source": item, "target": target_extensions.get(name)}
        for name, item in source_extensions.items()
        if target_extensions.get(name) != item
    }
    if extension_mismatches:
        problems.append(f"required_extension_mismatches={len(extension_mismatches)}")

    source_roles = item_map(source_singleton["application_roles"], "name")
    target_roles = item_map(target_singleton["application_roles"], "name")
    role_mismatches = {
        name: {"source": item, "target": target_roles.get(name)}
        for name, item in source_roles.items()
        if target_roles.get(name) != item
    }
    if role_mismatches:
        problems.append(f"application_role_mismatches={len(role_mismatches)}")

    performance_problems = []
    performance_rows = []
    for query, target_row in sorted(target["performance"].items()):
        target_ms = float(target_row.get("elapsed_ms") or 0)
        source_row = source["performance"].get(query)
        source_ms = float((source_row or {}).get("elapsed_ms") or 0)
        ratio = target_ms / max(source_ms, 5.0)
        row = {"query": query, "source_ms": source_ms, "target_ms": target_ms, "source_ratio": ratio}
        performance_rows.append(row)
        if source_row is None:
            performance_problems.append(f"{query}:missing_source_benchmark")
        if target_ms > max_target_query_ms:
            performance_problems.append(f"{query}:target_ms={target_ms:.3f}>{max_target_query_ms:.3f}")
        if ratio > max_target_source_ratio:
            performance_problems.append(f"{query}:source_ratio={ratio:.3f}>{max_target_source_ratio:.3f}")
    missing_target_benchmarks = sorted(set(source["performance"]) - set(target["performance"]))
    if missing_target_benchmarks:
        performance_problems.append("missing_target_benchmarks=" + ",".join(missing_target_benchmarks))
    if performance_problems:
        problems.append(f"performance_problems={len(performance_problems)}")

    details = {
        "source_database": source_identity.get("database"),
        "target_database": target_identity.get("database"),
        "source_connection": {
            "server_address": source_identity.get("server_address"),
            "server_port": source_identity.get("server_port"),
            "ssl": source_identity.get("ssl"),
            "ssl_version": source_identity.get("ssl_version"),
            "database_collation": source_identity.get("database_collation"),
            "database_ctype": source_identity.get("database_ctype"),
        },
        "target_connection": {
            "server_address": target_identity.get("server_address"),
            "server_port": target_identity.get("server_port"),
            "ssl": target_identity.get("ssl"),
            "ssl_version": target_identity.get("ssl_version"),
            "database_collation": target_identity.get("database_collation"),
            "database_ctype": target_identity.get("database_ctype"),
        },
        "source_table_count": len(source_tables),
        "target_table_count": len(target_tables),
        "source_total_rows": sum(int(row["row_count"]) for row in source_tables.values()),
        "target_total_rows": sum(int(row["row_count"]) for row in target_tables.values()),
        "table_mismatches": table_mismatches,
        "structural_hashes": structural_hashes,
        "required_extension_mismatches": extension_mismatches,
        "application_role_mismatches": role_mismatches,
        "target_extra_application_roles": sorted(set(target_roles) - set(source_roles)),
        "performance": performance_rows,
        "performance_problems": performance_problems,
    }
    return problems, details


def canonical_ip(address: str) -> str:
    return str(ipaddress.ip_interface(address).ip)


def is_rfc1918(address: str) -> bool:
    parsed = ipaddress.ip_address(canonical_ip(address))
    return any(
        parsed in network
        for network in (
            ipaddress.ip_network("10.0.0.0/8"),
            ipaddress.ip_network("172.16.0.0/12"),
            ipaddress.ip_network("192.168.0.0/16"),
        )
    )


def load_connectivity(
    path: Path | None,
    scope: str,
    target_identity: dict[str, Any],
) -> tuple[list[str], dict[str, Any]]:
    if scope == "local_restore":
        return [], {"required": False, "status": "not_applicable_local_restore"}
    if path is None:
        return ["gcp_private_connectivity_proof_missing"], {"required": True, "status": "missing"}
    proof = json.loads(path.read_text())
    problems = []
    if proof.get("status") != "pass":
        problems.append("gcp_private_connectivity_status_not_pass")
    if proof.get("private_connectivity") is not True:
        problems.append("gcp_private_connectivity_not_proven")
    if proof.get("public_ip_disabled") is not True:
        problems.append("gcp_public_ip_disabled_not_proven")
    if not proof.get("source_compute"):
        problems.append("gcp_source_compute_missing")
    target_database = target_identity.get("database")
    target_address = target_identity.get("server_address")
    target_ssl = target_identity.get("ssl")
    canonical_target_address = None
    canonical_proof_address = None
    if not target_database:
        problems.append("gcp_target_manifest_database_missing")
    if proof.get("target_database") != target_database:
        problems.append("gcp_connectivity_target_database_mismatch")
    if not target_address:
        problems.append("gcp_target_manifest_server_address_missing")
    else:
        try:
            canonical_target_address = canonical_ip(str(target_address))
            if not is_rfc1918(str(target_address)):
                problems.append("gcp_target_manifest_server_address_not_rfc1918")
        except ValueError:
            problems.append("gcp_target_manifest_server_address_invalid")
    if proof.get("server_address"):
        try:
            canonical_proof_address = canonical_ip(str(proof["server_address"]))
        except ValueError:
            problems.append("gcp_connectivity_server_address_invalid")
    else:
        problems.append("gcp_connectivity_server_address_missing")
    if canonical_proof_address != canonical_target_address:
        problems.append("gcp_connectivity_server_address_mismatch")
    if target_ssl is not True:
        problems.append("gcp_target_manifest_tls_not_proven")
    if proof.get("ssl") is not True:
        problems.append("gcp_connectivity_tls_not_proven")
    return problems, {"required": True, "path": str(path), "proof": proof}


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--source", required=True, type=Path)
    parser.add_argument("--target", required=True, type=Path)
    parser.add_argument("--scope", choices=("local_restore", "gcp_staging"), default="local_restore")
    parser.add_argument("--connectivity-proof", type=Path)
    parser.add_argument("--max-target-query-ms", type=float, default=2000.0)
    parser.add_argument("--max-target-source-ratio", type=float, default=20.0)
    parser.add_argument("--output", required=True, type=Path)
    args = parser.parse_args()

    try:
        source = load_manifest(args.source)
        target = load_manifest(args.target)
        problems, details = compare_manifests(
            source,
            target,
            max_target_query_ms=args.max_target_query_ms,
            max_target_source_ratio=args.max_target_source_ratio,
        )
        connectivity_problems, connectivity = load_connectivity(
            args.connectivity_proof,
            args.scope,
            target["singleton"]["identity"],
        )
        problems.extend(connectivity_problems)
        status = "pass" if not problems else "fail"
        error = None
    except (OSError, ValueError, KeyError, TypeError, json.JSONDecodeError) as exc:
        status = "fail"
        problems = [f"manifest_error={exc}"]
        details = {}
        connectivity = {"required": args.scope == "gcp_staging", "status": "not_checked"}
        error = str(exc)

    payload = {
        "artifact": "canonical_postgres_parity_verification",
        "generated_at_utc": datetime.now(UTC).isoformat(),
        "scope": args.scope,
        "source_manifest": str(args.source),
        "target_manifest": str(args.target),
        "status": status,
        "problems": problems,
        "details": details,
        "private_connectivity": connectivity,
        "error": error,
        "not_proven_by_this_artifact": [
            "GCP staging Leo composition replay",
            "ongoing replication after the manifest timestamps",
            "production cutover",
        ],
    }
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n")
    print(json.dumps(payload, indent=2, sort_keys=True))
    return 0 if status == "pass" else 1


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