#!/usr/bin/env python3
"""Guardrails for the store-domain split roadmap.

This script intentionally tracks only structural metrics. It must not enforce
runtime behavior; package-specific tests own behavior. The baseline is a ratchet:
oversized package metrics may go down, but should not grow while the split is in
progress unless the baseline is deliberately updated in a planning PR.
"""
from __future__ import annotations

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

REPO = Path(__file__).resolve().parents[2]
DEFAULT_BASELINE_PATH = REPO / "guardrails/store/store-domain-split-debt-baseline.json"
DEFAULT_WRITE_BASELINE_PATH = Path("/tmp/agentteam-store-domain-split-baseline.json")

TARGET_PACKAGES = [
    "internal/store/app",
    "internal/store/taskboard",
    "internal/store/organization",
    "internal/store/message",
    "internal/store/privateagent",
]
TARGET_METRICS = {
    "exported_methods",
    "exported_symbols",
    "external_prod_import_files",
    "go_files",
    "loc",
    "sibling_store_imports",
    "test_files",
    "test_loc",
}

GENERIC_CHILD_PACKAGE_NAMES = {"helpers", "ports", "repository"}
GENERIC_CHILD_PACKAGE_ALLOWLIST: set[str] = set()
FORBIDDEN_STORE_PATHS = {
    "internal/store/app/ports",
    "internal/store/message/diagnostics",
    "internal/store/message/readstate",
    "internal/store/message/search",
    "internal/store/runtimecatalog",
    "internal/store/tasksearch",
}

FORBIDDEN_STORE_IMPORT_PREFIXES = (
    "github.com/tidbcloud/AgentTeam/internal/api",
    "github.com/tidbcloud/AgentTeam/internal/router",
    "github.com/tidbcloud/AgentTeam/internal/server",
    "github.com/tidbcloud/AgentTeam/internal/usecase",
    "github.com/tidbcloud/AgentTeam/cmd",
    "github.com/tidbcloud/AgentTeam/web",
    "github.com/tidbcloud/AgentTeam/desktop",
)

MODULE_PREFIX = "github.com/tidbcloud/AgentTeam/"
STORE_IMPORT_PREFIX = "github.com/tidbcloud/AgentTeam/internal/store/"
SERVICE_IMPORT_PREFIX = "github.com/tidbcloud/AgentTeam/internal/service/"
DEBT_METADATA_FIELDS = ("owner", "acceptedReason", "retirementCondition")


def git_head() -> str:
    try:
        return subprocess.check_output(["git", "rev-parse", "HEAD"], cwd=REPO, text=True).strip()
    except Exception:
        return "unknown"


def go_files(pkg: str, tests: bool | None = None) -> list[Path]:
    root = REPO / pkg
    if not root.exists():
        return []
    files: list[Path] = []
    for path in root.glob("*.go"):
        is_test = path.name.endswith("_test.go")
        if tests is None or tests == is_test:
            files.append(path)
    return sorted(files)


def line_count(paths: list[Path]) -> int:
    return sum(len(path.read_text(errors="ignore").splitlines()) for path in paths)


def parse_imports(path: Path) -> list[str]:
    text = path.read_text(errors="ignore")
    try:
        tree = ast.parse(text)  # type: ignore[arg-type]
    except SyntaxError:
        # Go is not Python; use a lightweight regex fallback below.
        pass
    imports: list[str] = []
    in_block = False
    for raw in text.splitlines():
        line = raw.strip()
        if line.startswith("import ("):
            in_block = True
            continue
        if in_block and line == ")":
            in_block = False
            continue
        if in_block:
            match = re.search(r'"([^"]+)"', line)
            if match:
                imports.append(match.group(1))
        else:
            match = re.match(r'import\s+(?:[A-Za-z_][A-Za-z0-9_]*\s+)?"([^"]+)"', line)
            if match:
                imports.append(match.group(1))
    return imports


