#!/usr/bin/env python3
"""Audit unnecessary explicit aliases on local Go imports.

An explicit import alias is allowed only when removing it would create a name
conflict in the same file. The fast syntax audit detects:

- another import whose effective/default package name is the same; and
- the current file's own package name matching the imported package name.

The reviewed baseline is reserved for aliases that a separate `go/types`
cleanup proved are required by lexical local-name shadowing (or whose
build-constrained file cannot participate in the repository-wide type load).
It must never be regenerated as a blanket allowlist.

The audit intentionally ignores generated files, vendor/worktree/data folders,
and all code under x/ because those trees are generated, externally
synchronized, or local runtime snapshots.
"""

from __future__ import annotations

import argparse
import ast
import re
import sys
from dataclasses import dataclass
from pathlib import Path
from typing import Iterable, Sequence

DEFAULT_IMPORT_PREFIXES = ("github.com/tidbcloud/AgentTeam",)
IGNORE_DIRS = {
    ".claude",
    ".data",
    ".git",
    ".worktrees",
    "node_modules",
    "vendor",
    "x",
}
GENERATED_SUFFIXES = (
    ".pb.go",
    ".pb.gw.go",
    ".proto.mock.go",
    ".gen.go",
    "_generated.go",
)

SINGLE_IMPORT_RE = re.compile(
    r'^\s*import\s+(?:(?P<alias>[A-Za-z_]\w*|[._])\s+)?(?P<path>"(?:[^"\\]|\\.)*")'
)
BLOCK_START_RE = re.compile(r"^\s*import\s*\(\s*(?://.*)?$")
IMPORT_SPEC_RE = re.compile(
    r'^\s*(?:(?P<alias>[A-Za-z_]\w*|[._])\s+)?(?P<path>"(?:[^"\\]|\\.)*")'
)
PACKAGE_RE = re.compile(r"^\s*package\s+([A-Za-z_]\w*)\b")


@dataclass(frozen=True)
class ImportSpec:
    alias: str | None
    path: str
    line: int


@dataclass(frozen=True)
class Finding:
    file: Path
    line: int
    alias: str
    default_name: str
    path: str

    def baseline_key(self) -> str:
        return f"{self.file} alias={self.alias} default={self.default_name} path={self.path}"


@dataclass(frozen=True)
class RetainedAlias:
    file: Path
    line: int
    alias: str
    default_name: str
    path: str
    conflicts: tuple[str, ...]


def repo_root() -> Path:
    return Path(__file__).resolve().parents[1]


def has_ignored_part(path: Path) -> bool:
    return any(part in IGNORE_DIRS for part in path.parts)


def is_generated_file(path: Path) -> bool:
    name = path.name
    if name.endswith(GENERATED_SUFFIXES):
        return True
    try:
        head = "\n".join(path.read_text(errors="ignore").splitlines()[:20])
    except OSError:
        return False
    return "Code generated" in head and "DO NOT EDIT" in head


def parse_module_path(go_mod: Path) -> str | None:
    for raw_line in go_mod.read_text().splitlines():
        line = raw_line.strip()
        if line.startswith("module "):
            return line.split()[1]
    return None


def iter_module_roots(root: Path) -> list[tuple[Path, str]]:
    modules: list[tuple[Path, str]] = []
    for gomod in root.rglob("go.mod"):
        rel = gomod.relative_to(root)
        if has_ignored_part(rel):
            continue
        module_path = parse_module_path(gomod)
        if module_path:
            modules.append((gomod.parent, module_path))
    return sorted(modules, key=lambda item: len(item[0].relative_to(root).parts))


def is_relative_to(path: Path, parent: Path) -> bool:
    try:
        path.relative_to(parent)
    except ValueError:
        return False
    return True


def nearest_module_for(path: Path, modules: Sequence[tuple[Path, str]]) -> tuple[Path, str] | None:
    matches = [module for module in modules if is_relative_to(path, module[0])]
    if not matches:
        return None
    return max(matches, key=lambda item: len(item[0].parts))


