"""Run reproducible validation stages and maintain the upload gate report."""

from __future__ import annotations

import argparse
import json
import os
import subprocess
import sys
import xml.etree.ElementTree as ET
from datetime import datetime, timezone
from pathlib import Path

from .constants import NUMERICALLY_VALIDATED_MAX_TOKENS, SOURCE_REVISION


STAGE_COMMANDS = {
    "client_compatibility": ["tests/test_client_contract.py", "tests/test_export_builder.py"],
    "normal_parity": ["tests/test_parity.py", "-m", "parity and not slow"],
    "slow_65536_parity": ["tests/test_parity.py", "-m", "slow"],
    "structural_262144": ["tests/test_structure.py"],
}


def _run_stage(name: str, artifact: Path) -> dict:
    environment = os.environ.copy()
    environment["OPENTHAI_ONNX_ARTIFACT"] = str(artifact)
    if name == "normal_parity":
        environment["RUN_PARITY"] = "1"
    if name == "slow_65536_parity":
        environment["RUN_PARITY"] = "1"
        environment["RUN_SLOW_65536"] = "1"
    junit_path = artifact / f".{name}.junit.xml"
    command = [
        sys.executable,
        "-m",
        "pytest",
        "-q",
        *STAGE_COMMANDS[name],
        f"--junitxml={junit_path}",
    ]
    completed = subprocess.run(command, env=environment, text=True)
    counts = {"tests": 0, "failures": 0, "errors": 0, "skipped": 0}
    if junit_path.is_file():
        try:
            root = ET.parse(junit_path).getroot()
            suites = [root] if root.tag == "testsuite" else list(root.iter("testsuite"))
            for suite in suites:
                for key in counts:
                    counts[key] += int(suite.attrib.get(key, 0))
        finally:
            junit_path.unlink(missing_ok=True)
    passed = (
        completed.returncode == 0
        and counts["tests"] > 0
        and counts["failures"] == 0
        and counts["errors"] == 0
        and counts["skipped"] == 0
    )
    return {
        "passed": passed,
        "returncode": completed.returncode,
        "command": command,
        "test_counts": counts,
        "completed_at": datetime.now(timezone.utc).isoformat(),
    }


def validate(artifact: Path, stages: list[str]) -> Path:
    report_path = artifact / "validation-report.json"
    if report_path.exists():
        report = json.loads(report_path.read_text(encoding="utf-8"))
    else:
        report = {
            "format_version": 1,
            "source_revision": SOURCE_REVISION,
            "numerical_parity_certified_through_tokens": NUMERICALLY_VALIDATED_MAX_TOKENS,
            "lengths_above_65536": "structural validation only",
            "stages": {},
        }
    for stage in stages:
        print(f"\n=== {stage} ===", flush=True)
        report["stages"][stage] = _run_stage(stage, artifact)
        report["upload_gate"] = {
            "required_stages": list(STAGE_COMMANDS),
            "passed": all(report["stages"].get(item, {}).get("passed", False) for item in STAGE_COMMANDS),
        }
        report_path.write_text(json.dumps(report, indent=2, sort_keys=True) + "\n", encoding="utf-8")
        if not report["stages"][stage]["passed"]:
            raise SystemExit(f"validation stage failed: {stage}")
    return report_path


def build_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--artifact", type=Path, required=True)
    group = parser.add_mutually_exclusive_group(required=True)
    group.add_argument("--client", action="store_true")
    group.add_argument("--normal", action="store_true")
    group.add_argument("--slow", action="store_true")
    group.add_argument("--structural", action="store_true")
    group.add_argument("--all", action="store_true")
    return parser


def main(argv: list[str] | None = None):
    args = build_parser().parse_args(argv)
    selection = {
        "client": ["client_compatibility"],
        "normal": ["client_compatibility", "normal_parity"],
        "slow": ["slow_65536_parity"],
        "structural": ["structural_262144"],
        "all": list(STAGE_COMMANDS),
    }
    selected = next(key for key in selection if getattr(args, key))
    report = validate(args.artifact.resolve(), selection[selected])
    print(f"Validation report: {report}")


if __name__ == "__main__":
    main()