#!/usr/bin/env python3
from __future__ import annotations

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


REPO_ROOT = Path(__file__).resolve().parents[2]
FRONTEND_SRC = REPO_ROOT / "frontend" / "src"

FILE_GLOBS = ("**/*.ts", "**/*.tsx", "**/*.vue")
CALL_PATTERNS = {
    "localStorage.getItem": re.compile(r"localStorage\.getItem\(([^)]+)\)"),
    "localStorage.setItem": re.compile(r"localStorage\.setItem\(([^,)]+)"),
    "localStorage.removeItem": re.compile(r"localStorage\.removeItem\(([^)]+)\)"),
    "sessionStorage.getItem": re.compile(r"sessionStorage\.getItem\(([^)]+)\)"),
    "sessionStorage.setItem": re.compile(r"sessionStorage\.setItem\(([^,)]+)"),
    "sessionStorage.removeItem": re.compile(r"sessionStorage\.removeItem\(([^)]+)\)"),
}

BANNED_KEY_MARKERS = ("audit_local_",)


@dataclass
class Finding:
    file: str
    line: int
    call: str
    key_expr: str
    severity: str
    category: str
    note: str


def iter_source_files() -> Iterable[Path]:
    for pattern in FILE_GLOBS:
        for path in FRONTEND_SRC.glob(pattern):
            if path.is_file():
                yield path


def classify(file_path: str, key_expr: str) -> tuple[str, str, str]:
    normalized = key_expr.replace('"', "").replace("'", "").strip()
    lowered = normalized.lower()
    file_lower = file_path.lower()

    if "projectdetail.vue" in file_lower:
        if "publishedprojectsignaturekey" in lowered:
            return ("medium", "publish-helper", "发布比对签名缓存，属于辅助数据")
        if "storagekey" in lowered or "projectdraftstoragekey" in lowered:
            return ("high", "business-draft", "立项配置草稿本地兜底，需显式提示未落库")

    if "audit_local_" in lowered:
        return ("high", "audit-data", "审计日志本地缓存，不能作为正式审计数据源")
    if "ctms_setup_config_draft_" in lowered or "ctms_setup_project_draft_" in lowered:
        return ("high", "business-draft", "立项配置草稿本地兜底，需显式提示未落库")
    if "ctms_setup_published_project_sig_" in lowered:
        return ("medium", "publish-helper", "发布比对签名缓存，属于辅助数据")

    if any(k in lowered for k in ("token_key", "last_login_email_key", "logout_reason_storage_key")):
        return ("low", "auth-session", "认证/会话数据（非业务主数据）")
    if any(k in lowered for k in ("ctms_token", "ctms_last_login_email", "credential", "auth")):
        return ("low", "auth-session", "认证/会话数据（非业务主数据）")
    if any(k in lowered for k in ("ctms_sidebar_collapsed", "study", "site", "role")):
        return ("low", "ui-preference", "UI偏好/上下文缓存")

    if "personaltodo" in file_path.lower():
        return ("low", "feature-cache", "工作台本地待办缓存（非核心主数据）")

    return ("medium", "unknown", "未命中规则，建议人工复核是否为业务关键数据")


def scan_findings() -> list[Finding]:
    findings: list[Finding] = []
    for path in iter_source_files():
        try:
            content = path.read_text(encoding="utf-8")
        except UnicodeDecodeError:
            content = path.read_text(encoding="utf-8", errors="ignore")
        lines = content.splitlines()
        rel = str(path.relative_to(REPO_ROOT))
        for idx, line in enumerate(lines, start=1):
            for call_name, pattern in CALL_PATTERNS.items():
                for match in pattern.finditer(line):
                    key_expr = match.group(1).strip()
                    severity, category, note = classify(rel, key_expr)
                    findings.append(
                        Finding(
                            file=rel,
                            line=idx,
                            call=call_name,
                            key_expr=key_expr,
                            severity=severity,
                            category=category,
                            note=note,
                        )
                    )
    return findings


def markdown_report(findings: list[Finding]) -> str:
    counts: dict[str, int] = {"high": 0, "medium": 0, "low": 0}
    for item in findings:
        counts[item.severity] = counts.get(item.severity, 0) + 1

    lines = [
        "# 存储落库核查报告",
        "",
        f"- 扫描目录: `{FRONTEND_SRC}`",
        f"- 总发现数: `{len(findings)}`",
        f"- 高风险: `{counts['high']}` / 中风险: `{counts['medium']}` / 低风险: `{counts['low']}`",
        "",
        "## 明细",
        "",
        "| 风险 | 分类 | 文件 | 行号 | 调用 | key表达式 | 说明 |",
        "| --- | --- | --- | --- | --- | --- | --- |",
    ]
    for item in sorted(findings, key=lambda x: (x.severity, x.file, x.line)):
        lines.append(
            f"| {item.severity} | {item.category} | `{item.file}` | {item.line} | `{item.call}` | `{item.key_expr}` | {item.note} |"
        )
    return "\n".join(lines)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Scan frontend storage usage and classify persistence risks.")
    parser.add_argument("--format", choices=("markdown", "json"), default="markdown")
    parser.add_argument("--output", default="", help="Optional output file path.")
    parser.add_argument(
        "--fail-on-banned",
        action="store_true",
        help="Exit with code 2 when banned local key markers are detected.",
    )
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    findings = scan_findings()

    banned_hits = []
    if args.fail_on_banned:
        for item in findings:
            expr = item.key_expr.lower().replace('"', "").replace("'", "")
            if any(marker in expr for marker in BANNED_KEY_MARKERS):
                banned_hits.append(item)

    if args.format == "json":
        output_text = json.dumps([asdict(item) for item in findings], ensure_ascii=False, indent=2)
    else:
        output_text = markdown_report(findings)

    if args.output:
        out = Path(args.output)
        if not out.is_absolute():
            out = REPO_ROOT / out
        out.parent.mkdir(parents=True, exist_ok=True)
        out.write_text(output_text + "\n", encoding="utf-8")
    else:
        print(output_text)

    if banned_hits:
        return 2
    return 0


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