"""index1 性能基准测试

用法:
    python -m benchmarks.bench_search            # 完整基准（需要 Ollama）
    python -m benchmarks.bench_search --bm25     # 仅 BM25（无需 Ollama）
"""

from __future__ import annotations

import argparse
import asyncio
import os
import platform
import random
import shutil
import statistics
import sys
import tempfile
import time
from pathlib import Path

# 确保项目 src 在 path 中
sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "src"))

from index1.config import Config
from index1.db import Database
from index1.embed import OllamaEmbedder
from index1.indexer import index_files
from index1.search import search, _cache
from index1.cache import QueryCache

# ── 合成数据生成 ──

_MD_TOPICS = [
    ("Risk Management", "风控模块", ["liquidation", "margin call", "health factor", "collateral"]),
    ("Trading Engine", "交易引擎", ["order book", "matching", "settlement", "fee"]),
    ("User Authentication", "用户认证", ["OAuth", "JWT", "session", "2FA"]),
    ("Payment System", "支付系统", ["refund", "webhook", "reconciliation", "ledger"]),
    ("Data Pipeline", "数据管线", ["ETL", "streaming", "batch processing", "schema"]),
    ("API Gateway", "API 网关", ["rate limiting", "load balancing", "routing", "circuit breaker"]),
    ("Notification Service", "通知服务", ["email", "push", "webhook", "retry"]),
    ("Cache Layer", "缓存层", ["Redis", "invalidation", "TTL", "eviction"]),
    ("Search Engine", "搜索引擎", ["inverted index", "tokenizer", "ranking", "relevance"]),
    ("Monitoring", "监控系统", ["metrics", "alerting", "tracing", "dashboard"]),
]

_PY_CLASSES = [
    "Engine", "Manager", "Handler", "Service", "Controller",
    "Processor", "Validator", "Builder", "Factory", "Repository",
]

_PY_METHODS = [
    "initialize", "process", "validate", "execute", "transform",
    "calculate", "aggregate", "dispatch", "serialize", "configure",
]


def _gen_markdown(topic_idx: int, file_idx: int) -> str:
    title, cn_title, keywords = _MD_TOPICS[topic_idx % len(_MD_TOPICS)]
    sections = []
    sections.append(f"# {title} - Part {file_idx}\n\n{cn_title}模块文档第 {file_idx} 部分。\n")
    for i, kw in enumerate(keywords):
        body = f"This section covers {kw} in detail. "
        body += f"The {kw} subsystem handles critical operations. "
        body += f"{kw.title()} is essential for system reliability. "
        body += f"详细说明 {kw} 的实现逻辑和配置方式。" * 2
        sections.append(f"## {kw.title()}\n\n{body}\n")
        if i < len(keywords) - 1:
            sub = f"### {kw.title()} Configuration\n\n"
            sub += f"Configure {kw} via `config.yaml`. Default settings are production-ready.\n"
            sections.append(sub)
    return "\n".join(sections)


def _gen_python(file_idx: int) -> str:
    cls = _PY_CLASSES[file_idx % len(_PY_CLASSES)]
    lines = [f'"""Module {file_idx}: {cls} implementation"""\n']
    lines.append(f"CONSTANT_{file_idx} = {file_idx * 42}\n\n")
    lines.append(f"class {cls}{file_idx}:")
    lines.append(f'    """{cls} for component {file_idx}"""\n')
    lines.append(f"    def __init__(self):")
    lines.append(f"        self.state = {{}}")
    lines.append(f"        self.counter = 0\n")
    for j, method in enumerate(_PY_METHODS):
        lines.append(f"    def {method}_{j}(self, data):")
        lines.append(f'        """{method} operation"""')
        lines.append(f"        self.counter += 1")
        lines.append(f"        result = {{k: v for k, v in data.items() if v is not None}}")
        lines.append(f"        return result\n")
    lines.append(f"\ndef helper_{file_idx}():")
    lines.append(f'    """Utility function {file_idx}"""')
    lines.append(f"    return {cls}{file_idx}()\n")
    return "\n".join(lines)