def count_exported_symbols(paths: list[Path]) -> tuple[int, int]:
    exported = 0
    exported_methods = 0
    block_kind: str | None = None
    name_re = re.compile(r"[A-Z][A-Za-z0-9_]*")
    for path in paths:
        for raw in path.read_text(errors="ignore").splitlines():
            line = raw.strip()
            if line.startswith("//") or not line:
                continue
            if block_kind:
                if line.startswith(")"):
                    block_kind = None
                    continue
                if name_re.match(line):
                    exported += 1
                continue
            if re.match(r"^(const|var|type)\s*\(", line):
                block_kind = line.split("(", 1)[0].strip()
                continue
            if re.match(r"^(const|var|type)\s+[A-Z]", line):
                exported += 1
                continue
            if re.match(r"^func\s+[A-Z]", line):
                exported += 1
                continue
            if re.match(r"^func\s+\([^)]*\)\s+[A-Z]", line):
                exported += 1
                exported_methods += 1
                continue
    return exported, exported_methods


def package_metrics(pkg: str) -> dict[str, int]:
    prod = go_files(pkg, tests=False)
    test = go_files(pkg, tests=True)
    exported, exported_methods = count_exported_symbols(prod)
    unique_sibling_imports: set[str] = set()
    for path in prod:
        for imp in parse_imports(path):
            if imp.startswith(STORE_IMPORT_PREFIX) and not imp.startswith(MODULE_PREFIX + pkg + "/"):
                unique_sibling_imports.add(imp)
    return {
        "go_files": len(prod),
        "loc": line_count(prod),
        "test_files": len(test),
        "test_loc": line_count(test),
        "exported_symbols": exported,
        "exported_methods": exported_methods,
        "sibling_store_imports": len(unique_sibling_imports),
        "external_prod_import_files": count_external_import_files(pkg),
    }


def count_external_import_files(pkg: str) -> int:
    import_path = MODULE_PREFIX + pkg
    count = 0
    for path in REPO.rglob("*.go"):
        rel = path.relative_to(REPO).as_posix()
        if (
            rel.startswith(".git/")
            or rel.startswith(".worktrees/")
            or rel.startswith("third_party/")
            or rel.startswith("x/")
            or rel.startswith("vendor/")
            or rel.endswith("_test.go")
            or rel.startswith(pkg + "/")
        ):
            continue
        if f'"{import_path}"' in path.read_text(errors="ignore"):
            count += 1
    return count


def child_packages_under(pkg: str) -> list[str]:
    root = REPO / pkg
    if not root.exists():
        return []
    children: list[str] = []
    for directory in sorted(path for path in root.rglob("*") if path.is_dir()):
        if any(directory.glob("*.go")):
            children.append(directory.relative_to(REPO).as_posix())
    return children


def forbidden_store_paths() -> list[str]:
    forbidden: list[str] = []
    for rel in sorted(FORBIDDEN_STORE_PATHS):
        if (REPO / rel).exists():
            forbidden.append(rel)
    return forbidden


def generic_child_packages() -> list[str]:
    out: list[str] = []
    store_root = REPO / "internal/store"
    if not store_root.exists():
        return out
    for directory in sorted(path for path in store_root.rglob("*") if path.is_dir()):
        rel = directory.relative_to(REPO).as_posix()
        if directory.name not in GENERIC_CHILD_PACKAGE_NAMES:
            continue
        if rel in GENERIC_CHILD_PACKAGE_ALLOWLIST:
            continue
        if any(directory.glob("*.go")):
            out.append(rel)
    return out


def forbidden_store_import_files() -> list[str]:
    out: list[str] = []
    for path in (REPO / "internal/store").rglob("*.go"):
        rel = path.relative_to(REPO).as_posix()
        imports = parse_imports(path)
        if any(any(imp.startswith(prefix) for prefix in FORBIDDEN_STORE_IMPORT_PREFIXES) for imp in imports):
            out.append(rel)
    return sorted(out)


