from __future__ import annotations

import hashlib
import json
import os
import subprocess
import sys
import zipfile
from pathlib import Path

os.environ.setdefault("KERAS_BACKEND", "tensorflow")
os.environ.setdefault("TF_CPP_MIN_LOG_LEVEL", "2")

import keras
import numpy as np


ROOT = Path(__file__).resolve().parent

VARIANTS = [
    {
        "name": "activation",
        "artifact": ROOT / "variants" / "activation" / "activation_lambda_manual_config.keras",
        "marker": ROOT / "variants" / "activation" / "activation_lambda_marker.txt",
        "expected_marker": "KERAS_NATIVE_ACTIVATION_LAMBDA_MARKER",
        "inference_input": np.array([[2.0]], dtype="float32"),
    },
    {
        "name": "ops_map",
        "artifact": ROOT / "variants" / "ops_map" / "ops_map_inbound_lambda.keras",
        "marker": ROOT / "variants" / "ops_map" / "ops_map_marker.txt",
        "expected_marker": "OPS_MAP_MARKER",
        "inference_input": np.ones((1, 2, 1), dtype="float32"),
    },
    {
        "name": "vectorized_map",
        "artifact": ROOT / "variants" / "vectorized_map" / "vectorized_map_lambda.keras",
        "marker": ROOT / "variants" / "vectorized_map" / "vectorized_map_marker.txt",
        "expected_marker": "VECTORIZED_MAP_MARKER",
        "inference_input": np.ones((1, 2, 1), dtype="float32"),
    },
]


def sha256(path: Path) -> str:
    return hashlib.sha256(path.read_bytes()).hexdigest()


def lambda_paths(path: Path) -> list[str]:
    with zipfile.ZipFile(path) as zf:
        config = json.loads(zf.read("config.json"))
    hits: list[str] = []

    def walk(value, trail: str = "$") -> None:
        if isinstance(value, dict):
            if value.get("class_name") == "__lambda__":
                hits.append(trail)
            for key, child in value.items():
                walk(child, f"{trail}.{key}")
        elif isinstance(value, list):
            for index, child in enumerate(value):
                walk(child, f"{trail}[{index}]")

    walk(config)
    return hits


def run_modelscan(artifact: Path) -> dict[str, object]:
    output_path = artifact.with_name(f"{artifact.stem}.modelscan.json")
    output_path.unlink(missing_ok=True)
    proc = subprocess.run(
        [
            sys.executable,
            "-m",
            "modelscan.cli",
            "-p",
            str(artifact),
            "-r",
            "json",
            "-o",
            str(output_path),
            "--show-skipped",
        ],
        text=True,
        capture_output=True,
        check=False,
        timeout=60,
    )
    parsed: dict[str, object]
    try:
        parsed = json.loads(output_path.read_text())
    except (json.JSONDecodeError, FileNotFoundError):
        parsed = {"stdout": proc.stdout, "stderr": proc.stderr}
    parsed["_returncode"] = proc.returncode
    parsed["_stderr"] = proc.stderr
    return parsed


def verify_variant(variant: dict[str, object]) -> dict[str, object]:
    artifact = Path(variant["artifact"])
    marker = Path(variant["marker"])
    expected_marker = str(variant["expected_marker"])
    inference_input = variant["inference_input"]

    marker.unlink(missing_ok=True)
    old_cwd = Path.cwd()
    os.chdir(artifact.parent)
    try:
        result: dict[str, object] = {
            "name": variant["name"],
            "artifact": str(artifact.relative_to(ROOT)),
            "sha256": sha256(artifact),
            "lambda_paths": lambda_paths(artifact),
        }

        try:
            keras.saving.load_model(artifact, safe_mode=True)
            result["safe_mode_true"] = "loaded_unexpectedly"
        except Exception as exc:  # noqa: BLE001
            result["safe_mode_true"] = (
                f"blocked:{type(exc).__name__}:{str(exc).splitlines()[0]}"
            )
        result["marker_after_safe_mode_true"] = marker.exists()

        marker.unlink(missing_ok=True)
        model = keras.saving.load_model(artifact, safe_mode=False)
        result["marker_after_safe_mode_false_load"] = marker.exists()
        if not marker.exists():
            model(inference_input)
        result["marker_after_unsafe_inference"] = marker.exists()
        result["marker_content"] = marker.read_text().strip() if marker.exists() else ""
        result["marker_matched"] = result["marker_content"] == expected_marker

        scan = run_modelscan(artifact)
        summary = scan.get("summary", {}) if isinstance(scan, dict) else {}
        result["modelscan_returncode"] = scan.get("_returncode")
        result["modelscan_version"] = (
            summary.get("modelscan_version") if isinstance(summary, dict) else None
        )
        result["modelscan_total_issues"] = (
            summary.get("total_issues") if isinstance(summary, dict) else None
        )
        result["modelscan_issues"] = scan.get("issues", [])
        result["passed"] = (
            str(result["safe_mode_true"]).startswith("blocked:")
            and result["marker_after_safe_mode_true"] is False
            and result["marker_matched"] is True
            and result["modelscan_total_issues"] == 0
        )
        return result
    finally:
        os.chdir(old_cwd)


def main() -> int:
    results = {
        "python": sys.version.split()[0],
        "keras": keras.__version__,
        "variants": [verify_variant(variant) for variant in VARIANTS],
    }
    results["all_passed"] = all(item["passed"] for item in results["variants"])
    print(json.dumps(results, indent=2, sort_keys=True))
    return 0 if results["all_passed"] else 1


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