def generate_test_data(base_dir: Path, n_md: int = 100, n_py: int = 50) -> int:
    """生成合成测试文件，返回文件总数"""
    docs_dir = base_dir / "docs"
    src_dir = base_dir / "src"
    docs_dir.mkdir(parents=True, exist_ok=True)
    src_dir.mkdir(parents=True, exist_ok=True)

    try:
        for i in range(n_md):
            content = _gen_markdown(i, i)
            (docs_dir / f"doc_{i:03d}.md").write_text(content, encoding="utf-8")

        for i in range(n_py):
            content = _gen_python(i)
            (src_dir / f"module_{i:03d}.py").write_text(content, encoding="utf-8")
    except OSError as e:
        print(f"生成测试数据失败: {e}")
        raise

    return n_md + n_py


# ── 计时工具 ──

def _percentile(data: list[float], pct: int) -> float:
    """计算百分位数"""
    if not data:
        return 0.0
    k = (len(data) - 1) * pct / 100
    f = int(k)
    c = f + 1
    if c >= len(data):
        return data[-1]
    return data[f] + (k - f) * (data[c] - data[f])


def _format_ms(ms: float) -> str:
    if ms < 1:
        return f"{ms:.2f}ms"
    return f"{ms:.1f}ms"


# ── 搜索查询 ──

_QUERIES = [
    "liquidation mechanism",
    "how to configure rate limiting",
    "数据管线 ETL 流程",
    "OAuth authentication flow",
    "Redis cache invalidation",
    "order book matching engine",
    "webhook retry strategy",
    "monitoring alerting rules",
    "payment reconciliation process",
    "search ranking algorithm",
]


# ── 基准测试 ──

async def bench_index(data_dir: Path, config: Config, db: Database, embedder: OllamaEmbedder) -> dict:
    """索引性能测试"""
    start = time.perf_counter()
    stats = await index_files(
        [str(data_dir)], config, db, embedder, force=True, base_dir=data_dir,
    )
    elapsed = time.perf_counter() - start
    return {"elapsed_s": round(elapsed, 2), **stats}


def bench_bm25(db: Database, queries: list[str], runs: int = 20) -> list[float]:
    """BM25 搜索延迟 (ms)"""
    latencies = []
    for _ in range(runs):
        q = random.choice(queries)
        start = time.perf_counter()
        db.fts5_search(q, collection=None, limit=10)
        latencies.append((time.perf_counter() - start) * 1000)
    latencies.sort()
    return latencies


async def bench_vector(
    db: Database, embedder: OllamaEmbedder, queries: list[str], runs: int = 20
) -> list[float]:
    """向量搜索延迟 (ms)"""
    latencies = []
    # 预热 embedding
    embeddings = {}
    for q in queries:
        embeddings[q] = await embedder.embed_one(q)

    for _ in range(runs):
        q = random.choice(queries)
        start = time.perf_counter()
        db.vector_search(embeddings[q], collection=None, limit=10)
        latencies.append((time.perf_counter() - start) * 1000)
    latencies.sort()
    return latencies


async def bench_hybrid(
    config: Config, db: Database, embedder: OllamaEmbedder,
    queries: list[str], runs: int = 20,
) -> list[float]:
    """混合搜索延迟 (ms) — 完整管线"""
    latencies = []
    for _ in range(runs):
        q = random.choice(queries)
        _cache.clear()  # 每次查询前清除，确保不命中
        start = time.perf_counter()
        await search(q, config, db, embedder)
        latencies.append((time.perf_counter() - start) * 1000)
    latencies.sort()
    return latencies


async def bench_cache(
    config: Config, db: Database, embedder: OllamaEmbedder,
    queries: list[str], runs: int = 20,
) -> list[float]:
    """缓存命中延迟 (ms)"""
    _cache.clear()
    # 预热：先搜一遍填充 L1
    for q in queries:
        await search(q, config, db, embedder)

    latencies = []
    for _ in range(runs):
        q = random.choice(queries)
        start = time.perf_counter()
        await search(q, config, db, embedder)
        latencies.append((time.perf_counter() - start) * 1000)
    latencies.sort()
    return latencies


# ── 报告输出 ──

def _get_env_info(config: Config) -> str:
    """收集环境信息"""
    uname = platform.uname()
    py_ver = platform.python_version()
    try:
        import psutil
        ram_gb = round(psutil.virtual_memory().total / (1024**3), 1)
    except ImportError:
        try:
            ram_gb = round(os.sysconf("SC_PAGE_SIZE") * os.sysconf("SC_PHYS_PAGES") / (1024**3), 1)
        except (ValueError, OSError):
            ram_gb = "?"
    return (
        f"环境: {uname.system} {uname.release}, {uname.machine}, {ram_gb}GB RAM, "
        f"Python {py_ver}\n"
        f"模型: {config.embedding_model} ({config.embedding_dim}d)"
    )


