server.py

"""
TestGenius Flask Web Server
提供 HTTP API，让浏览器前端能够调用 TestGenius 的静态分析和测试生成功能。
"""
import os
import sys
import time
import tempfile
from pathlib import Path

# 加载 .env 文件（如果存在）
try:
    from dotenv import load_dotenv
    load_dotenv()
except ImportError:
    pass

from flask import Flask, request, jsonify, send_from_directory
from flask_cors import CORS

# 将项目根目录加入 sys.path
PROJECT_ROOT = Path(__file__).parent
sys.path.insert(0, str(PROJECT_ROOT))

from concurrent.futures import ThreadPoolExecutor, as_completed

from core.ast_parser import PythonASTParser
from core.cpp_ast_parser import CppASTParser
from core.utils import (
    build_function_id,
    build_function_qualified_name,
    build_function_signature,
    normalize_file_path,
)
from agents.orchestrator import OrchestratorAgent
from agents.cpp_generator_agent import ALL_EXTENSIONS as CPP_EXTENSIONS

# 所有支持的源文件扩展名
SUPPORTED_EXTENSIONS = {".py"} | CPP_EXTENSIONS

# 文件浏览/读取的根目录约束（环境变量可覆盖，默认为项目根目录）
BROWSE_ROOT = Path(os.environ.get("TESTGENIUS_BROWSE_ROOT", str(PROJECT_ROOT))).resolve()

app = Flask(__name__, static_folder=str(PROJECT_ROOT))

# CORS: 允许 localhost 和 file:// 协议（origin=null）访问
_debug_mode = os.environ.get("FLASK_DEBUG", "0") == "1"
if _debug_mode:
    CORS(app)
else:
    CORS(app, origins=["http://localhost:*", "http://127.0.0.1:*", "null"])


# ─── 安全工具函数 ─────────────────────────────────────────────────────────────

def _is_within_browse_root(target: Path) -> bool:
    """检查路径是否在允许的浏览根目录内"""
    try:
        target.resolve().relative_to(BROWSE_ROOT)
        return True
    except ValueError:
        return False


def _function_info_to_dict(fn, file_path: str = "") -> dict:
    """将 FunctionInfo dataclass 转换为可序列化的 dict"""
    resolved_file_path = normalize_file_path(
        file_path or getattr(fn, "file_path", "") or ""
    )
    return {
        "name": fn.name,
        "class_name": fn.class_name,
        "namespace": getattr(fn, "namespace", None),
        "params": fn.params,
        "return_type": fn.return_type,
        "docstring": fn.docstring,
        "decorators": fn.decorators,
        "start_line": fn.start_line,
        "end_line": fn.end_line,
        "template_params": getattr(fn, "template_params", []) or [],
        "access_specifier": getattr(fn, "access_specifier", None),
        "file_path": resolved_file_path,
        "signature": build_function_signature(fn),
        "qualified_name": build_function_qualified_name(fn),
        "function_id": build_function_id(fn, resolved_file_path),
    }


def _read_test_file(test_file_path: str) -> str:
    """读取生成的测试文件内容"""
    try:
        return Path(test_file_path).read_text(encoding="utf-8")
    except Exception:
        return ""


def _normalize_name_filter(raw_filter) -> list[str]:
    """兼容 string/list 形式的 function_filter。"""
    if isinstance(raw_filter, str):
        raw_filter = [raw_filter]
    if not isinstance(raw_filter, list):
        return []
    normalized = []
    for item in raw_filter:
        if not isinstance(item, str):
            continue
        s = item.strip()
        if s:
            normalized.append(s)
    return normalized


