#!/usr/bin/env python3
"""赛题生成工具：generate + calibrate。

用法:
    conda activate Video-Tree-TRM
    python tools/generate_questions.py generate --store-dir store ...
    python tools/generate_questions.py calibrate ...

app/core/adapters 不 import 此脚本。
"""

from __future__ import annotations

import argparse
import asyncio
import json
import os
import random
import sys
from pathlib import Path

PROJECT_ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(PROJECT_ROOT))

from dotenv import load_dotenv
from loguru import logger

load_dotenv(PROJECT_ROOT / ".env")

import numpy as np

from app.question_gen.loader import load_benchmark
from app.question_gen.synthesizer import (
    TASK_TYPE_LEVEL_MAP,
    generate_one,
    is_duplicate,
)
from core.types import GeneratedQuestion  # noqa: TCH001 — runtime use in _append_to_json

# ---------------------------------------------------------------------------
# 日志配置：不缓存，立即输出
# ---------------------------------------------------------------------------

logger.remove()
logger.add(
    sys.stderr,
    format="{time:HH:mm:ss} | {level:<7} | {message}",
    level="DEBUG",
    colorize=True,
)
logger.add(
    PROJECT_ROOT / "logs" / "generate_questions.log",
    format="{time:YYYY-MM-DD HH:mm:ss} | {level:<7} | {message}",
    level="DEBUG",
    rotation="50 MB",
)


# ---------------------------------------------------------------------------
# 断点续跑 — progress 文件管理
# ---------------------------------------------------------------------------


def _load_or_init_progress(output_dir: Path) -> dict:
    """加载 progress.json，不存在则返回初始结构。

    参数:
        output_dir: 输出目录路径。

    返回:
        {"completed": {task_type: [question_id, ...]}, "output_dir": str}。
        文件损坏时返回初始结构并记录警告。
    """
    progress_path = output_dir / "progress.json"
    if progress_path.exists():
        try:
            data = json.loads(progress_path.read_text(encoding="utf-8"))
            if not isinstance(data.get("completed"), dict):
                raise ValueError("completed 字段不是 dict")
            return data
        except (json.JSONDecodeError, ValueError, KeyError, TypeError) as exc:
            logger.warning("progress.json 损坏，重新初始化: {}", exc)
    return {"completed": {}, "output_dir": str(output_dir)}


def _save_progress(output_dir: Path, progress: dict) -> None:
    """原子写入 progress.json。

    参数:
        output_dir: 输出目录路径。
        progress: 进度数据。
    """
    tmp = output_dir / "progress.json.tmp"
    tmp.write_text(
        json.dumps(progress, ensure_ascii=False, indent=2),
        encoding="utf-8",
    )
    os.replace(str(tmp), str(output_dir / "progress.json"))


# ---------------------------------------------------------------------------
# Embedding 池重建（断点续跑时从已生成 JSON 重建）
# ---------------------------------------------------------------------------


def _rebuild_embedding_pool(
    output_dir: Path,
    embed_fn,
    benchmark_questions: list[GeneratedQuestion],
) -> dict[str, np.ndarray]:
    """从已生成 JSON + benchmark 题目重建每个题型的 embedding 池。

    断点续跑时调用，确保去重池包含所有已有题目。

    参数:
        output_dir: 包含 {video_id}.json 的输出目录。
        embed_fn: 文本嵌入函数（str | list[str] → [N, D] ndarray）。
        benchmark_questions: benchmark 题目列表（也要加入去重池）。

    返回:
        {task_type: [N, D] ndarray}，空题型的 ndarray 为 shape (0,)。
    """
    pools: dict[str, list[str]] = {}

    # Phase 1: 收集 benchmark 题目文本
    for q in benchmark_questions:
        pools.setdefault(q.task_type, []).append(q.question)

    # Phase 2: 收集已生成题目文本
    for json_path in sorted(output_dir.glob("*.json")):
        if json_path.name == "progress.json":
            continue
        try:
            items = json.loads(json_path.read_text(encoding="utf-8"))
        except (json.JSONDecodeError, OSError) as exc:
            logger.warning("跳过损坏文件 {}: {}", json_path, exc)
            continue
        if not isinstance(items, list):
            continue
        for item in items:
            task_type = item.get("task_type", "")
            question = item.get("question", "")
            if task_type and question:
                pools.setdefault(task_type, []).append(question)

    # Phase 3: 批量嵌入
    result: dict[str, np.ndarray] = {}
    for task_type, texts in pools.items():
        if texts:
            result[task_type] = embed_fn(texts)
        else:
            result[task_type] = np.empty(0)

    # Phase 4: 确保所有 12 题型都有条目
    for task_type in TASK_TYPE_LEVEL_MAP:
        if task_type not in result:
            result[task_type] = np.empty(0)

    logger.info(
        "embedding 池重建完成: {}",
        {k: v.shape[0] if v.ndim == 2 else 0 for k, v in result.items()},
    )
    return result