def package_name_for_dir(directory: Path) -> str | None:
    names: list[str] = []
    for go_file in sorted(directory.glob("*.go")):
        if is_generated_file(go_file):
            continue
        name = parse_package_name(go_file)
        if not name:
            continue
        if not go_file.name.endswith("_test.go") and not name.endswith("_test"):
            return name
        names.append(name.removesuffix("_test"))
    return names[0] if names else None


def load_package_names(root: Path) -> dict[str, str]:
    # Resolve declared package names from local source files instead of `go list`.
    # This keeps the guardrail fast and avoids module-specific setup such as
    # generated web dist assets needed by some embed packages.
    modules = iter_module_roots(root)
    package_names: dict[str, str] = {}
    for directory in sorted({path.parent for path in iter_go_files(root)}):
        module = nearest_module_for(directory, modules)
        if module is None:
            continue
        module_root, module_path = module
        rel = directory.relative_to(module_root)
        import_path = module_path if rel == Path(".") else f"{module_path}/{rel.as_posix()}"
        name = package_name_for_dir(directory)
        if name:
            package_names[import_path] = name
    return package_names


def fallback_package_name(import_path: str) -> str:
    return import_path.rsplit("/", 1)[-1].replace("-", "_")


def package_name(import_path: str, package_names: dict[str, str]) -> str:
    return package_names.get(import_path) or fallback_package_name(import_path)


def unquote_go_string(value: str) -> str:
    return ast.literal_eval(value)


def remove_line_comment(line: str) -> str:
    # Go import paths are quoted and repository import paths do not contain //.
    # This keeps parsing simple while allowing trailing comments on import specs.
    return line.split("//", 1)[0]


def parse_package_name(path: Path) -> str | None:
    for raw_line in path.read_text(errors="ignore").splitlines():
        match = PACKAGE_RE.match(raw_line)
        if match:
            return match.group(1)
    return None


def parse_imports(path: Path) -> list[ImportSpec]:
    imports: list[ImportSpec] = []
    in_import_block = False
    for line_no, raw_line in enumerate(path.read_text(errors="ignore").splitlines(), 1):
        line = re.sub(r"/\*.*?\*/", " ", raw_line)
        if not in_import_block:
            single = SINGLE_IMPORT_RE.match(remove_line_comment(line))
            if single:
                imports.append(
                    ImportSpec(
                        alias=single.group("alias"),
                        path=unquote_go_string(single.group("path")),
                        line=line_no,
                    )
                )
                continue
            if BLOCK_START_RE.match(line):
                in_import_block = True
            continue

        if line.strip().startswith(")"):
            in_import_block = False
            continue
        spec = IMPORT_SPEC_RE.match(remove_line_comment(line))
        if spec:
            imports.append(
                ImportSpec(
                    alias=spec.group("alias"),
                    path=unquote_go_string(spec.group("path")),
                    line=line_no,
                )
            )
    return imports


def iter_go_files(root: Path) -> Iterable[Path]:
    for path in root.rglob("*.go"):
        relative = path.relative_to(root)
        if has_ignored_part(relative):
            continue
        if is_generated_file(path):
            continue
        yield path


def is_included_import(import_path: str, prefixes: Sequence[str]) -> bool:
    return any(import_path == prefix or import_path.startswith(prefix + "/") for prefix in prefixes)


def audit(
    root: Path,
    package_names: dict[str, str],
    include_prefixes: Sequence[str] = DEFAULT_IMPORT_PREFIXES,
) -> tuple[list[Finding], list[RetainedAlias], int]:
    findings: list[Finding] = []
    retained: list[RetainedAlias] = []
    explicit_local_aliases = 0

    for path in iter_go_files(root):
        imports = parse_imports(path)
        if not imports:
            continue
        own_package_name = parse_package_name(path)

        for spec in imports:
            if (
                not is_included_import(spec.path, include_prefixes)
                or spec.alias is None
                or spec.alias in {"_", "."}
            ):
                continue
            explicit_local_aliases += 1
            default_name = package_name(spec.path, package_names)
            conflicts = [
                other.path
                for other in imports
                if other is not spec
                and other.path != spec.path
                and package_name(other.path, package_names) == default_name
            ]
            if own_package_name == default_name:
                conflicts.append(f"file package {own_package_name}")
            if conflicts:
                retained.append(
                    RetainedAlias(
                        file=path.relative_to(root),
                        line=spec.line,
                        alias=spec.alias,
                        default_name=default_name,
                        path=spec.path,
                        conflicts=tuple(conflicts),
                    )
                )
                continue
            findings.append(
                Finding(
                    file=path.relative_to(root),
                    line=spec.line,
                    alias=spec.alias,
                    default_name=default_name,
                    path=spec.path,
                )
            )
    return findings, retained, explicit_local_aliases