def _normalize_target_functions(raw_targets) -> list[dict]:
    """兼容 string/list[dict] 的 target_functions。"""
    if isinstance(raw_targets, (str, dict)):
        raw_targets = [raw_targets]
    if not isinstance(raw_targets, list):
        return []

    normalized = []
    for item in raw_targets:
        if isinstance(item, str):
            value = item.strip()
            if not value:
                continue
            if value.startswith("fn_"):
                normalized.append({"function_id": value})
            else:
                normalized.append({"name": value})
            continue
        if not isinstance(item, dict):
            continue

        target = {}
        for key in (
            "function_id", "name", "class_name", "namespace", "qualified_name",
            "signature", "file_path", "start_line", "end_line",
        ):
            if item.get(key) not in (None, ""):
                target[key] = item.get(key)
        if "file_path" in target:
            target["file_path"] = normalize_file_path(target["file_path"])
        if target:
            normalized.append(target)

    return normalized


def _select_parser_for_filename(filename: str):
    suffix = Path(filename or "").suffix.lower()
    if suffix in CPP_EXTENSIONS:
        return CppASTParser()
    return PythonASTParser()


# ─── 路由 ────────────────────────────────────────────────────────────────────

@app.route("/")
def index():
    """返回前端页面"""
    return send_from_directory(str(PROJECT_ROOT), "TestGenius.html")


@app.route("/api/health")
def health():
    """服务健康检查"""
    api_key_set = bool(os.environ.get("ANTHROPIC_API_KEY"))
    return jsonify({
        "status": "ok",
        "api_key_configured": api_key_set,
    })


@app.route("/api/analyze", methods=["POST"])
def analyze():
    """dry-run 分析：仅解析代码，返回函数列表，不调用 Claude。"""
    data = request.get_json(force=True, silent=True) or {}
    code = data.get("code", "").strip()
    filename = data.get("filename", "sample.py") or "sample.py"

    if not code:
        return jsonify({"success": False, "error": "代码不能为空"}), 400

    try:
        parser = _select_parser_for_filename(filename)
        if Path(filename).suffix.lower() in CPP_EXTENSIONS:
            functions, imports = parser.parse_source(code, filename)
            for fn in functions:
                fn.file_path = normalize_file_path(filename)
        else:
            functions, imports = parser.parse_source(code)
        return jsonify({
            "success": True,
            "functions": [_function_info_to_dict(fn, filename) for fn in functions],
            "imports": imports,
        })
    except Exception as e:
        return jsonify({"success": False, "error": str(e)}), 500