# ---------------------------------------------------------------------------
# Exemplar 选取
# ---------------------------------------------------------------------------


def _select_exemplars(
    benchmark: list[GeneratedQuestion],
    task_type: str,
    n: int,
    rng: random.Random,
) -> list[GeneratedQuestion]:
    """从 benchmark 中选取同题型示例，优先跨视频多样性。

    参数:
        benchmark: benchmark 题目全集。
        task_type: 目标题型。
        n: 期望选取数量。
        rng: 可控随机数生成器。

    返回:
        min(n, 可用数) 个示例，尽量来自不同 video_id。
    """
    # Phase 1: 过滤同题型
    candidates = [q for q in benchmark if q.task_type == task_type]
    if not candidates:
        return []

    take = min(n, len(candidates))

    # Phase 2: 按 video_id 分桶，轮询取样保证跨视频多样性
    by_video: dict[str, list[GeneratedQuestion]] = {}
    for q in candidates:
        by_video.setdefault(q.video_id, []).append(q)

    # 每桶内部打乱
    for bucket in by_video.values():
        rng.shuffle(bucket)

    # 轮询选取
    video_ids = list(by_video.keys())
    rng.shuffle(video_ids)
    selected: list[GeneratedQuestion] = []
    idx = 0
    while len(selected) < take:
        vid = video_ids[idx % len(video_ids)]
        bucket = by_video[vid]
        if bucket:
            selected.append(bucket.pop(0))
        else:
            # 桶空了，从 video_ids 中移除
            video_ids.remove(vid)
            if not video_ids:
                break
            # 不递增 idx，因为移除后当前位置是下一个
            continue
        idx += 1

    return selected


# ---------------------------------------------------------------------------
# 客户端构建（从 .env）
# ---------------------------------------------------------------------------


def _build_vlm_client():
    """构建 GovernedVLMClient，复用 repair_trees.py 的模式。

    从 .env 读取 VL_LLM_MODEL / VL_LLM_BASE_URL / VL_LLM_API_KEY
    和 LLM 韧性参数，构造治理栈。

    返回:
        GovernedVLMClient 实例。
    """
    from adapters.breaker import CircuitBreaker
    from adapters.llm import GovernedLLMClient
    from adapters.telemetry import SQLiteTelemetryRecorder
    from adapters.vlm import GovernedVLMClient

    (PROJECT_ROOT / "logs").mkdir(exist_ok=True)
    telemetry = SQLiteTelemetryRecorder(
        str(PROJECT_ROOT / "logs" / "generate_questions_telemetry.db")
    )

    breaker_threshold = int(os.getenv("LLM_CIRCUIT_BREAKER_THRESHOLD", "5"))
    breaker_cooldown = int(os.getenv("LLM_CIRCUIT_BREAKER_COOLDOWN", "60"))
    timeout_s = float(os.getenv("LLM_TIMEOUT", "120"))
    max_retries = int(os.getenv("LLM_MAX_RETRIES", "3"))
    base_delay = float(os.getenv("LLM_RETRY_BASE_DELAY", "2.0"))
    max_delay = float(os.getenv("LLM_RETRY_MAX_DELAY", "30.0"))
    ttft = float(os.getenv("LLM_TTFT_TIMEOUT", "30"))
    inter_token = float(os.getenv("LLM_INTER_TOKEN_TIMEOUT", "15"))

    vlm_base = GovernedLLMClient(
        model=os.environ["VL_LLM_MODEL"],
        base_url=os.environ["VL_LLM_BASE_URL"],
        api_key=os.environ["VL_LLM_API_KEY"],
        provider="qwen",
        thinking=False,
        breaker=CircuitBreaker(fail_threshold=breaker_threshold, cooldown_s=breaker_cooldown),
        cache=None,
        telemetry=telemetry,
        timeout_s=timeout_s,
        ttft_timeout_s=ttft,
        inter_token_timeout_s=inter_token,
        max_retries=max_retries,
        retry_base_delay_s=base_delay,
        retry_max_delay_s=max_delay,
    )
    return GovernedVLMClient(vlm_base)


