#!/usr/bin/env python3
"""Validate fold-rust across installed nightly target-toolchain skew.

The analyzer itself is built against rustc_private from its pinned nightly. Cargo
may select a different target toolchain; fold-rust intentionally keeps using the
analyzer's linked compiler while preserving Cargo's crate graph and arguments.
This check records same-version and newer-target evidence without downloading
or silently substituting toolchains.
"""

from __future__ import annotations

import argparse
import json
import os
import re
import subprocess
import tempfile
from dataclasses import dataclass
from json import JSONDecodeError
from pathlib import Path
from typing import Any

MINIMUM_ANALYZER = (1, 98, 0)
REQUIRED_COMPONENTS = ("rustc-dev", "rustfmt", "llvm-tools-preview")


def component_name(installed: str) -> str:
    if installed.startswith("rustc-dev-"):
        return "rustc-dev"
    if installed.startswith("rustfmt-"):
        return "rustfmt"
    if installed.startswith("llvm-tools-"):
        return "llvm-tools-preview"
    return installed


@dataclass(frozen=True)
class Toolchain:
    name: str
    release: str
    version: tuple[int, int, int]
    nightly: bool
    components: tuple[str, ...]


def version_tuple(release: str) -> tuple[int, int, int] | None:
    match = re.match(r"^(\d+)\.(\d+)\.(\d+)(?:-|$)", release)
    if not match:
        return None
    try:
        major, minor, patch = (int(value) for value in match.groups())
    except ValueError:
        return None
    return major, minor, patch


def run(command: list[str], **kwargs: Any) -> subprocess.CompletedProcess[str]:
    return subprocess.run(
        command,
        check=False,
        capture_output=True,
        text=True,
        **kwargs,
    )


def rustup_toolchains() -> list[Toolchain]:
    listed = run(["rustup", "toolchain", "list"])
    if listed.returncode != 0:
        raise RuntimeError(f"rustup toolchain list failed: {listed.stderr.strip()}")
    result: list[Toolchain] = []
    for line in listed.stdout.splitlines():
        name = line.split(maxsplit=1)[0]
        details = run(["rustup", "run", name, "rustc", "-Vv"])
        if details.returncode != 0:
            continue
        release = next(
            (
                entry.removeprefix("release: ")
                for entry in details.stdout.splitlines()
                if entry.startswith("release: ")
            ),
            "unknown",
        )
        parsed = version_tuple(release)
        if parsed is None:
            continue
        installed = run(
            ["rustup", "component", "list", "--toolchain", name, "--installed"]
        )
        components = tuple(
            sorted(
                {
                    component_name(entry.split()[0])
                    for entry in installed.stdout.splitlines()
                    if entry.strip()
                }
            )
        )
        result.append(
            Toolchain(
                name=name,
                release=release,
                version=parsed,
                nightly="nightly" in release,
                components=components,
            )
        )
    return sorted(result, key=lambda item: (item.version, item.name))


def analyzer_exclusion(toolchain: Toolchain) -> str | None:
    if not toolchain.nightly:
        return "stable rustc cannot provide supported rustc_private APIs"
    if toolchain.version < MINIMUM_ANALYZER:
        return (
            f"nightly {toolchain.release} is below the minimum supported "
            "analyzer release 1.98.0-nightly"
        )
    missing = [
        component
        for component in REQUIRED_COMPONENTS
        if component not in toolchain.components
    ]
    if missing:
        return "missing analyzer components: " + ", ".join(missing)
    return None


def analyzer_info(binary: Path) -> dict[str, Any]:
    completed = run([str(binary), "toolchain-info"])
    if completed.returncode != 0:
        raise RuntimeError(
            f"fold-rust toolchain-info failed: {completed.stderr.strip()}"
        )
    try:
        info = json.loads(completed.stdout)
    except json.JSONDecodeError as error:
        raise RuntimeError(f"invalid fold-rust toolchain-info JSON: {error}") from error
    release = str(info.get("release", ""))
    parsed = version_tuple(release)
    if parsed is None or "nightly" not in release:
        raise RuntimeError(
            "fold-rust requires a nightly rustc_private analyzer; "
            f"binary reports unsupported release {release or '<missing>'}"
        )
    if parsed < MINIMUM_ANALYZER:
        raise RuntimeError(
            f"fold-rust analyzer {release} is unsupported; minimum is 1.98.0-nightly"
        )
    return info


