#!/usr/bin/env python3
"""Select relevant Python E2E files for the local gate."""

from __future__ import annotations

import argparse
import fnmatch
import json
import os
import re
import subprocess
import sys
from pathlib import Path
import tomllib

ROOT = Path(__file__).resolve().parents[2]
E2E = ROOT / "test" / "e2e"
# run.py records each test's last duration here (test/scratch.py's root)
DURATIONS = (
    Path(os.environ.get("ALBEDO_TEST_TMP", "/tmp/albedo-tests")) / "e2e-durations.json"
)
CAP_SECONDS = 90.0
SMOKE = [
    "daemon_test.py:DaemonTest.test_history_pages_recover_older_turns_once",
    "python_test.py:PythonToolsTests.test_cells_get_the_api_a_model_reaches_for",
    "agents_test.py:AgentsTests.test_reply_and_external_id_mail",
    "lifecycle_test.py:LifecycleTest.test_restart_confirmation_cannot_replace_a_newer_daemon",
]
IGNORED_TOKENS = {
    "src",
    "priv",
    "albedo",
    "daemon",
    "harness",
    "extensions",
    "gleam",
    "py",
    "go",
    "test",
    "internal",
}


def changed():
    commands = [
        ["git", "diff", "--name-only", "origin/main...HEAD"],
        ["git", "diff", "--name-only"],
        ["git", "diff", "--cached", "--name-only"],
        ["git", "ls-files", "--others", "--exclude-standard"],
    ]
    paths = set()
    for command in commands:
        result = subprocess.run(
            command, cwd=ROOT, text=True, capture_output=True, check=False
        )
        paths.update(line for line in result.stdout.splitlines() if line)
    return sorted(paths)


def tokens(path):
    return {
        token.lower()
        for token in re.findall(r"[a-zA-Z0-9]+", path)
        if token.lower() not in IGNORED_TOKENS
    }


def durations():
    try:
        data = json.loads(DURATIONS.read_text())
    except (OSError, ValueError):
        return {}
    result = {}
    for name in E2E.glob("*_test.py"):
        prefix = name.stem + "."
        result[name.name] = sum(
            value for key, value in data.items() if key.startswith(prefix)
        )
        for key, value in data.items():
            if key.startswith(prefix):
                result[f"{name.name}:{key[len(prefix) :]}"] = value
    return result


def select(paths):
    rules = tomllib.loads((E2E / "areas.toml").read_text())["areas"]
    selected, reasons = set(SMOKE), {name: "required smoke coverage" for name in SMOKE}
    candidates = {}
    must_keep = set()
    for path in paths:
        if path in {"test/e2e/harness.py", "test/e2e/run.py"}:
            for name in SMOKE:
                reasons[name] = f"E2E infrastructure change {path}"
            continue
        if path.startswith("test/e2e/") and path.endswith("_test.py"):
            name = Path(path).name
            candidates[name] = f"changed {path}"
            must_keep.add(name)
            continue
        for pattern, files in rules.items():
            if fnmatch.fnmatch(path, pattern):
                for name in files:
                    candidates[name] = f"{path} matches {pattern}"
                    must_keep.add(name)
                break
        path_tokens = tokens(path)
        for test in E2E.glob("*_test.py"):
            shared = path_tokens & tokens(test.stem)
            if shared:
                candidates.setdefault(
                    test.name, f"token match {', '.join(sorted(shared))} from {path}"
                )
    selected.update(candidates)
    selected = {
        name
        for name in selected
        if ":" in name or not any(target.startswith(name + ":") for target in selected)
    }
    must_keep = {
        name
        for name in must_keep
        if not any(target.startswith(name + ":") for target in selected)
    }
    if not candidates and not paths:
        print("no changed paths; using smoke coverage", file=sys.stderr)
    timing = durations()
    total = sum(timing.get(name.split(":", 1)[0], 0) for name in selected)
    if total > CAP_SECONDS:
        kept = set(SMOKE) | must_keep
        budget = max(
            0, CAP_SECONDS - sum(timing.get(name.split(":", 1)[0], 0) for name in kept)
        )
        ranked = sorted(
            (name for name in selected if name not in kept),
            key=lambda name: timing.get(name.split(":", 1)[0], 0),
        )
        for name in ranked:
            cost = timing.get(name.split(":", 1)[0], 0)
            if cost <= budget or not timing:
                kept.add(name)
                budget -= cost
            else:
                print(
                    f"dropped {name} ({cost:.1f}s): local E2E budget is {CAP_SECONDS:.0f}s; run --full",
                    file=sys.stderr,
                )
        selected = kept
    for name in selected:
        reasons.setdefault(name, candidates.get(name, "required smoke coverage"))
    return sorted(selected), reasons


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--print-reasons", action="store_true")
    args = parser.parse_args()
    names, reasons = select(changed())
    for name in names:
        print(name)
        if args.print_reasons:
            print(f"# {name}: {reasons[name]}", file=sys.stderr)
    print(f"selected e2e files ({len(names)}): {', '.join(names)}", file=sys.stderr)


if __name__ == "__main__":
    main()