@app.route("/api/generate", methods=["POST"])
def generate():
    """完整流程：解析 → 生成 → 验证，返回生成的测试代码和指标。"""
    data = request.get_json(force=True, silent=True) or {}
    code = data.get("code", "").strip()
    filename = data.get("filename", "sample.py") or "sample.py"
    file_path_param = data.get("file_path", "").strip() if data.get("file_path") else ""
    workers = int(data.get("workers", 5))
    workers = max(1, min(10, workers))
    run_pytest = bool(data.get("run_pytest", False))
    include_paths = data.get("include_paths", []) or []
    function_filter = _normalize_name_filter(data.get("function_filter", []) or [])
    target_functions = _normalize_target_functions(data.get("target_functions", []) or [])

    # 如果提供了 file_path，直接从磁盘读取
    use_disk_file = False
    if file_path_param:
        disk_path = Path(file_path_param).resolve()
        if not _is_within_browse_root(disk_path):
            return jsonify({"success": False, "error": "访问路径超出允许范围"}), 403
        if not disk_path.exists() or not disk_path.is_file():
            return jsonify({"success": False, "error": f"文件不存在: {disk_path}"}), 404
        use_disk_file = True
        filename = disk_path.name
        try:
            code = disk_path.read_text(encoding="utf-8")
        except Exception as e:
            return jsonify({"success": False, "error": f"读取文件失败: {e}"}), 500
    elif not code:
        return jsonify({"success": False, "error": "代码不能为空"}), 400

    filename = Path(filename).name
    suffix = Path(filename).suffix.lower()
    if not suffix:
        filename += ".py"
        suffix = ".py"
    if suffix not in SUPPORTED_EXTENSIONS:
        filename = Path(filename).stem + ".py"
        suffix = ".py"

    is_cpp = suffix in CPP_EXTENSIONS

    TESTS_DIR = PROJECT_ROOT / "tests"
    TESTS_DIR.mkdir(exist_ok=True)

    def _run_orchestrator(src_path: str, selected_filter=None, selected_targets=None):
        orchestrator = OrchestratorAgent(
            output_dir=str(TESTS_DIR),
            max_workers=workers,
            run_pytest=run_pytest,
            include_paths=include_paths,
        )
        orchestrator.client.reset_usage()
        t_start = time.time()
        result = orchestrator.run(
            src_path,
            function_filter=selected_filter or None,
            target_functions=selected_targets or None,
        )
        elapsed = round(time.time() - t_start, 2)
        token_usage = orchestrator.client.get_usage()
        return orchestrator, result, elapsed, token_usage

    def _run_with_fallback(src_path: str):
        orchestrator, result, elapsed, token_usage = _run_orchestrator(
            src_path,
            function_filter,
            target_functions,
        )
        retried_with_name_filter = False
        if (
            result.get("status") == "skipped"
            and target_functions
            and function_filter
        ):
            fallback_orchestrator, fallback_result, fallback_elapsed, fallback_usage = _run_orchestrator(
                src_path,
                function_filter,
                None,
            )
            if fallback_result.get("status") != "skipped":
                orchestrator = fallback_orchestrator
                result = fallback_result
                elapsed = fallback_elapsed
                token_usage = fallback_usage
                retried_with_name_filter = True
        return orchestrator, result, elapsed, token_usage, retried_with_name_filter

    try:
        source_path_for_meta = normalize_file_path(file_path_param or str(PROJECT_ROOT / filename))
        if use_disk_file:
            source_path = str(disk_path)
            source_path_for_meta = normalize_file_path(str(disk_path))
            orchestrator, result, elapsed, token_usage, retried_with_name_filter = _run_with_fallback(
                source_path,
            )
        else:
            with tempfile.TemporaryDirectory(prefix="testgenius_") as tmpdir:
                tmpdir_path = Path(tmpdir)
                src_file = tmpdir_path / filename
                src_file.write_text(code, encoding="utf-8")
                source_path = str(src_file)
                orchestrator, result, elapsed, token_usage, retried_with_name_filter = _run_with_fallback(
                    source_path,
                )

        test_code = ""
        if result.get("test_file"):
            test_code = _read_test_file(result["test_file"])

        # 解析函数总数（用于覆盖率计算）
        functions_total = 0
        functions = []
        if not is_cpp:
            try:
                parser = PythonASTParser()
                fns, _ = parser.parse_source(code)
                functions = [_function_info_to_dict(fn, source_path_for_meta) for fn in fns]
                functions_total = len(fns)
            except Exception:
                pass
        else:
            try:
                parser = CppASTParser()
                fns, _ = parser.parse_source(code, source_path_for_meta)
                for fn in fns:
                    fn.file_path = source_path_for_meta
                functions = [_function_info_to_dict(fn, source_path_for_meta) for fn in fns]
                functions_total = len(fns)
            except Exception:
                functions_total = result.get("functions_processed", 0)

        details = result.get("details", [])
        passed = sum(1 for d in details if d.get("pytest_passed") is True)
        skipped = sum(1 for d in details if d.get("pytest_passed") is None)
        failed = len(details) - passed - skipped
        functions_tested = result.get("functions_processed", len(details))

        # 覆盖率
        func_coverage = round(functions_tested / functions_total * 100, 1) if functions_total > 0 else 0

        status = result.get("status")
        success = status not in {"error", "skipped"}
        error_message = None
        if not success:
            error_message = result.get("reason") or "未生成测试代码"
        elif retried_with_name_filter:
            error_message = None

        return jsonify({
            "success": success,
            "error": error_message,
            "test_code": test_code,
            "functions": functions,
            "lang": "cpp" if is_cpp else "python",
            "stats": {
                "total": len(details),
                "passed": passed,
                "failed": failed,
                "skipped": skipped,
            },
            "details": details,
            "metrics": {
                "total_time_s": elapsed,
                "input_tokens": token_usage["input_tokens"],
                "output_tokens": token_usage["output_tokens"],
                "cache_creation_input_tokens": token_usage["cache_creation_input_tokens"],
                "cache_read_input_tokens": token_usage["cache_read_input_tokens"],
                "total_tokens": token_usage["total_tokens"],
                "functions_total": functions_total,
                "functions_tested": functions_tested,
                "function_coverage_pct": func_coverage,
                "model": orchestrator.client.model,
            },
        })

    except Exception as e:
        return jsonify({"success": False, "error": str(e)}), 500


