"""task-121 N+1 扫描审计工具。

对 backend-java modules 下的 Java Service 源码做机械化扫描：
识别循环（for/while/forEach/map）内对 MyBatis-Plus 单行查询
（selectById/selectOne/selectList/selectCount/selectObjs/selectMaps/getById/getOne）
的调用，产出 N+1 候选清单（位置/对象/影响），供人工审计与 task-122 修复。
本工具只读，不改任何代码。
"""

from __future__ import annotations

import argparse
import json
import re
import sys
from dataclasses import dataclass, field
from pathlib import Path

HIGH = "HIGH"
MEDIUM = "MEDIUM"
LOW = "LOW"
IMPACT_ORDER = {HIGH: 0, MEDIUM: 1, LOW: 2}

MODULES_ROOT = (
    Path(__file__).resolve().parent.parent / "src/main/java/com/nanri/aiimage/modules"
)

# 单行查询调用（selectBatchIds/selectByIds 天然批量，不匹配此正则）
QUERY_RE = re.compile(
    r"\.(selectById|selectOne|selectList|selectCount|selectObjs|selectMaps|getById|getOne)\("
)
# 循环头：for/while 关键字 + 圆括号跨度，或 .forEach( / .map( 方法调用
LOOP_RE = re.compile(r"\b(for|while)\s*\(|\.(forEach|map)\s*\(")
# 返回集合的查询调用（循环遍历这类结果集 → 影响 HIGH）
COLLECTION_QUERY_RE = re.compile(
    r"(selectList|selectObjs|selectMaps|selectBatchIds|selectByIds|\.list\()"
)
ASSIGN_RE = re.compile(r"(\w+)\s*=\s*([^;{}=]+);")
FIELD_RE = re.compile(r"\b[A-Za-z_]\w*\b")


def _strip_comments(text: str) -> str:
    """用空格替换注释（保留换行与偏移），简化后续扫描。"""
    out = list(text)
    i, n = 0, len(text)
    while i < n:
        if text.startswith("/*", i):
            j = text.find("*/", i + 2)
            j = n if j < 0 else j + 2
            for k in range(i, j):
                if text[k] != "\n":
                    out[k] = " "
            i = j
        elif text.startswith("//", i):
            j = text.find("\n", i)
            j = n if j < 0 else j
            for k in range(i, j):
                out[k] = " "
            i = j
        else:
            i += 1
    return "".join(out)


def _match_span(text: str, open_pos: int) -> int | None:
    """从 text[open_pos] 的 ( 出发，返回匹配的 ) 偏移（含），不匹配返回 None。"""
    depth = 0
    for i in range(open_pos, len(text)):
        c = text[i]
        if c == "(":
            depth += 1
        elif c == ")":
            depth -= 1
            if depth == 0:
                return i
    return None


def _line_of(text: str, offset: int) -> int:
    return text.count("\n", 0, offset) + 1


class _Frame:
    __slots__ = ("body_depth", "expr", "loop_type", "start")

    def __init__(self, start: int, body_depth: int, loop_type: str, expr: str):
        self.start = start
        self.body_depth = body_depth
        self.loop_type = loop_type
        self.expr = expr


def _query_results(stripped: str) -> set[str]:
    """收集被赋值给集合查询结果（selectList 等）的变量名。"""
    names = set()
    for lhs, rhs in ASSIGN_RE.findall(stripped):
        if COLLECTION_QUERY_RE.search(rhs):
            names.add(lhs)
    return names


def _loop_headers(stripped: str) -> list[dict]:
    """预取全部循环头：(start, end, type, expr)（start 为 ( 偏移，end 为匹配的 ) 偏移）。"""
    headers: list[dict] = []
    pos = 0
    while True:
        m = LOOP_RE.search(stripped, pos)
        if not m:
            break
        paren_pos = m.end() - 1
        if stripped[paren_pos] == "(":
            close = _match_span(stripped, paren_pos)
            if close is not None:
                headers.append(
                    {
                        "start": paren_pos,
                        "end": close,
                        "type": m.group(1) or m.group(2),
                        "expr": stripped[paren_pos + 1 : close].strip(),
                    }
                )
        pos = m.end()
    return headers


def _header_context(stripped: str, h: dict, offset: int) -> tuple[bool, str, str]:
    """查询落在循环头括号跨度内时的归属判定。

    返回 (是否循环内, 循环类型, 循环表达式)：
    - lambda 箭头之后 → 循环内（无花括号 lambda 体）
    - 增强 for 的 ` : ` 之后（迭代源 selectList 等）→ 集合装载，非循环内
    - for/while 初始化或条件 → 视为循环内（条件逐次求值）
    """
    expr = h["expr"]
    if "->" in stripped[h["start"] : h["end"]]:
        arrow = stripped.find("->", h["start"], h["end"])
        if arrow != -1 and offset > arrow:
            return True, h["type"], expr
        return False, h["type"], expr
    if " : " in expr:
        colon = stripped.find(" : ", h["start"], h["end"])
        if colon != -1 and offset > colon:
            return False, h["type"], expr
        return True, h["type"], expr
    return True, h["type"], expr