def _build_embed_provider():
    """构建 EmbeddingProvider，从 .env 决定 local 或 remote。

    环境变量:
        EMBED_API_KEY + EMBED_API_URL 都非空 → RemoteEmbeddingProvider
        否则 → LocalEmbeddingProvider

    模型名称和维度通过 EMBED_MODEL / EMBED_DIM 环境变量配置。

    返回:
        LocalEmbeddingProvider 或 RemoteEmbeddingProvider 实例。
    """
    from adapters.embedding import LocalEmbeddingProvider, RemoteEmbeddingProvider

    model_name = os.environ.get("EMBED_MODEL", "BAAI/bge-base-zh-v1.5")
    embed_dim = int(os.environ.get("EMBED_DIM", "768"))
    api_key = os.environ.get("EMBED_API_KEY", "")
    api_url = os.environ.get("EMBED_API_URL", "")

    if api_key and api_url:
        logger.info("使用远程嵌入: model={}, url={}", model_name, api_url)
        return RemoteEmbeddingProvider(
            model_name=model_name,
            embed_dim=embed_dim,
            api_key=api_key,
            api_url=api_url,
        )

    logger.info("使用本地嵌入: model={}, dim={}", model_name, embed_dim)
    device = os.environ.get("EMBED_DEVICE", "cpu")
    return LocalEmbeddingProvider(
        model_name=model_name,
        embed_dim=embed_dim,
        device=device,
    )


# ---------------------------------------------------------------------------
# JSON 追加写入
# ---------------------------------------------------------------------------


def _append_to_json(output_dir: Path, question: GeneratedQuestion) -> None:
    """将生成的题目追加到对应 video_id 的 JSON 文件。

    文件格式：[{...}, {...}, ...]，每个 video_id 一个文件。

    参数:
        output_dir: 输出目录。
        question: 待写入的题目。
    """
    json_path = output_dir / f"{question.video_id}.json"
    existing: list[dict] = []
    if json_path.exists():
        try:
            existing = json.loads(json_path.read_text(encoding="utf-8"))
        except (json.JSONDecodeError, OSError):
            logger.warning("读取 {} 失败，覆盖写入", json_path)
            existing = []

    entry = {
        "question_id": question.question_id,
        "task_type": question.task_type,
        "question": question.question,
        "options": list(question.options),
        "answer": question.answer,
        "source_nodes": list(question.source_nodes),
        "difficulty": question.difficulty,
    }
    existing.append(entry)

    # 原子写入
    tmp = json_path.with_suffix(".json.tmp")
    tmp.write_text(
        json.dumps(existing, ensure_ascii=False, indent=2),
        encoding="utf-8",
    )
    os.replace(str(tmp), str(json_path))


# ---------------------------------------------------------------------------
# generate 主流程
# ---------------------------------------------------------------------------


