#!/usr/bin/env python3
"""Heuristic complexity hotspot scanner for mixed-language repositories.

Original: https://github.com/Kappaemme-git/codex-complexity-optimizer
"""

from __future__ import annotations

import argparse
import ast
import json
import os
import re
from dataclasses import asdict, dataclass
from pathlib import Path
from typing import Iterable


DEFAULT_EXCLUDES = {
    ".git",
    ".hg",
    ".svn",
    "node_modules",
    "vendor",
    "dist",
    "build",
    ".next",
    ".nuxt",
    "coverage",
    "__pycache__",
    ".venv",
    "venv",
    "target",
    ".turbo",
}

TEXT_EXTENSIONS = {
    ".py",
    ".js",
    ".jsx",
    ".ts",
    ".tsx",
    ".mjs",
    ".cjs",
    ".java",
    ".go",
    ".rb",
    ".php",
    ".cs",
    ".cpp",
    ".cc",
    ".c",
    ".h",
    ".hpp",
    ".swift",
    ".rs",
    ".kt",
    ".kts",
    ".scala",
    ".lua",
    ".zig",
    ".ex",
    ".exs",
    ".erl",
    ".hrl",
    ".dart",
    ".r",
    ".R",
    ".jl",
    ".ml",
    ".mli",
    ".clj",
    ".cljs",
    ".el",
    ".vim",
}

LOOP_RE = re.compile(
    r"\b(for|while|forEach|map|filter|reduce|some|every|find|findIndex|each|collect|select|reject|flat_map|flat_map!)\b"
)
MEMBERSHIP_RE = re.compile(
    r"(\.includes\s*\(|\.indexOf\s*\(|\.find\s*\(|\.findIndex\s*\(|\bin_array\s*\(|\bcontains\s*\(|\.include\?\s*\(|\.contains\s*\(|\.has\s*\()"
)
SORT_RE = re.compile(r"(\.sort\s*\(|\bsorted\s*\(|\bsort\s*\(|\.sort_by\s*[\({]|\.order\s*\()")
# Strong signals match on the method name alone; generic verbs (get, post, find, ...)
# only count when called on a client-looking receiver, so Map.get()/Array.find() stay quiet.
QUERY_IN_LOOP_RE = re.compile(
    r"(\bfetch\s*\(|\baxios\.|\brequest\s*\(|\bquery\s*\(|\bexecute\s*\(|\.findMany\s*\(|\.findOne\s*\(|\.findUnique\s*\(|\.find_by\s*\("
    r"|\b(?:db|database|client|session|conn|connection|api|http|repo|repository|knex|prisma|supabase|requests)\.\w+\s*\()",
    re.IGNORECASE,
)
# Component names need at least one lowercase letter (PascalCase) so UPPER_CASE constants
# don't open a component zone; const-form additionally requires a function-looking right side.
RENDER_HINT_RE = re.compile(
    r"\bfunction\s+[A-Z]\w*[a-z]\w*\s*\("
    r"|\bconst\s+[A-Z]\w*[a-z]\w*\s*=\s*(?:async\s+)?(?:\(|function\b|React\.|memo\s*\(|forwardRef\s*\(|styled[.(]|[A-Za-z_$][\w$]*\s*=>)"
)


@dataclass
class Finding:
    path: str
    line: int
    severity: str
    kind: str
    message: str
    suggestion: str
    confidence: str = "low"


def iter_files(root: Path, excludes: set[str]) -> Iterable[Path]:
    for dirpath, dirnames, filenames in os.walk(root):
        dirnames[:] = [d for d in dirnames if d not in excludes]
        for filename in filenames:
            path = Path(dirpath) / filename
            if path.suffix in TEXT_EXTENSIONS:
                yield path


def read_text(path: Path) -> str | None:
    try:
        return path.read_text(encoding="utf-8")
    except UnicodeDecodeError:
        try:
            return path.read_text(encoding="latin-1")
        except Exception:
            return None
    except Exception:
        return None


def rel(path: Path, root: Path) -> str:
    try:
        return str(path.relative_to(root))
    except ValueError:
        return str(path)