def production_store_to_service_edges() -> list[dict[str, str]]:
    edges: set[tuple[str, str]] = set()
    for path in (REPO / "internal/store").rglob("*.go"):
        if path.name.endswith("_test.go"):
            continue
        importer = path.parent.relative_to(REPO).as_posix()
        for imp in parse_imports(path):
            if imp.startswith(SERVICE_IMPORT_PREFIX):
                edges.add((importer, imp.removeprefix(MODULE_PREFIX)))
    return [{"from": source, "to": target} for source, target in sorted(edges)]


def collect() -> dict[str, Any]:
    child_baseline: dict[str, list[str]] = {pkg: child_packages_under(pkg) for pkg in TARGET_PACKAGES}
    return {
        "schemaVersion": 2,
        "baseCommit": git_head(),
        "targets": {pkg: package_metrics(pkg) for pkg in TARGET_PACKAGES},
        "existingChildPackages": child_baseline,
        "plannedChildPackages": [],
        "forbiddenStoreImportFiles": forbidden_store_import_files(),
        "forbiddenStorePaths": forbidden_store_paths(),
        "genericChildPackages": generic_child_packages(),
        "productionStoreToServiceEdges": production_store_to_service_edges(),
    }


def write_baseline(path: Path) -> None:
    data = collect()
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(json.dumps(data, ensure_ascii=False, indent=2, sort_keys=True) + "\n")


def load_baseline(path: Path | None) -> dict[str, Any]:
    baseline_path = path or DEFAULT_BASELINE_PATH
    try:
        return json.loads(baseline_path.read_text())
    except FileNotFoundError:
        print(f"error: baseline not found: {baseline_path}", file=sys.stderr)
        return {}


def baseline_validation_errors(baseline: dict[str, Any]) -> list[str]:
    errors: list[str] = []
    if baseline.get("schemaVersion") != 2:
        errors.append("baseline schemaVersion must be 2")
    if not baseline.get("baseCommit"):
        errors.append("baseline baseCommit must identify its inventory provenance")

    metadata = baseline.get("debtMetadata", {})
    for pkg in TARGET_PACKAGES:
        metrics = baseline.get("targets", {}).get(pkg)
        if metrics is None:
            errors.append(f"baseline target missing: {pkg}")
        else:
            metric_names = set(metrics)
            for metric in sorted(TARGET_METRICS - metric_names):
                errors.append(f"baseline target metric missing: {pkg} {metric}")
            for metric in sorted(metric_names - TARGET_METRICS):
                errors.append(f"baseline target metric unknown: {pkg} {metric}")
            for metric in sorted(TARGET_METRICS & metric_names):
                value = metrics[metric]
                if type(value) is not int or value < 0:
                    errors.append(
                        f"baseline target metric must be a non-negative integer: "
                        f"{pkg} {metric}={value!r}"
                    )
        item = metadata.get(pkg, {})
        for field in DEBT_METADATA_FIELDS:
            if not str(item.get(field, "")).strip():
                errors.append(f"baseline debt metadata missing {field}: {pkg}")

    for edge in baseline.get("productionStoreToServiceEdges", []):
        label = f"{edge.get('from', '?')} -> {edge.get('to', '?')}"
        for field in ("from", "to", *DEBT_METADATA_FIELDS):
            if not str(edge.get(field, "")).strip():
                errors.append(f"store-to-service edge metadata missing {field}: {label}")
    return errors


def edge_pairs(entries: list[dict[str, Any]]) -> set[tuple[str, str]]:
    return {
        (str(entry.get("from", "")), str(entry.get("to", "")))
        for entry in entries
        if entry.get("from") and entry.get("to")
    }


def store_to_service_edge_failures(
    expected_entries: list[dict[str, Any]], current_entries: list[dict[str, Any]]
) -> list[str]:
    expected = edge_pairs(expected_entries)
    current = edge_pairs(current_entries)
    failures = [
        f"new production store-to-service edge: {source} -> {target}"
        for source, target in sorted(current - expected)
    ]
    failures.extend(
        f"stale production store-to-service debt entry: {source} -> {target}"
        for source, target in sorted(expected - current)
    )
    return failures