def print_report(
    env_info: str,
    file_count: int,
    chunk_count: int,
    index_stats: dict,
    results: dict[str, list[float]],
) -> None:
    """打印格式化报告"""
    print("\nindex1 性能基准测试")
    print("=" * 40)
    print(env_info)
    print(f"数据: {file_count} files, {chunk_count} chunks")

    print(f"\n索引性能")
    print(f"{'操作':<20} {'耗时':>10}")
    print(f"{'-'*20} {'-'*10}")
    print(f"{'索引 ' + str(index_stats['total']) + ' 文件':<20} {index_stats['elapsed_s']:>8.1f}s")
    print(f"  已索引: {index_stats['indexed']}, 跳过: {index_stats['skipped']}, 失败: {index_stats['failed']}")

    print(f"\n搜索性能")
    print(f"{'操作':<20} {'P50':>10} {'P95':>10} {'P99':>10}")
    print(f"{'-'*20} {'-'*10} {'-'*10} {'-'*10}")
    for name, latencies in results.items():
        if latencies:
            p50 = _format_ms(_percentile(latencies, 50))
            p95 = _format_ms(_percentile(latencies, 95))
            p99 = _format_ms(_percentile(latencies, 99))
            print(f"{name:<20} {p50:>10} {p95:>10} {p99:>10}")


# ── 主流程 ──

async def main(bm25_only: bool = False) -> None:
    tmp_dir = Path(tempfile.mkdtemp(prefix="index1_bench_"))
    db_path = tmp_dir / "bench.db"

    config = Config(
        db_path=str(db_path),
        embedding_model="nomic-embed-text",
        embedding_dim=768,
        top_k=10,
        collection="bench",
    )

    try:
        db = Database(db_path, config.embedding_dim)
        db.connect()
    except Exception as e:
        print(f"数据库初始化失败: {e}")
        sys.exit(1)

    embedder = OllamaEmbedder(
        model=config.embedding_model,
        ollama_url=config.ollama_url,
        dim=config.embedding_dim,
    )

    try:
        # 1. 生成测试数据
        print("生成测试数据...")
        file_count = generate_test_data(tmp_dir, n_md=100, n_py=50)
        print(f"  {file_count} 个文件已生成")

        # 2. 索引
        print("索引中...")
        index_stats = await bench_index(tmp_dir, config, db, embedder)
        print(f"  索引完成: {index_stats['elapsed_s']}s")

        stats = db.get_stats()
        chunk_count = stats["chunk_count"]

        # 3. 搜索基准
        results: dict[str, list[float]] = {}
        queries = _QUERIES

        print("BM25 搜索基准...")
        results["BM25 搜索"] = bench_bm25(db, queries)

        ollama_ok = not bm25_only and await embedder.check_availability()

        if ollama_ok:
            print("向量搜索基准...")
            results["向量搜索"] = await bench_vector(db, embedder, queries)

            print("混合搜索基准...")
            results["混合搜索"] = await bench_hybrid(config, db, embedder, queries)

            print("缓存命中基准...")
            results["缓存命中"] = await bench_cache(config, db, embedder, queries)
        else:
            if not bm25_only:
                print("  Ollama 不可用，跳过向量/混合/缓存测试")
            # BM25-only 也测缓存
            print("BM25 缓存命中基准...")
            _cache.clear()
            for q in queries:
                await search(q, config, db, embedder)
            cache_latencies = []
            for _ in range(20):
                q = random.choice(queries)
                start = time.perf_counter()
                await search(q, config, db, embedder)
                cache_latencies.append((time.perf_counter() - start) * 1000)
            cache_latencies.sort()
            results["缓存命中"] = cache_latencies

        # 4. 报告
        env_info = _get_env_info(config)
        print_report(env_info, file_count, chunk_count, index_stats, results)

    finally:
        await embedder.close()
        db.close()
        try:
            shutil.rmtree(tmp_dir)
        except OSError as e:
            print(f"警告: 无法删除临时目录 {tmp_dir}: {e}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="index1 性能基准测试")
    parser.add_argument("--bm25", action="store_true", help="仅测试 BM25（无需 Ollama）")
    args = parser.parse_args()
    asyncio.run(main(bm25_only=args.bm25))