# Method names that are query-like on their own vs. only when called on a client-looking object.
STRONG_QUERY_NAMES = {"fetch", "query", "execute", "find_one", "find_many", "find_unique", "find_by", "findone", "findmany", "findunique"}
GENERIC_QUERY_NAMES = {"find", "select", "where", "get", "post", "put", "delete", "request", "all", "first", "last"}
CLIENT_RECEIVER_HINTS = {
    "db",
    "database",
    "session",
    "client",
    "conn",
    "connection",
    "cursor",
    "requests",
    "http",
    "api",
    "repo",
    "repository",
    "collection",
    "engine",
    "orm",
    "supabase",
    "prisma",
}

# Constructors whose results support O(1) membership checks.
CONSTANT_TIME_CONSTRUCTORS = {"set", "frozenset", "dict", "Counter", "defaultdict"}


class PythonVisitor(ast.NodeVisitor):
    def __init__(self, path: Path, root: Path) -> None:
        self.path = path
        self.root = root
        self.loop_depth = 0
        self.findings: list[Finding] = []
        self.constant_time_names: set[str] = set()

    def add(self, node: ast.AST, severity: str, kind: str, message: str, suggestion: str) -> None:
        self.findings.append(
            Finding(
                rel(self.path, self.root),
                getattr(node, "lineno", 1),
                severity,
                kind,
                message,
                suggestion,
                confidence="high",
            )
        )

    def visit_For(self, node: ast.For) -> None:
        self._visit_loop(node)

    def visit_While(self, node: ast.While) -> None:
        self._visit_loop(node)

    def visit_ListComp(self, node: ast.ListComp) -> None:
        self._visit_loop(node)

    def visit_SetComp(self, node: ast.SetComp) -> None:
        self._visit_loop(node)

    def visit_DictComp(self, node: ast.DictComp) -> None:
        self._visit_loop(node)

    def visit_GeneratorExp(self, node: ast.GeneratorExp) -> None:
        self._visit_loop(node)

    def _visit_loop(self, node: ast.AST) -> None:
        if self.loop_depth >= 1:
            self.add(
                node,
                "high",
                "nested-loop",
                "Nested loop may create O(n^2) or worse behavior.",
                "Check whether a map/set index, sort+two-pointer pass, grouping, or batching can replace the inner scan.",
            )
        self.loop_depth += 1
        self.generic_visit(node)
        self.loop_depth -= 1

    def visit_Assign(self, node: ast.Assign) -> None:
        self._track_assignment(node.targets, node.value)
        self.generic_visit(node)

    def visit_AnnAssign(self, node: ast.AnnAssign) -> None:
        if node.value is not None:
            self._track_assignment([node.target], node.value)
        self.generic_visit(node)

    def _track_assignment(self, targets: list[ast.expr], value: ast.expr) -> None:
        constant_time = isinstance(value, (ast.Set, ast.Dict, ast.SetComp, ast.DictComp)) or (
            isinstance(value, ast.Call) and call_name(value.func) in CONSTANT_TIME_CONSTRUCTORS
        )
        for target in targets:
            if isinstance(target, ast.Name):
                if constant_time:
                    self.constant_time_names.add(target.id)
                else:
                    self.constant_time_names.discard(target.id)

    def _constant_time_lookup(self, node: ast.AST) -> bool:
        if isinstance(node, (ast.Set, ast.Dict, ast.SetComp, ast.DictComp)):
            return True
        if isinstance(node, ast.Call) and call_name(node.func) in CONSTANT_TIME_CONSTRUCTORS:
            return True
        if isinstance(node, ast.Name) and node.id in self.constant_time_names:
            return True
        if isinstance(node, ast.Constant) and isinstance(node.value, str) and len(node.value) <= 32:
            return True
        if isinstance(node, (ast.Tuple, ast.List)) and len(node.elts) <= 10:
            return True
        return False

    def visit_Compare(self, node: ast.Compare) -> None:
        if self.loop_depth:
            for op, comparator in zip(node.ops, node.comparators):
                if isinstance(op, (ast.In, ast.NotIn)) and not self._constant_time_lookup(comparator):
                    self.add(
                        node,
                        "medium",
                        "membership-in-loop",
                        "Membership check inside a loop can become O(n*m) when the right side is a list or computed sequence.",
                        "If semantics allow it, build a set or dict once before the loop.",
                    )
                    break
        self.generic_visit(node)

    def visit_Call(self, node: ast.Call) -> None:
        name = call_name(node.func)
        if self.loop_depth and name in {"sorted", "sort"}:
            self.add(
                node,
                "high",
                "sort-in-loop",
                "Sorting inside a loop is often avoidable repeated O(n log n) work.",
                "Sort once outside the loop, maintain a heap, or use binary search/insertion if intermediate ordering is required.",
            )
        if self.loop_depth and name in {"filter", "map"}:
            self.add(
                node,
                "medium",
                "repeated-scan",
                f"{name}() inside a loop may repeatedly scan a collection.",
                "Consider precomputing an index/grouping or combining passes.",
            )
        if self.loop_depth and self._is_query_call(node, name.lower()):
            self.add(
                node,
                "high",
                "io-or-query-in-loop",
                "Potential database/API/file operation inside a loop.",
                "Look for N+1 behavior; batch or preload while preserving auth, filters, ordering, and error handling.",
            )
        self.generic_visit(node)

    def _is_query_call(self, node: ast.Call, name: str) -> bool:
        if name in STRONG_QUERY_NAMES:
            return True
        if name in GENERIC_QUERY_NAMES:
            return receiver_name(node.func) in CLIENT_RECEIVER_HINTS
        return False