async def _run_generate(args: argparse.Namespace) -> None:
    """generate 子命令主流程。

    按题型顺序生成题目，每个 slot 串行生成并去重，
    断点续跑通过 progress.json 跳过已完成 slot。

    参数:
        args: CLI 参数（store_dir, output_dir, per_type, similarity_threshold,
              max_retries, concurrency, seed）。
    """
    store_dir = Path(args.store_dir)
    output_dir = Path(args.output_dir)
    per_type = args.per_type
    similarity_threshold = args.similarity_threshold
    max_retries = args.max_retries
    concurrency = args.concurrency
    seed = args.seed

    output_dir.mkdir(parents=True, exist_ok=True)

    # Phase 1: 加载视频列表
    videos_dir = store_dir / "videos"
    if not videos_dir.exists():
        logger.error("视频目录不存在: {}", videos_dir)
        sys.exit(1)

    video_ids = sorted(
        d.name for d in videos_dir.iterdir() if d.is_dir() and (d / "tree.json").exists()
    )
    if not video_ids:
        logger.error("未找到任何有 tree.json 的视频目录")
        sys.exit(1)
    logger.info("发现 {} 个视频", len(video_ids))

    # Phase 2: 加载 benchmark 题目（用于 exemplars + 去重池初始化）
    benchmark_dir = store_dir / "questions" / "benchmarks" / "Video-MME"
    benchmark: list[GeneratedQuestion] = []
    if benchmark_dir.exists():
        benchmark = load_benchmark(benchmark_dir)
        logger.info("加载 {} 道 benchmark 题目", len(benchmark))
    else:
        logger.warning("benchmark 目录不存在: {}", benchmark_dir)

    # Phase 3: 构建客户端
    vlm = _build_vlm_client()
    embed_provider = _build_embed_provider()
    embed_fn = embed_provider.embed

    # Phase 4: 初始化或恢复 progress
    progress = _load_or_init_progress(output_dir)

    # Phase 5: 重建 embedding 池
    pools = _rebuild_embedding_pool(output_dir, embed_fn, benchmark)

    # Phase 6: 加载视频树索引（延迟按需加载）
    from app.tree.index import TreeIndex

    rng = random.Random(seed)
    sem = asyncio.Semaphore(concurrency)
    task_types = list(TASK_TYPE_LEVEL_MAP.keys())
    total_generated = 0
    total_failed = 0

    async def _generate_with_sem(
        vlm_client,
        embed_fn_inner,
        tree,
        video_id,
        task_type,
        seq,
        *,
        exemplars,
        used_node_ids,
        max_retries_inner,
        similarity_threshold_inner,
        rng_inner,
        session_id,
    ):
        """Semaphore 包装的 generate_one 调用。"""
        async with sem:
            return await generate_one(
                vlm_client,
                embed_fn_inner,
                tree,
                video_id,
                task_type,
                seq,
                exemplars=exemplars,
                used_node_ids=used_node_ids,
                max_retries=max_retries_inner,
                similarity_threshold=similarity_threshold_inner,
                rng=rng_inner,
                session_id=session_id,
            )

    # Phase 7: 逐题型、逐 slot 生成
    for task_type in task_types:
        completed_ids = set(progress["completed"].get(task_type, []))
        start_seq = len(completed_ids)

        if start_seq >= per_type:
            logger.info("题型 {} 已完成 {}/{}", task_type, start_seq, per_type)
            continue

        logger.info(
            "题型 {} 开始生成: 已完成 {}, 目标 {}",
            task_type,
            start_seq,
            per_type,
        )

        # 选取 exemplars
        exemplars = _select_exemplars(benchmark, task_type, 3, rng)

        for seq in range(start_seq, per_type):
            # 随机选一个视频
            video_id = rng.choice(video_ids)
            tree_path = videos_dir / video_id / "tree.json"

            try:
                tree = TreeIndex.load_json(str(tree_path))
            except Exception as exc:
                logger.warning("加载树 {} 失败: {}", tree_path, exc)
                total_failed += 1
                continue

            used_node_ids: set[str] = set()
            session_id = f"gen-{task_type}-{seq}"
            generated = False

            for _attempt in range(max_retries):
                candidate = await _generate_with_sem(
                    vlm,
                    embed_fn,
                    tree,
                    video_id,
                    task_type,
                    seq,
                    exemplars=exemplars,
                    used_node_ids=used_node_ids,
                    max_retries_inner=1,
                    similarity_threshold_inner=similarity_threshold,
                    rng_inner=rng,
                    session_id=session_id,
                )

                if candidate is None:
                    continue

                # 去重检查（单线程原子操作）
                pool = pools.get(task_type, np.empty(0))
                if (
                    pool.ndim == 2
                    and pool.shape[0] > 0
                    and is_duplicate(candidate.question, pool, embed_fn, similarity_threshold)
                ):
                    logger.warning("去重: {} 与池中题目相似", candidate.question_id)
                    continue

                # 原子操作：更新池 + 写 JSON + 更新 progress
                new_emb = embed_fn(candidate.question)  # [1, D]
                if pool.ndim == 2 and pool.shape[0] > 0:
                    pools[task_type] = np.vstack([pool, new_emb])
                else:
                    pools[task_type] = new_emb

                _append_to_json(output_dir, candidate)
                progress["completed"].setdefault(task_type, []).append(candidate.question_id)
                _save_progress(output_dir, progress)

                total_generated += 1
                generated = True
                logger.debug(
                    "生成: {} (题型={}, 序号={})",
                    candidate.question_id,
                    task_type,
                    seq,
                )
                break

            if not generated:
                logger.error("题型 {} seq {} 耗尽 {} 次重试", task_type, seq, max_retries)
                total_failed += 1

    # Phase 8: 汇总
    logger.info("=" * 60)
    logger.info("生成完成: 成功 {}, 失败 {}", total_generated, total_failed)
    logger.info("=" * 60)

    if total_failed > 0:
        logger.error("{} 个 slot 生成失败", total_failed)
        sys.exit(1)

    # 全部完成，删除 progress.json
    progress_path = output_dir / "progress.json"
    if progress_path.exists():
        progress_path.unlink()
        logger.info("已删除 progress.json（全部完成）")


