#!/usr/bin/env python3
"""Inventory and ratchet static performance red flags in first-party code."""
import argparse
import re
import sys
from collections import Counter
from collections import namedtuple
from pathlib import Path
from typing import List, Pattern

REPO = Path(__file__).resolve().parents[2]
PRODUCTION_ROOTS = ("internal", "cmd", "desktop/internal", "web/src")
EXCLUDED_PREFIXES = (
    "cmd/agentteam-e2e-server/",
    "internal/coredb/testkit/",
)
EXCLUDED_PARTS = {
    ".git",
    "node_modules",
    "dist",
    "coverage",
    ".data",
}
SOURCE_SUFFIXES = (".go", ".ts", ".tsx", ".js", ".jsx", ".mjs")

SELECT_STAR_RE = re.compile(r"\bselect\s+\*", re.IGNORECASE)
TIME_SLEEP_RE = re.compile(r"\btime\.Sleep\s*\(")
TIME_NEW_TIMER_RE = re.compile(r"\b(?P<var>\w+)\s*:=\s*time\.NewTimer\s*\(")
DB_CALL_RE = re.compile(
    r"\b(?P<receiver>[A-Za-z_]\w*(?:\.[A-Za-z_]\w*)*)\."
    r"(?P<method>Query|QueryContext|Exec|ExecContext)\s*\("
)
NON_DB_CALL_RECEIVERS = {
    "Gin",
    "URL",
}
IGNORED_DB_AGENTTEAM_HITS = {
    (
        "desktop/internal/appbuild/templates/storage/sqlite/pkg/store/sqlite.go",
        "if _, err := db.Exec(s); err != nil {",
    ),
}


Hit = namedtuple("Hit", "file line code")


def rel(path: Path) -> str:
    return path.relative_to(REPO).as_posix()


def is_excluded(path: Path) -> bool:
    path_rel = rel(path)
    if any(path_rel.startswith(prefix) for prefix in EXCLUDED_PREFIXES):
        return True
    return bool(EXCLUDED_PARTS.intersection(path.parts))


def is_test_file(path: Path) -> bool:
    name = path.name
    return name.endswith("_test.go") or ".test." in name or ".spec." in name


def source_files() -> List[Path]:
    files: List[Path] = []
    for root_name in PRODUCTION_ROOTS:
        root = REPO / root_name
        if not root.exists():
            continue
        for path in root.rglob("*"):
            if not path.is_file() or is_excluded(path) or is_test_file(path):
                continue
            if path.suffix in SOURCE_SUFFIXES:
                files.append(path)
    return sorted(files)


def go_files() -> List[Path]:
    return [path for path in source_files() if path.suffix == ".go"]


def scan_pattern(pattern: Pattern[str], files: List[Path]) -> List[Hit]:
    hits: List[Hit] = []
    for path in files:
        for lineno, raw in enumerate(path.read_text(errors="ignore").splitlines(), 1):
            if pattern.search(raw):
                hits.append(Hit(rel(path), lineno, raw.strip()))
    return hits


def line_has_db_call(raw: str) -> bool:
    for match in DB_CALL_RE.finditer(raw):
        receiver_parts = match.group("receiver").split(".")
        if receiver_parts[-1] in NON_DB_CALL_RECEIVERS:
            continue
        return True
    return False


def scan_db_calls_in_agentteams() -> List[Hit]:
    hits: List[Hit] = []
    for path in go_files():
        path_rel = rel(path)
        agentteam_stack: List[int] = []
        depth = 0
        for lineno, raw in enumerate(path.read_text(errors="ignore").splitlines(), 1):
            stripped = raw.strip()
            opens = raw.count("{")
            closes = raw.count("}")

            if stripped.startswith("for ") or stripped in {"for {", "for{"}:
                agentteam_stack.append(depth + opens)

            if agentteam_stack and line_has_db_call(raw):
                if (path_rel, stripped) not in IGNORED_DB_AGENTTEAM_HITS:
                    hits.append(Hit(path_rel, lineno, stripped))

            depth += opens - closes
            while agentteam_stack and depth < agentteam_stack[-1]:
                agentteam_stack.pop()
    return hits


def scan_uncancellable_timer_waits() -> List[Hit]:
    hits: List[Hit] = []
    for path in go_files():
        lines = path.read_text(errors="ignore").splitlines()
        for index, raw in enumerate(lines):
            match = TIME_NEW_TIMER_RE.search(raw)
            if not match:
                continue
            timer_var = match.group("var")
            direct_receive = re.compile(rf"^(?:\w+\s*(?::=|=)\s*)?<-\s*{re.escape(timer_var)}\.C\b")
            for lookahead_index in range(index + 1, min(index + 8, len(lines))):
                stripped = lines[lookahead_index].strip()
                if not stripped or stripped.startswith("//") or stripped.startswith("defer "):
                    continue
                if stripped.startswith("select"):
                    break
                if direct_receive.search(stripped):
                    hits.append(Hit(rel(path), lookahead_index + 1, stripped))
                    break
                break
    return hits


def print_inventory(name: str, hits: List[Hit], *, max_examples: int = 10) -> None:
    print(f"{name}: {len(hits)}")
    if not hits:
        return
    by_file = Counter(hit.file for hit in hits)
    top_files = by_file.most_common(max_examples)
    for file, count in top_files:
        print(f"  {count:4d} {file}")
    print("  examples:")
    for file, _ in top_files:
        hit = next(hit for hit in hits if hit.file == file)
        print(f"       {hit.file}:{hit.line}: {hit.code}")


def fail_if_over(name: str, actual: int, maximum: int, failures: List[str]) -> None:
    print(f"{name}: {actual} (max {maximum})")
    if actual > maximum:
        failures.append(f"{name} grew above baseline: {actual} > {maximum}")


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--max-select-star", type=int, required=True)
    parser.add_argument("--max-production-sleeps", type=int, required=True)
    parser.add_argument("--max-uncancellable-timer-waits", type=int, required=True)
    parser.add_argument("--max-db-calls-in-agentteams", type=int, required=True)
    args = parser.parse_args()

    files = source_files()
    select_star = scan_pattern(SELECT_STAR_RE, files)
    production_sleeps = scan_pattern(TIME_SLEEP_RE, go_files())
    uncancellable_timer_waits = scan_uncancellable_timer_waits()
    db_calls_in_agentteams = scan_db_calls_in_agentteams()

    print("== Static performance red-flag inventory ==")
    print_inventory("select_star", select_star)
    print_inventory("production_sleeps", production_sleeps)
    print_inventory("uncancellable_timer_waits", uncancellable_timer_waits)
    print_inventory("db_calls_in_agentteams", db_calls_in_agentteams)

    failures: List[str] = []
    print("\n== Ratchet baselines ==")
    fail_if_over("select_star", len(select_star), args.max_select_star, failures)
    fail_if_over("production_sleeps", len(production_sleeps), args.max_production_sleeps, failures)
    fail_if_over(
        "uncancellable_timer_waits",
        len(uncancellable_timer_waits),
        args.max_uncancellable_timer_waits,
        failures,
    )
    fail_if_over("db_calls_in_agentteams", len(db_calls_in_agentteams), args.max_db_calls_in_agentteams, failures)

    if failures:
        for failure in failures:
            print(f"FAIL: {failure}", file=sys.stderr)
        print(f"\nStatic performance guardrail failed with {len(failures)} issue(s).", file=sys.stderr)
        return 1

    print("\nStatic performance guardrail passed.")
    return 0


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