def call_name(func: ast.AST) -> str:
    if isinstance(func, ast.Name):
        return func.id
    if isinstance(func, ast.Attribute):
        return func.attr
    return ""


def receiver_name(func: ast.AST) -> str:
    """Innermost attribute receiver: `self.db.get(...)` -> "db", `requests.get(...)` -> "requests"."""
    if isinstance(func, ast.Attribute):
        value = func.value
        if isinstance(value, ast.Name):
            return value.id.lower()
        if isinstance(value, ast.Attribute):
            return value.attr.lower()
    return ""


def scan_python(path: Path, root: Path, text: str) -> list[Finding]:
    try:
        tree = ast.parse(text)
    except SyntaxError as exc:
        return [
            Finding(
                rel(path, root),
                exc.lineno or 1,
                "info",
                "parse-error",
                "Python file could not be parsed; falling back to textual scanning only.",
                "Inspect manually if this file is on a hot path.",
            )
        ] + scan_text(path, root, text)
    visitor = PythonVisitor(path, root)
    visitor.visit(tree)
    return visitor.findings


def scan_text(path: Path, root: Path, text: str) -> list[Finding]:
    findings: list[Finding] = []
    lines = text.splitlines()
    loop_stack: list[tuple[int, int]] = []
    function_component_ranges = component_ranges(lines) if path.suffix in {".jsx", ".tsx", ".js", ".ts"} else set()

    for idx, line in enumerate(lines, start=1):
        stripped = line.strip()
        if not stripped or stripped.startswith(("//", "#", "*")):
            continue
        indent = len(line) - len(line.lstrip(" \t"))
        # A loop is still open only while the current line is indented deeper than it;
        # a line at the same or shallower indent closes it (otherwise sequential loops look nested).
        loop_stack = [(level, lno) for level, lno in loop_stack if level < indent]

        if LOOP_RE.search(stripped):
            if loop_stack:
                findings.append(
                    Finding(
                        rel(path, root),
                        idx,
                        "high",
                        "nested-or-callback-loop",
                        "Loop or array iteration appears inside another loop/callback.",
                        "Check whether indexing, grouping, batching, or a single-pass algorithm can remove repeated scans.",
                    )
                )
            loop_stack.append((indent, idx))

        if loop_stack and MEMBERSHIP_RE.search(stripped):
            findings.append(
                Finding(
                    rel(path, root),
                    idx,
                    "medium",
                    "membership-in-loop",
                    "Membership/search operation appears inside iterative code.",
                    "Consider a Set/Map or precomputed lookup if equality and ordering semantics allow it.",
                )
            )

        if loop_stack and SORT_RE.search(stripped):
            findings.append(
                Finding(
                    rel(path, root),
                    idx,
                    "high",
                    "sort-in-loop",
                    "Sort appears inside iterative code.",
                    "Move sorting out of the loop or use a heap/binary-search strategy if intermediate order is needed.",
                )
            )

        if loop_stack and QUERY_IN_LOOP_RE.search(stripped):
            findings.append(
                Finding(
                    rel(path, root),
                    idx,
                    "high",
                    "io-or-query-in-loop",
                    "Potential database/API/file operation inside a loop.",
                    "Look for N+1 behavior; batch or preload while preserving auth, filters, ordering, and error handling.",
                )
            )

        if idx in function_component_ranges and any(
            token in stripped for token in [".filter(", ".map(", ".sort(", ".reduce("]
        ):
            findings.append(
                Finding(
                    rel(path, root),
                    idx,
                    "medium",
                    "render-derived-work",
                    "Collection transform appears in a likely UI component render path.",
                    "For large collections, consider memoized selectors, server-side derivation, or virtualization.",
                )
            )

    return findings