def scenario(binary: Path, toolchain: Toolchain, kind: str) -> dict[str, Any]:
    with tempfile.TemporaryDirectory(prefix="fold-rust-toolchain-") as directory:
        crate = Path(directory)
        (crate / "src").mkdir()
        (crate / "Cargo.toml").write_text(
            '[package]\nname = "fold-rust-toolchain-check"\nversion = "0.1.0"\nedition = "2021"\n',
            encoding="utf-8",
        )
        (crate / "src" / "lib.rs").write_text(
            "pub unsafe fn read_raw(p: *const i32) -> i32 { *p }\n",
            encoding="utf-8",
        )
        env = os.environ.copy()
        env["RUSTUP_TOOLCHAIN"] = toolchain.name
        command = [str(binary), "cargo-check", "--crate-dir", str(crate)]
        try:
            completed = run(command, env=env, timeout=120)
        except subprocess.TimeoutExpired as error:
            return {
                "kind": kind,
                "target_toolchain": toolchain.name,
                "target_release": toolchain.release,
                "status": "failed",
                "reason": f"cargo-check timed out after {error.timeout} seconds",
            }
        records: list[dict[str, Any]] = []
        malformed = 0
        for line in completed.stdout.splitlines():
            if not line.strip():
                continue
            try:
                record = json.loads(line)
            except JSONDecodeError:
                malformed += 1
                continue
            if isinstance(record, dict):
                records.append(record)
        functions = sum(record.get("kind") == "function" for record in records)
        findings = sum(record.get("kind") == "finding" for record in records)
        passed = completed.returncode == 0 and malformed == 0 and functions > 0
        return {
            "kind": kind,
            "target_toolchain": toolchain.name,
            "target_release": toolchain.release,
            "status": "passed" if passed else "failed",
            "exit_status": completed.returncode,
            "records": len(records),
            "functions": functions,
            "findings": findings,
            "malformed_records": malformed,
            "stderr": completed.stderr[-4000:],
            **(
                {}
                if passed
                else {"reason": "cargo-check did not emit a complete function record"}
            ),
        }


def skipped(kind: str, reason: str) -> dict[str, str]:
    return {"kind": kind, "status": "skipped", "reason": reason}


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


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--fold-rust",
        type=Path,
        default=Path("tools/fold-rust/target/release/fold-rust"),
    )
    parser.add_argument(
        "--output", type=Path, default=Path("/tmp/fold-rust-toolchain-skew.json")
    )
    parser.add_argument("--require-skew", action="store_true")
    args = parser.parse_args()

    document: dict[str, Any] = {"schema_version": 1, "scenarios": []}
    try:
        info = analyzer_info(args.fold_rust.resolve())
    except (OSError, RuntimeError) as error:
        document.update({"status": "failed", "error": str(error)})
        write_json(args.output, document)
        print(str(error))
        return 1
    try:
        installed = rustup_toolchains()
    except (OSError, RuntimeError) as error:
        document.update(
            {
                "analyzer": info,
                "status": "failed" if args.require_skew else "skipped",
                "scenarios": [
                    skipped(
                        "toolchain-discovery",
                        f"installed toolchains unavailable: {error}",
                    )
                ],
            }
        )
        write_json(args.output, document)
        print(json.dumps(document, indent=2))
        return 1 if args.require_skew else 0

    analyzer_version = version_tuple(str(info["release"]))
    assert analyzer_version is not None
    document["analyzer"] = info
    document["installed_toolchains"] = [
        {
            "name": item.name,
            "release": item.release,
            "nightly": item.nightly,
            "components": list(item.components),
            "analyzer_eligible": analyzer_exclusion(item) is None,
            **(
                {}
                if analyzer_exclusion(item) is None
                else {"analyzer_exclusion": analyzer_exclusion(item)}
            ),
        }
        for item in installed
    ]

    same = [
        item for item in installed if item.nightly and item.version == analyzer_version
    ]
    if same:
        missing = [
            component
            for component in REQUIRED_COMPONENTS
            if component not in same[-1].components
        ]
        if missing:
            document["scenarios"].append(
                skipped(
                    "same-version",
                    "analyzer toolchain is missing components: " + ", ".join(missing),
                )
            )
        else:
            document["scenarios"].append(
                scenario(args.fold_rust.resolve(), same[-1], "same-version")
            )
    else:
        document["scenarios"].append(
            skipped(
                "same-version",
                f"no installed nightly matches analyzer release {info['release']}",
            )
        )

    newer = [
        item for item in installed if item.nightly and item.version > analyzer_version
    ]
    if newer:
        document["scenarios"].append(
            scenario(args.fold_rust.resolve(), newer[-1], "newer-target")
        )
    else:
        document["scenarios"].append(
            skipped(
                "newer-target",
                f"no installed nightly is newer than analyzer {info['release']}",
            )
        )

    failures = [item for item in document["scenarios"] if item["status"] == "failed"]
    skew_skipped = document["scenarios"][-1]["status"] == "skipped"
    if args.require_skew and skew_skipped:
        failures.append(
            {"reason": "--require-skew requested but no newer nightly is installed"}
        )
    scenario_statuses = {item["status"] for item in document["scenarios"]}
    document["status"] = (
        "failed"
        if failures
        else "passed"
        if "passed" in scenario_statuses
        else "skipped"
    )
    write_json(args.output, document)
    print(
        json.dumps(
            {
                "status": document["status"],
                "artifact": str(args.output),
                "scenarios": document["scenarios"],
            },
            indent=2,
        )
    )
    return 1 if failures else 0


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