# 搜索时跳过的目录名
_SKIP_DIRS = {"build", ".git", "__pycache__", "test", "tests", "doc", "docs",
              "debug", "Debug", "Release", "release", ".svn", "node_modules",
              "CMakeFiles", ".cache"}

# 搜索时使用的解析器（模块级单例，避免重复初始化 tree-sitter）
_search_py_parser = PythonASTParser()
_search_cpp_parser = CppASTParser()


def _collect_source_files(directory: Path, max_files: int = 500) -> list[Path]:
    """递归收集目录下所有支持的源文件，跳过无关目录，最多返回 max_files 个。"""
    result = []
    stack = [directory]
    while stack and len(result) < max_files:
        current = stack.pop()
        try:
            entries = sorted(current.iterdir(), key=lambda x: x.name.lower())
        except (PermissionError, OSError):
            continue
        dirs = []
        for entry in entries:
            if len(result) >= max_files:
                break
            if entry.is_dir():
                if entry.name not in _SKIP_DIRS and not entry.name.startswith("."):
                    dirs.append(entry)
            elif entry.is_file() and entry.suffix.lower() in SUPPORTED_EXTENSIONS:
                result.append(entry)
        stack.extend(reversed(dirs))
    return result


def _parse_file_for_search(file_path: Path):
    """解析单个文件提取函数列表，出错返回空列表。"""
    try:
        suffix = file_path.suffix.lower()
        if suffix == ".py":
            functions, _ = _search_py_parser.parse_file(str(file_path))
        else:
            functions, _ = _search_cpp_parser.parse_file(str(file_path))
        return file_path, functions
    except Exception:
        return file_path, []


@app.route("/api/search-functions", methods=["POST"])
def search_functions():
    """在指定目录中搜索函数名，支持模糊匹配。"""
    data = request.get_json(force=True, silent=True) or {}
    directory = data.get("directory", "").strip()
    query = data.get("query", "").strip()
    max_results = int(data.get("max_results", 50))
    max_results = max(1, min(200, max_results))

    if not directory:
        return jsonify({"success": False, "error": "directory 参数不能为空"}), 400
    if not query:
        return jsonify({"success": False, "error": "query 参数不能为空"}), 400

    dir_path = Path(directory).resolve()
    if not _is_within_browse_root(dir_path):
        return jsonify({"success": False, "error": "访问路径超出允许范围"}), 403
    if not dir_path.exists() or not dir_path.is_dir():
        return jsonify({"success": False, "error": f"目录不存在: {dir_path}"}), 400

    try:
        source_files = _collect_source_files(dir_path)
        files_scanned = len(source_files)

        # 多线程解析
        all_matches = []
        query_lower = query.lower()
        with ThreadPoolExecutor(max_workers=8) as executor:
            futures = [executor.submit(_parse_file_for_search, f) for f in source_files]
            for future in as_completed(futures):
                file_path, functions = future.result()
                file_abs_path = normalize_file_path(str(file_path))
                language = "python" if file_path.suffix.lower() == ".py" else "cpp"
                for fn in functions:
                    if query_lower in fn.name.lower():
                        try:
                            rel = str(file_path.relative_to(dir_path))
                        except ValueError:
                            rel = str(file_path)
                        item = _function_info_to_dict(fn, file_abs_path)
                        item["relative_path"] = rel
                        item["language"] = language
                        all_matches.append(item)

        total_matches = len(all_matches)
        # 按文件路径 + 行号排序，截断到 max_results
        all_matches.sort(key=lambda m: (m["relative_path"], m["start_line"]))
        truncated = all_matches[:max_results]

        return jsonify({
            "success": True,
            "functions": truncated,
            "files_scanned": files_scanned,
            "total_matches": total_matches,
        })

    except Exception as e:
        return jsonify({"success": False, "error": str(e)}), 500