def component_ranges(lines: list[str]) -> set[int]:
    active_until = 0
    interesting: set[int] = set()
    brace_balance = 0
    in_component = False

    for idx, line in enumerate(lines, start=1):
        if RENDER_HINT_RE.search(line):
            in_component = True
            active_until = idx + 120
            brace_balance = 0
        if in_component:
            interesting.add(idx)
            brace_balance += line.count("{") - line.count("}")
            # Close the zone as soon as braces balance out, so code below a small
            # component doesn't inherit its render-path findings; cap as a fallback.
            if idx > active_until or (brace_balance <= 0 and "}" in line):
                in_component = False
    return interesting


def dedupe(findings: list[Finding]) -> list[Finding]:
    seen: set[tuple[str, int, str]] = set()
    result: list[Finding] = []
    for finding in findings:
        key = (finding.path, finding.line, finding.kind)
        if key not in seen:
            seen.add(key)
            result.append(finding)
    return result


def severity_rank(finding: Finding) -> tuple[int, str, int]:
    order = {"high": 0, "medium": 1, "info": 2}
    return (order.get(finding.severity, 3), finding.path, finding.line)


def render_markdown(findings: list[Finding]) -> str:
    if not findings:
        return "No obvious complexity hotspots found by heuristic scanning.\n"
    lines = ["# Complexity Hotspots", ""]
    for finding in findings:
        lines.extend(
            [
                f"## {finding.severity.upper()} {finding.kind}",
                f"- Location: `{finding.path}:{finding.line}`",
                f"- Confidence: {finding.confidence}",
                f"- Finding: {finding.message}",
                f"- Suggestion: {finding.suggestion}",
                "",
            ]
        )
    return "\n".join(lines)


def main() -> int:
    parser = argparse.ArgumentParser(description="Scan a repository for likely complexity hotspots.")
    parser.add_argument("root", nargs="?", default=".", help="Repository or directory to scan.")
    parser.add_argument("--format", choices=["markdown", "json"], default="markdown")
    parser.add_argument("--exclude", action="append", default=[], help="Additional directory name to exclude.")
    parser.add_argument("--max-findings", type=int, default=80)
    args = parser.parse_args()

    root = Path(args.root).resolve()
    excludes = DEFAULT_EXCLUDES | set(args.exclude)
    findings: list[Finding] = []

    for path in iter_files(root, excludes):
        text = read_text(path)
        if text is None:
            continue
        if path.suffix == ".py":
            findings.extend(scan_python(path, root, text))
        else:
            findings.extend(scan_text(path, root, text))

    findings = sorted(dedupe(findings), key=severity_rank)
    total = len(findings)
    findings = findings[: args.max_findings]
    if args.format == "json":
        print(json.dumps([asdict(f) for f in findings], indent=2))
    else:
        output = render_markdown(findings)
        if total > len(findings):
            output += f"\nShowing {len(findings)} of {total} findings. Raise --max-findings to see the rest.\n"
        print(output)
    return 0


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