#!/usr/bin/env python3
"""Check the pinned Matplotlib inventory and generate compatibility docs."""

from __future__ import annotations

import argparse
import ast
import json
import re
import subprocess
import sys
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
SNAPSHOT = ROOT / "tests/pyplot/matplotlib_311_plotting.json"
METADATA = ROOT / "tests/pyplot/compatibility.json"
CORPUS = ROOT / "tests/pyplot/corpus"
OUTPUT = ROOT / "spec/matplotlib/compat-matrix.md"


def _load(path: Path) -> dict:
    return json.loads(path.read_text(encoding="utf-8"))


def _upstream_inventory(checkout: Path) -> tuple[dict[str, list[str]], str, str]:
    snapshot = _load(SNAPSHOT)
    source = checkout / snapshot["upstream"]["source"]
    text = source.read_text(encoding="utf-8")
    plotting = text.split("\nPlotting\n========\n", 1)[1].split("\nClearing\n========\n", 1)[0]
    excluded = set(snapshot["excluded"])
    families: dict[str, list[str]] = {}
    family = ""
    lines = plotting.splitlines()
    for index, line in enumerate(lines):
        if index + 1 < len(lines) and re.fullmatch(r"-+", lines[index + 1]):
            family = line
            families[family] = []
        match = re.fullmatch(r"\s+Axes\.([A-Za-z0-9_]+)", line)
        if match and match.group(1) not in excluded:
            families[family].append(match.group(1))
    revision = subprocess.run(
        ["git", "rev-parse", "--short=10", "HEAD"],
        cwd=checkout,
        check=True,
        capture_output=True,
        text=True,
    ).stdout.strip()
    describe = subprocess.run(
        ["git", "describe", "--tags", "--always"],
        cwd=checkout,
        check=True,
        capture_output=True,
        text=True,
    ).stdout.strip()
    return families, revision, describe


def _root_name(node: ast.AST) -> str | None:
    while isinstance(node, (ast.Attribute, ast.Subscript)):
        node = node.value
    return node.id if isinstance(node, ast.Name) else None


def _axes_receivers(tree: ast.AST) -> set[str]:
    """Find names that are demonstrably populated with Axes objects."""
    names = {"ax", "axes"}
    for node in ast.walk(tree):
        if isinstance(node, (ast.Assign, ast.AnnAssign)):
            value = node.value
            targets = node.targets if isinstance(node, ast.Assign) else [node.target]
            if isinstance(value, ast.Call) and isinstance(value.func, ast.Attribute):
                if value.func.attr in {
                    "subplot",
                    "gca",
                    "axes",
                    "add_subplot",
                    "add_axes",
                    "inset_axes",
                    "twinx",
                    "twiny",
                }:
                    for target in targets:
                        if isinstance(target, ast.Name):
                            names.add(target.id)
                if value.func.attr == "subplots":
                    for target in targets:
                        if isinstance(target, (ast.Tuple, ast.List)) and len(target.elts) >= 2:
                            axes_target = target.elts[1]
                            if isinstance(axes_target, ast.Name):
                                names.add(axes_target.id)
        elif isinstance(node, ast.For) and isinstance(node.target, ast.Name):
            if _root_name(node.iter) in names:
                names.add(node.target.id)
    return names


def _corpus_calls() -> dict[str, list[str]]:
    calls: dict[str, list[str]] = {}
    for path in sorted(CORPUS.glob("[0-9][0-9]_*.py")):
        tree = ast.parse(path.read_text(encoding="utf-8"), filename=str(path))
        axes_names = _axes_receivers(tree)
        for node in ast.walk(tree):
            if isinstance(node, ast.Call) and isinstance(node.func, ast.Attribute):
                receiver = _root_name(node.func.value)
                if receiver != "plt" and receiver not in axes_names:
                    continue
                names = calls.setdefault(node.func.attr, [])
                if path.name not in names:
                    names.append(path.name)
    return calls


def _render() -> str:
    snapshot = _load(SNAPSHOT)
    metadata = _load(METADATA)
    calls = _corpus_calls()
    lines = [
        "<!-- Generated by scripts/sync_matplotlib_compat.py; do not edit. -->",
        "# Matplotlib compatibility matrix",
        "",
        f"Pinned upstream: `{snapshot['upstream']['describe']}` (`{snapshot['upstream']['revision']}`).",
        "",
        "The level describes the intended compatibility contract, not pixel identity.",
        "Corpus links are executable examples and are checked for every supported method.",
        "",
        "| Family | Level | Methods | Executable corpus |",
        "|---|---|---:|---:|",
    ]
    for family, methods in snapshot["families"].items():
        examples = sorted({name for method in methods for name in calls.get(method, [])})
        level = metadata["families"][family]["level"]
        lines.append(f"| {family} | {level} | {len(methods)} | {len(examples)} |")
    lines += ["", "## Method inventory", ""]
    for family, methods in snapshot["families"].items():
        level = metadata["families"][family]["level"]
        lines += [f"### {family}", "", f"Approximation level: **{level}**.", ""]
        for method in methods:
            examples = calls.get(method, [])
            refs = ", ".join(f"[`{name}`](../../tests/pyplot/corpus/{name})" for name in examples)
            lines.append(f"- `{method}` — {refs or '**missing executable corpus coverage**'}")
        lines.append("")
    return "\n".join(lines)


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--upstream", type=Path)
    parser.add_argument("--check", action="store_true")
    parser.add_argument(
        "--update-snapshot",
        action="store_true",
        help="replace the reviewed snapshot from --upstream after moving the pin",
    )
    args = parser.parse_args()
    snapshot = _load(SNAPSHOT)
    errors: list[str] = []
    if args.update_snapshot and not args.upstream:
        parser.error("--update-snapshot requires --upstream")
    if args.upstream:
        families, revision, describe = _upstream_inventory(args.upstream)
        if args.update_snapshot:
            snapshot["upstream"]["revision"] = revision
            snapshot["upstream"]["describe"] = describe
            snapshot["families"] = families
            SNAPSHOT.write_text(json.dumps(snapshot, indent=2) + "\n", encoding="utf-8")
        else:
            if revision != snapshot["upstream"]["revision"]:
                errors.append(
                    f"upstream revision is {revision}, expected {snapshot['upstream']['revision']}"
                )
            if families != snapshot["families"]:
                errors.append("upstream Axes Plotting inventory differs from the reviewed snapshot")
    rendered = _render()
    if args.check:
        if not OUTPUT.exists() or OUTPUT.read_text(encoding="utf-8") != rendered:
            errors.append(f"{OUTPUT.relative_to(ROOT)} is stale; run {Path(__file__).name}")
    else:
        OUTPUT.write_text(rendered, encoding="utf-8")
    if errors:
        print("\n".join(errors), file=sys.stderr)
        return 1
    return 0


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