# ---------------------------------------------------------------------------
# CLI 解析
# ---------------------------------------------------------------------------


def _parse_args() -> argparse.Namespace:
    """解析命令行参数。"""
    parser = argparse.ArgumentParser(description="赛题生成工具：generate + calibrate")
    subparsers = parser.add_subparsers(dest="command", required=True)

    # generate 子命令
    gen_parser = subparsers.add_parser("generate", help="生成新题目")
    gen_parser.add_argument(
        "--store-dir",
        type=str,
        required=True,
        help="store 根目录（包含 videos/ 和 questions/）",
    )
    gen_parser.add_argument(
        "--output-dir",
        type=str,
        required=True,
        help="输出目录（生成的 JSON 写入此处）",
    )
    gen_parser.add_argument(
        "--per-type",
        type=int,
        required=True,
        help="每种题型生成数量",
    )
    gen_parser.add_argument(
        "--similarity-threshold",
        type=float,
        required=True,
        help="embedding 去重阈值（余弦相似度）",
    )
    gen_parser.add_argument(
        "--max-retries",
        type=int,
        required=True,
        help="每个 slot 最大重试次数",
    )
    gen_parser.add_argument(
        "--concurrency",
        type=int,
        required=True,
        help="VLM 调用并发数（Semaphore 容量）",
    )
    gen_parser.add_argument(
        "--seed",
        type=int,
        required=True,
        help="随机种子",
    )

    # calibrate 子命令（占位，后续任务实现）
    subparsers.add_parser("calibrate", help="校准题目难度（待实现）")

    return parser.parse_args()


def main() -> None:
    """同步入口。"""
    args = _parse_args()
    (PROJECT_ROOT / "logs").mkdir(exist_ok=True)

    if args.command == "generate":
        asyncio.run(_run_generate(args))
    elif args.command == "calibrate":
        logger.error("calibrate 子命令尚未实现")
        sys.exit(1)


if __name__ == "__main__":
    main()