#!/usr/bin/env python3
"""Sanitize and validate local GCP IAP operator result artifacts."""

from __future__ import annotations

import argparse
import hashlib
import json
from pathlib import Path
from typing import Any

ALLOWED_REMOTE_FIELDS = {
    "operation",
    "request_id",
    "target_db",
    "status",
    "instance",
    "expected_internal_ip",
    "active_state",
    "sub_state",
    "nrestarts",
    "dispatcher_version",
    "dispatcher_sha256",
    "bundle_commit",
    "bundle_sha256",
    "result_sha256",
    "result_bytes",
    "full_result_export_ready",
    "fixed_baseline_validation",
    "no_message_send",
    "live_service_unchanged",
    "live_profile_unchanged",
    "clone_database_remaining",
    "run_directory_remaining",
    "legacy_directory_remaining",
    "export_file_remaining",
    "leftover_temp_directory_count",
    "rollback_database_datallowconn",
    "rollback_database_connections",
}


def write_sanitized_result(args: argparse.Namespace) -> int:
    raw = args.raw_stdout.read_bytes()
    exit_code = args.exit_code
    remote: dict[str, Any] = {}
    if exit_code == 0:
        try:
            candidate = json.loads(raw.decode("utf-8").splitlines()[-1])
            remote = {key: value for key, value in candidate.items() if key in ALLOWED_REMOTE_FIELDS}
        except (AttributeError, UnicodeDecodeError, json.JSONDecodeError, IndexError):
            exit_code = 90

    payload = {
        "schema": "livingip.gcpIapOperatorSanitizedResult.v1",
        "operation": args.operation,
        "request_id": args.request_id,
        "target_db": args.target_db,
        "status": "pass" if exit_code == 0 and remote.get("status") == "pass" else "fail",
        "ssh_exit_code": exit_code,
        "remote_result": remote,
        "remote_stdout_sha256": hashlib.sha256(raw).hexdigest(),
        "remote_stdout_bytes": len(raw),
        "raw_stdout_retained": False,
        "raw_stderr_retained": False,
        "credential_values_logged": False,
        "ssh_debug_enabled": False,
    }
    args.result_path.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n", encoding="utf-8")
    return 0


def validate_downloaded_replay(args: argparse.Namespace) -> int:
    result = json.loads(args.result_path.read_text(encoding="utf-8"))
    full_bytes = args.full_result.read_bytes()
    full = json.loads(full_bytes)
    remote = result.get("remote_result") or {}
    checks = full.get("checks") or {}
    required = {
        "result_hash": hashlib.sha256(full_bytes).hexdigest() == remote.get("result_sha256"),
        "result_size": len(full_bytes) == remote.get("result_bytes"),
        "status": full.get("status") == "pass",
        "target": full.get("target_database") == args.target_db,
        "strict_six_of_six": (full.get("score") or {}).get("passes") == 6
        and (full.get("score") or {}).get("expected_prompt_count") == 6,
        "all_checks": bool(checks) and all(checks.values()),
    }
    if args.target_db == args.expected_cory_replay_db:
        expected_counts = {
            "claims": 1837,
            "sources": 4145,
            "claim_edges": 4916,
            "claim_evidence": 4670,
            "kb_proposals": 26,
        }
        before_counts = (full.get("canonical_status_before") or {}).get("high_signal_rows") or {}
        after_counts = (full.get("canonical_status_after") or {}).get("high_signal_rows") or {}
        required.update(
            {
                "fixed_fingerprint_before": (full.get("database_fingerprint_before") or {}).get("sha256")
                == args.expected_cory_replay_fingerprint,
                "fixed_fingerprint_after": (full.get("database_fingerprint_after") or {}).get("sha256")
                == args.expected_cory_replay_fingerprint,
                "fixed_counts_before": all(before_counts.get(key) == value for key, value in expected_counts.items()),
                "fixed_counts_after": all(after_counts.get(key) == value for key, value in expected_counts.items()),
            }
        )
    failed = sorted(name for name, passed in required.items() if not passed)
    if failed:
        raise SystemExit("downloaded replay receipt failed: " + ", ".join(failed))
    result["full_result_exported_and_validated"] = True
    result["full_result_retained"] = False
    result["full_result_sha256"] = hashlib.sha256(full_bytes).hexdigest()
    result["full_result_bytes"] = len(full_bytes)
    result["full_result_fixed_checks"] = required
    args.result_path.write_text(json.dumps(result, indent=2, sort_keys=True) + "\n", encoding="utf-8")
    args.full_result.unlink()
    return 0


def mark_export_failure(args: argparse.Namespace) -> int:
    payload = json.loads(args.result_path.read_text(encoding="utf-8"))
    payload["status"] = "fail"
    payload["full_result_exported_and_validated"] = False
    payload["full_result_retained"] = False
    payload["full_result_export_failure"] = "receipt_download_validation_or_remote_export_cleanup_failed"
    payload["ssh_exit_code"] = args.exit_code
    payload["remote_full_result_export_removed"] = args.cleanup_exit_code == 0
    args.result_path.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n", encoding="utf-8")
    return 0


def mark_export_removed(args: argparse.Namespace) -> int:
    payload = json.loads(args.result_path.read_text(encoding="utf-8"))
    payload["remote_full_result_export_removed"] = True
    args.result_path.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n", encoding="utf-8")
    return 0


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    subparsers = parser.add_subparsers(dest="command", required=True)

    sanitize = subparsers.add_parser("sanitize")
    sanitize.add_argument("--result-path", type=Path, required=True)
    sanitize.add_argument("--raw-stdout", type=Path, required=True)
    sanitize.add_argument("--operation", required=True)
    sanitize.add_argument("--request-id", required=True)
    sanitize.add_argument("--target-db", required=True)
    sanitize.add_argument("--exit-code", type=int, required=True)
    sanitize.set_defaults(handler=write_sanitized_result)

    replay = subparsers.add_parser("validate-replay")
    replay.add_argument("--result-path", type=Path, required=True)
    replay.add_argument("--full-result", type=Path, required=True)
    replay.add_argument("--target-db", required=True)
    replay.add_argument("--expected-cory-replay-db", required=True)
    replay.add_argument("--expected-cory-replay-fingerprint", required=True)
    replay.set_defaults(handler=validate_downloaded_replay)

    export_failure = subparsers.add_parser("mark-export-failure")
    export_failure.add_argument("--result-path", type=Path, required=True)
    export_failure.add_argument("--exit-code", type=int, required=True)
    export_failure.add_argument("--cleanup-exit-code", type=int, required=True)
    export_failure.set_defaults(handler=mark_export_failure)

    export_removed = subparsers.add_parser("mark-export-removed")
    export_removed.add_argument("--result-path", type=Path, required=True)
    export_removed.set_defaults(handler=mark_export_removed)
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    return args.handler(args)


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