def load_baseline(path: Path | None) -> set[str]:
    if path is None:
        return set()
    baseline: set[str] = set()
    for raw_line in path.read_text().splitlines():
        line = raw_line.strip()
        if not line or line.startswith("#"):
            continue
        baseline.add(line)
    return baseline


def write_baseline(path: Path, findings: Sequence[Finding]) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    lines = [
        "# Baseline for scripts/check-import-aliases.py.",
        "# Every committed entry must be backed by a go/types review showing that",
        "# the default import name is shadowed in a selector's lexical scope, or",
        "# that build constraints exclude the file from the repository type load.",
        "# Do not commit syntax-only findings as blanket exceptions.",
        "",
    ]
    lines.extend(sorted(finding.baseline_key() for finding in findings))
    path.write_text("\n".join(lines) + "\n")


def print_findings(title: str, findings: Sequence[Finding]) -> None:
    if not findings:
        return
    print(title, file=sys.stderr)
    for finding in findings:
        print(
            f"{finding.file}:{finding.line}: alias {finding.alias!r} on {finding.path} "
            f"can use default package name {finding.default_name!r}",
            file=sys.stderr,
        )


def parse_args(argv: Sequence[str]) -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--root", type=Path, default=repo_root())
    parser.add_argument(
        "--baseline",
        type=Path,
        default=None,
        help="line-based allowlist of existing findings; stale entries fail the audit",
    )
    parser.add_argument(
        "--write-baseline",
        type=Path,
        default=None,
        help="write the current findings to this baseline path and exit",
    )
    parser.add_argument(
        "--include-prefix",
        action="append",
        default=None,
        help="local import prefix to audit; may be repeated",
    )
    parser.add_argument(
        "--show-retained",
        action="store_true",
        help="print aliases retained because they resolve a same-file conflict",
    )
    return parser.parse_args(argv)


def main(argv: Sequence[str] | None = None) -> int:
    args = parse_args(argv or sys.argv[1:])
    root = args.root.resolve()
    prefixes = tuple(args.include_prefix or DEFAULT_IMPORT_PREFIXES)
    package_names = load_package_names(root)
    findings, retained, explicit_local_aliases = audit(root, package_names, prefixes)

    if args.write_baseline is not None:
        baseline_path = args.write_baseline
        if not baseline_path.is_absolute():
            baseline_path = root / baseline_path
        write_baseline(baseline_path, findings)
        print(f"wrote import alias baseline: {baseline_path} entries={len(findings)}")
        return 0

    baseline_path = args.baseline
    if baseline_path is not None and not baseline_path.is_absolute():
        baseline_path = root / baseline_path
    baseline = load_baseline(baseline_path)
    finding_by_key = {finding.baseline_key(): finding for finding in findings}
    finding_keys = set(finding_by_key)
    new_keys = sorted(finding_keys - baseline)
    stale_keys = sorted(baseline - finding_keys)
    new_findings = [finding_by_key[key] for key in new_keys]

    if new_findings:
        print_findings("non-conflicting explicit aliases remain:", new_findings)
    if stale_keys:
        print("stale import alias baseline entries:", file=sys.stderr)
        for key in stale_keys:
            print(key, file=sys.stderr)

    print(
        f"import alias audit summary: explicit_local_aliases={explicit_local_aliases} "
        f"findings={len(findings)} baseline_allowed={len(finding_keys & baseline)} "
        f"new={len(new_keys)} stale_baseline={len(stale_keys)} "
        f"retained_conflict={len(retained)}"
    )
    if args.show_retained:
        for item in retained:
            print(
                f"RETAIN {item.file}:{item.line} alias={item.alias} "
                f"default={item.default_name} path={item.path} "
                f"conflicts={','.join(item.conflicts)}"
            )

    return 1 if new_keys or stale_keys else 0


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