def unexpected_child_packages(
    baseline: dict[str, Any], current: dict[str, Any]
) -> list[str]:
    baseline_children = {
        child
        for children in baseline.get("existingChildPackages", {}).values()
        for child in children
    }
    approved_new_children = set(baseline.get("plannedChildPackages", []))
    current_children = {
        child
        for children in current.get("existingChildPackages", {}).values()
        for child in children
    }
    return sorted(current_children - baseline_children - approved_new_children)


def check(baseline: dict[str, Any], current: dict[str, Any]) -> int:
    failures = baseline_validation_errors(baseline)
    for pkg, expected in baseline.get("targets", {}).items():
        actual = current["targets"].get(pkg)
        if actual is None:
            failures.append(f"{pkg}: target package missing")
            continue
        for metric, max_value in expected.items():
            actual_value = actual.get(metric, 0)
            print(f"{pkg} {metric}: {actual_value} (max {max_value})")
            if actual_value > max_value:
                failures.append(f"{pkg} {metric} grew: {actual_value} > {max_value}")

    unexpected_children = unexpected_child_packages(baseline, current)
    print(f"unexpected store split child packages: {len(unexpected_children)}")
    for child in unexpected_children:
        print(f"  {child}")
    if unexpected_children:
        failures.append("unexpected child packages under oversized store domains")

    forbidden = current.get("forbiddenStoreImportFiles", [])
    print(f"internal/store upward imports: {len(forbidden)} (max 0)")
    for rel in forbidden:
        print(f"  {rel}")
    if forbidden:
        failures.append("internal/store has forbidden upward imports")

    forbidden_paths = current.get("forbiddenStorePaths", [])
    print(f"forbidden store paths: {len(forbidden_paths)} (max 0)")
    for rel in forbidden_paths:
        print(f"  {rel}")
    if forbidden_paths:
        failures.append("internal/store contains forbidden provider-side path")

    generic_children = current.get("genericChildPackages", [])
    print(f"generic store child packages: {len(generic_children)} (max 0)")
    for rel in generic_children:
        print(f"  {rel}")
    if generic_children:
        failures.append("internal/store contains generic child package names")

    expected_edges = baseline.get("productionStoreToServiceEdges", [])
    current_edges = current.get("productionStoreToServiceEdges", [])
    print(
        "production internal/store -> internal/service edges: "
        f"{len(current_edges)} (allowed debt {len(expected_edges)}, target 0)"
    )
    for edge in current_edges:
        print(f"  {edge['from']} -> {edge['to']}")
    failures.extend(store_to_service_edge_failures(expected_edges, current_edges))

    if failures:
        print("\nFAILED store domain split guardrails:", file=sys.stderr)
        for failure in failures:
            print(f"- {failure}", file=sys.stderr)
        return 1
    print("store domain split guardrails: ok")
    return 0


def print_markdown(data: dict[str, Any]) -> None:
    print("| Package | Go files | LOC | Test files | Exported symbols | Exported methods | Sibling store imports | External prod import files |")
    print("|---|---:|---:|---:|---:|---:|---:|---:|")
    for pkg in TARGET_PACKAGES:
        m = data["targets"][pkg]
        print(
            f"| `{pkg}` | {m['go_files']} | {m['loc']} | {m['test_files']} | "
            f"{m['exported_symbols']} | {m['exported_methods']} | {m['sibling_store_imports']} | {m['external_prod_import_files']} |"
        )


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--baseline",
        type=Path,
        default=None,
        help=f"Optional baseline JSON path. Defaults to {DEFAULT_BASELINE_PATH.relative_to(REPO)}.",
    )
    parser.add_argument("--write-baseline", action="store_true")
    parser.add_argument("--markdown", action="store_true")
    args = parser.parse_args()

    if args.write_baseline:
        baseline_path = args.baseline or DEFAULT_WRITE_BASELINE_PATH
        write_baseline(baseline_path)
        print(f"wrote {baseline_path}")
        return 0

    current = collect()
    if args.markdown:
        print_markdown(current)
        return 0

    baseline = load_baseline(args.baseline)
    if not baseline:
        return 1
    return check(baseline, current)


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