def scan_text(text: str, filename: str = "fixture.java") -> list[dict]:
    """扫描单份 Java 源码，返回 N+1 候选清单（按影响降序、文件行号升序）。"""
    stripped = _strip_comments(text)
    query_results = _query_results(stripped)
    headers = _loop_headers(stripped)

    depth = 0
    frames: list[_Frame] = []
    single: tuple[str, str, int, int] | None = None  # (loop_type, expr, start, end)
    findings: list[dict] = []

    def in_loop(offset: int) -> bool:
        if single is not None and single[2] <= offset < single[3]:
            return True
        if any(f.start <= offset and depth >= f.body_depth for f in frames):
            return True
        for h in reversed(headers):
            if h["start"] <= offset <= h["end"]:
                return _header_context(stripped, h, offset)[0]
        return False

    def enclosing(offset: int) -> tuple[str, str]:
        if single is not None and single[2] <= offset < single[3]:
            return single[0], single[1]
        for h in reversed(headers):
            if h["start"] <= offset <= h["end"]:
                in_loop, loop_type, expr = _header_context(stripped, h, offset)
                if in_loop:
                    return loop_type, expr
        for f in reversed(frames):
            if f.start <= offset:
                return f.loop_type, f.expr
        return "loop", ""

    i = 0
    n = len(stripped)
    while i < n:
        c = stripped[i]
        if c == "{":
            depth += 1
            i += 1
            continue
        if c == "}":
            depth -= 1
            if frames and depth == frames[-1].body_depth - 1:
                frames.pop()
            i += 1
            continue

        loop_match = LOOP_RE.match(stripped, i)
        if loop_match:
            loop_type = loop_match.group(1) or loop_match.group(2)
            paren_pos = loop_match.end() - 1
            close = (
                _match_span(stripped, paren_pos) if stripped[paren_pos] == "(" else None
            )
            if close is not None:
                expr = stripped[paren_pos + 1 : close].strip()
                body_start = close + 1
                while body_start < n and stripped[body_start].isspace():
                    body_start += 1
                if body_start < n and stripped[body_start] == "{":
                    frames.append(_Frame(body_start + 1, depth + 1, loop_type, expr))
                else:
                    end = stripped.find(";", body_start)
                    single = (loop_type, expr, body_start, n if end < 0 else end)
                # 继续扫描括号跨度内容（无花括号 lambda 体的查询落在其中）
                i = paren_pos + 1
                continue
            i = loop_match.end()
            continue

        query_match = QUERY_RE.match(stripped, i)
        if query_match:
            if in_loop(i):
                line_start = stripped.rfind("\n", 0, i) + 1
                line_end = stripped.find("\n", i)
                line_end = n if line_end < 0 else line_end
                statement = stripped[line_start:line_end].strip()
                loop_type, expr = enclosing(i)
                impact = _classify(loop_type, [expr], query_results)
                findings.append(
                    {
                        "file": filename,
                        "line": _line_of(text, i),
                        "statement": statement,
                        "loop_type": loop_type,
                        "impact": impact,
                    }
                )
            i = query_match.end()
            continue

        i += 1

    findings.sort(key=lambda f: (IMPACT_ORDER[f["impact"]], f["line"]))
    return findings


def _classify(loop_type: str, exprs: list[str], query_results: set[str]) -> str:
    """影响分级：遍历查询结果集 → HIGH；forEach/map lambda → MEDIUM；其余 → LOW。"""
    for expr in exprs:
        if expr and COLLECTION_QUERY_RE.search(expr):
            return HIGH
        if expr and any(t in query_results for t in FIELD_RE.findall(expr)):
            return HIGH
    if loop_type in ("forEach", "map"):
        return MEDIUM
    return LOW


@dataclass
class Report:
    scanned_files: int = 0
    modules: list[str] = field(default_factory=list)
    findings: list[dict] = field(default_factory=list)

    def to_dict(self) -> dict:
        return {
            "scanned_files": self.scanned_files,
            "modules": self.modules,
            "findings": self.findings,
        }


def scan_modules(root: Path) -> Report:
    """扫描 modules 根目录下全部 Java 文件，产出 N+1 候选报告。"""
    report = Report()
    module_set: set[str] = set()
    for java_file in sorted(root.rglob("*.java")):
        report.scanned_files += 1
        rel = java_file.relative_to(root)
        module_set.add(rel.parts[0])
        for finding in scan_text(java_file.read_text(encoding="utf-8"), str(rel)):
            finding["module"] = rel.parts[0]
            report.findings.append(finding)
    report.findings.sort(
        key=lambda f: (IMPACT_ORDER[f["impact"]], f["file"], f["line"])
    )
    report.modules = sorted(module_set)
    return report


def main() -> None:
    parser = argparse.ArgumentParser(description="N+1 扫描审计")
    parser.add_argument("--json", help="输出 JSON 报告路径（默认 stdout）")
    args = parser.parse_args()

    report = scan_modules(MODULES_ROOT)
    payload = report.to_dict()
    if args.json:
        Path(args.json).write_text(
            json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8"
        )
    else:
        print(json.dumps(payload, ensure_ascii=False, indent=2))
    print(
        f"scanned {report.scanned_files} files, "
        f"{len(report.findings)} N+1 candidates "
        f"across {len(report.modules)} modules",
        file=sys.stderr,
    )


if __name__ == "__main__":
    main()