@app.route("/api/browse")
def browse():
    """浏览目录内容，限制在 BROWSE_ROOT 内。"""
    raw = request.args.get("path", "").strip()
    target = Path(raw).resolve() if raw else BROWSE_ROOT

    if not _is_within_browse_root(target):
        return jsonify({"success": False, "error": "访问路径超出允许范围"}), 403

    if not target.exists() or not target.is_dir():
        return jsonify({"success": False, "error": f"路径不存在: {target}"}), 400

    try:
        items = []
        for item in sorted(target.iterdir(), key=lambda x: (x.is_file(), x.name.lower())):
            if item.name.startswith(".") or item.name == "__pycache__":
                continue
            entry = {
                "name": item.name,
                "path": str(item),
                "is_dir": item.is_dir(),
                "is_py": item.is_file() and item.suffix == ".py",
                "is_cpp": item.is_file() and item.suffix.lower() in CPP_EXTENSIONS,
            }
            if item.is_file():
                entry["size"] = item.stat().st_size
            items.append(entry)

        parent = str(target.parent) if target != target.parent and _is_within_browse_root(target.parent) else None

        return jsonify({
            "success": True,
            "path": str(target),
            "parent": parent,
            "items": items,
        })
    except PermissionError:
        return jsonify({"success": False, "error": "权限不足，无法访问该目录"}), 403
    except Exception as e:
        return jsonify({"success": False, "error": str(e)}), 500


@app.route("/api/read-file")
def read_file():
    """读取指定源文件内容，限制在 BROWSE_ROOT 内。"""
    raw = request.args.get("path", "").strip()
    if not raw:
        return jsonify({"success": False, "error": "path 参数不能为空"}), 400

    target = Path(raw).resolve()

    if not _is_within_browse_root(target):
        return jsonify({"success": False, "error": "访问路径超出允许范围"}), 403

    if not target.exists() or not target.is_file():
        return jsonify({"success": False, "error": "文件不存在"}), 404
    if target.suffix.lower() not in SUPPORTED_EXTENSIONS:
        return jsonify({"success": False, "error": f"不支持的文件类型: {target.suffix}"}), 400

    try:
        content = target.read_text(encoding="utf-8")
        return jsonify({
            "success": True,
            "content": content,
            "filename": target.name,
            "path": str(target),
        })
    except PermissionError:
        return jsonify({"success": False, "error": "权限不足，无法读取该文件"}), 403
    except Exception as e:
        return jsonify({"success": False, "error": str(e)}), 500


# ─── 启动 ────────────────────────────────────────────────────────────────────

if __name__ == "__main__":
    port = int(os.environ.get("PORT", 5000))
    debug = os.environ.get("FLASK_DEBUG", "0") == "1"
    print(f"[TestGenius] 服务启动: http://localhost:{port}")
    print(f"[TestGenius] API Key 已配置: {bool(os.environ.get('ANTHROPIC_API_KEY'))}")
    print(f"[TestGenius] 文件浏览根目录: {BROWSE_ROOT}")
    app.run(host="0.0.0.0", port=port, debug=debug)