# -*- coding: utf-8 -*-
"""
Loop task execution engine with pause/resume support.
"""
import asyncio
import sys
from datetime import datetime
from typing import Optional

# Fix Windows GBK encoding issue
sys.stdout.reconfigure(encoding='utf-8', errors='replace')
sys.stderr.reconfigure(encoding='utf-8', errors='replace')

from models.db import get_db, _id, _now
from service.loop_recall_md import DEFAULT_LLM_NOTE, append_recall_md_section


# Module-level control dictionary for pause/resume/stop
# key=loop_task_id, value={"pause_event": asyncio.Event, "stop": bool}
_loop_controls: dict[str, dict] = {}


async def _check_pause(loop_task_id: str) -> bool:
    """Check if task should pause. Returns True if stopped."""
    ctrl = _loop_controls.get(loop_task_id)
    if not ctrl:
        return False

    if ctrl["stop"]:
        return True

    # Wait for pause_event (will block if event is cleared)
    await ctrl["pause_event"].wait()
    return ctrl["stop"]


def _init_control(loop_task_id: str) -> None:
    """Initialize control structure for a loop task."""
    event = asyncio.Event()
    event.set()  # Initially not paused
    _loop_controls[loop_task_id] = {
        "pause_event": event,
        "stop": False,
    }


def _clear_control(loop_task_id: str) -> None:
    """Clean up control structure."""
    _loop_controls.pop(loop_task_id, None)


async def pause_loop(loop_task_id: str) -> bool:
    """Pause a running loop task."""
    ctrl = _loop_controls.get(loop_task_id)
    if not ctrl:
        return False

    # 立即写数据库，让前端看到"已暂停"状态
    async with get_db() as db:
        await db.execute(
            "UPDATE loop_task SET status='paused', paused_at=? WHERE id=?",
            (_now(), loop_task_id),
        )
        await db.commit()

    # Clear event，后台会在阶段边界停下来
    ctrl["pause_event"].clear()
    return True


async def resume_loop(loop_task_id: str) -> bool:
    """Resume a paused loop task."""
    ctrl = _loop_controls.get(loop_task_id)
    if not ctrl:
        return False

    ctrl["pause_event"].set()
    return True


async def stop_loop(loop_task_id: str) -> bool:
    """Stop a loop task permanently."""
    ctrl = _loop_controls.get(loop_task_id)
    if not ctrl:
        return False

    ctrl["stop"] = True
    ctrl["pause_event"].set()  # Unblock if paused

    async with get_db() as db:
        await db.execute(
            "UPDATE loop_task SET status='stopped', finished_at=? WHERE id=?",
            (_now(), loop_task_id),
        )
        await db.commit()

    return True


async def run_loop_task(
    loop_task_id: str,
    org_id: str,
    file_ids: list[str],
    judge_config_id: str,
    questions_per_section: int,
    quality_threshold: float,
    include_multimodal: bool,
    env_url: str,
    d_user_id: str,
    agent_id: str,
    top_k: int,
    recall_top_k: int,
    concurrency: int,
    cross_chunk: bool,
    max_rounds: int,
    max_questions: int,
    global_dedup: bool = False,  # 是否使用全局去重（跨任务）
):
    """
    Main loop execution engine.

    Each round:
    1. Fetch existing questions from all previous rounds
    2. Generate new questions (avoiding existing angles)
    3. Deduplicate with LLM
    4. Create single-jump test
    5. Wait for test completion
    6. Update stats and check termination conditions
    """
    _init_control(loop_task_id)

    try:
        await _do_run_loop(
            loop_task_id, org_id, file_ids, judge_config_id,
            questions_per_section, quality_threshold, include_multimodal,
            env_url, d_user_id, agent_id, top_k, recall_top_k, concurrency, cross_chunk,
            max_rounds, max_questions, global_dedup
        )
    except Exception as e:
        # Mark as failed
        async with get_db() as db:
            await db.execute(
                "UPDATE loop_task SET status='failed', error_message=? WHERE id=?",
                (str(e), loop_task_id),
            )
            await db.commit()
    finally:
        _clear_control(loop_task_id)


async def _do_run_loop(
    loop_task_id: str,
    org_id: str,
    file_ids: list[str],
    judge_config_id: str,
    questions_per_section: int,
    quality_threshold: float,
    include_multimodal: bool,
    env_url: str,
    d_user_id: str,
    agent_id: str,
    top_k: int,
    recall_top_k: int,
    concurrency: int,
    cross_chunk: bool,
    max_rounds: int,
    max_questions: int,
    global_dedup: bool = False,
):
    """Internal loop implementation."""

    # Get loop task name与批次期望切片数（与 chunk_batches_plan.chunk_count 对齐，用于拉取完整性校验）
    async with get_db() as db:
        task_rows = await db.execute_fetchall(
            "SELECT name, expected_chunk_count FROM loop_task WHERE id=?", (loop_task_id,)
        )
    _tr = dict(task_rows[0]) if task_rows else {}
    loop_task_name = _tr.get("name") or loop_task_id[:8]
    _ecc = _tr.get("expected_chunk_count")
    try:
        expected_chunk_count = int(_ecc) if _ecc is not None and int(_ecc) > 0 else None
    except (TypeError, ValueError):
        expected_chunk_count = None

    # Get judge config for LLM client
    async with get_db() as db:
        cfg_rows = await db.execute_fetchall(
            "SELECT * FROM judge_config WHERE id=?", (judge_config_id,)
        )
    if not cfg_rows:
        raise ValueError("judge_config not found")
    judge_cfg = dict(cfg_rows[0])

    # Initialize Embedding client for dedup (向量相似度查重，不再使用 LLM)
    from openai import AsyncOpenAI
    embed_base = (judge_cfg.get("embed_base_url") or judge_cfg["base_url"]).rstrip("/")
    embed_key = judge_cfg.get("embed_api_key") or judge_cfg["api_key"]
    embed_client = AsyncOpenAI(
        base_url=embed_base,
        api_key=embed_key,
    )
    embed_model = judge_cfg.get("embed_model") or "text-embedding-3-small"

    # Update status to running
    async with get_db() as db:
        await db.execute(
            "UPDATE loop_task SET status='running' WHERE id=?",
            (loop_task_id,),
        )
        await db.commit()

    consecutive_empty_rounds = 0

    def stop_check():
        ctrl = _loop_controls.get(loop_task_id)
        if ctrl is None or ctrl.get("stop", False):
            return True
        return False

    async def async_pause_check():
        """Check if paused and wait for resume. Returns True if should stop."""
        ctrl = _loop_controls.get(loop_task_id)
        if not ctrl:
            return False
        if ctrl.get("stop", False):
            return True
        # Check pause and wait if needed
        if not ctrl["pause_event"].is_set():
            await ctrl["pause_event"].wait()
            if ctrl.get("stop", False):
                return True
        return False

    async def check_pause_between_stages() -> bool:
        """在阶段边界等待暂停信号，返回 True 表示应该停止。"""
        ctrl = _loop_controls.get(loop_task_id)
        if not ctrl:
            return False
        if ctrl["stop"]:
            return True
        # 如果 pause_event 已被 clear，说明用户点了暂停
        # pause_loop 已经写了数据库，这里只需要等待 resume
        if not ctrl["pause_event"].is_set():
            await ctrl["pause_event"].wait()  # 阻塞直到 resume
            if ctrl["stop"]:
                return True
            # resume 后把状态改回 running
            async with get_db() as db:
                await db.execute(
                    "UPDATE loop_task SET status='running', paused_at=NULL WHERE id=?",
                    (loop_task_id,),
                )
                await db.commit()
        return False

    # 确定从哪一轮、哪个阶段开始
    # 查最后一轮的状态，决定是继续该轮还是开新轮
    async with get_db() as db:
        rows = await db.execute_fetchall(
            """SELECT id, round_number, status, qa_gen_task_id, single_jump_task_id
               FROM loop_round
               WHERE loop_task_id=?
               ORDER BY round_number DESC LIMIT 1""",
            (loop_task_id,),
        )

    # resume_round: 需要继续执行的轮次信息，None 表示从新轮开始
    resume_round = None
    if rows:
        last = dict(rows[0])
        if last["status"] != "done":
            resume_round = last  # 需要从这一轮的某个阶段继续
            round_number = last["round_number"] - 1  # 循环会 +1 回到这一轮
        else:
            round_number = last["round_number"]  # 从下一轮开始
    else:
        round_number = 0

    while True:
        # 阶段边界：检查暂停/停止
        if await check_pause_between_stages():
            return

        round_number += 1

        # Check max_rounds
        if max_rounds > 0 and round_number > max_rounds:
            break

        # Check max_questions
        if max_questions > 0:
            async with get_db() as db:
                row = await db.execute_fetchall(
                    "SELECT total_approved FROM loop_task WHERE id=?", (loop_task_id,)
                )
                current_total = row[0]["total_approved"] if row else 0
            if current_total >= max_questions:
                break

        # 判断是继续上次中断的轮次，还是创建新轮次
        if resume_round and resume_round["round_number"] == round_number:
            # 继续上次中断的轮次，复用已有的 round_id 和 qa_gen_task_id
            round_id = resume_round["id"]
            resume_stage = resume_round["status"]  # qa_generating / deduplicating / testing
            qa_task_id = resume_round["qa_gen_task_id"]
            resume_round = None  # 只用一次
        else:
            # 创建新轮次
            resume_stage = None
            round_id = _id()
            qa_task_id = None
            async with get_db() as db:
                await db.execute(
                    """INSERT INTO loop_round
                       (id, loop_task_id, round_number, status, started_at)
                       VALUES (?,?,?,?,?)""",
                    (round_id, loop_task_id, round_number, "qa_generating", _now()),
                )
                await db.execute(
                    "UPDATE loop_task SET current_round=? WHERE id=?",
                    (round_number, loop_task_id),
                )
                await db.commit()

        # 1. Get existing questions from all previous rounds
        section_existing_questions = await _get_existing_questions(loop_task_id, global_dedup=global_dedup)
        all_existing_questions = []
        for questions in section_existing_questions.values():
            all_existing_questions.extend(questions)

        # For QA generation, only pass question text (not ids)
        section_existing_text = {
            sp: [q["question"] for q in qs]
            for sp, qs in section_existing_questions.items()
        }

        # 2. QA 生成阶段
        # 如果是从 deduplicating 或 testing 阶段 resume，跳过 QA 生成
        if resume_stage in ("deduplicating", "testing"):
            # qa_task_id 已经有了，直接跳过生成
            pass
        else:
            # 需要运行 QA 生成（新轮次，或从 qa_generating 阶段 resume）
            if qa_task_id is None:
                qa_task_id = _id()
                async with get_db() as db:
                    await db.execute(
                        """INSERT INTO qa_gen_task
                           (id,name,status,judge_config_id,questions_per_section,quality_threshold,
                            progress,total,created_at)
                           VALUES (?,?,?,?,?,?,?,?,?)""",
                        (qa_task_id, f"{loop_task_name}-问题生成-第{round_number}轮", "pending",
                         judge_config_id, questions_per_section, quality_threshold,
                         0, 0, _now()),
                    )
                    await db.execute(
                        "UPDATE loop_round SET qa_gen_task_id=?, status='qa_generating' WHERE id=?",
                        (qa_task_id, round_id),
                    )
                    await db.commit()
            else:
                # resume_stage == 'qa_generating'：qa_task 已存在但未完成，重新跑
                async with get_db() as db:
                    await db.execute(
                        "UPDATE loop_round SET status='qa_generating' WHERE id=?",
                        (round_id,),
                    )
                    await db.commit()

            from api.qa_gen_dagent import _run_dagent_task
            try:
                await _run_dagent_task(
                    task_id=qa_task_id,
                    org_id=org_id,
                    file_id_list=file_ids,
                    judge_config_id=judge_config_id,
                    questions_per_section=questions_per_section,
                    quality_threshold=quality_threshold,
                    include_multimodal=include_multimodal,
                    section_existing_questions=section_existing_text,
                    stop_check=stop_check,
                    pause_check=async_pause_check,
                    env_url=env_url,
                    expected_chunk_count=expected_chunk_count,
                )
            except Exception as e:
                async with get_db() as db:
                    await db.execute(
                        "UPDATE loop_round SET status='failed', finished_at=? WHERE id=?",
                        (_now(), round_id),
                    )
                    await db.commit()
                raise

        # 阶段边界：QA 生成完成后检查暂停
        if await check_pause_between_stages():
            return

        # 3. 去重阶段
        if resume_stage != "testing":
            async with get_db() as db:
                await db.execute(
                    "UPDATE loop_round SET status='deduplicating' WHERE id=?",
                    (round_id,),
                )
                await db.commit()

            # 按切片分组获取新问题
            new_questions_by_chunk = await _get_new_questions_by_chunk(qa_task_id)

            # 按切片分组获取已有问题（用于查重），排除本轮 qa_task_id 避免自查自
            existing_by_chunk = await _get_existing_questions_by_chunk(
                loop_task_id,
                exclude_qa_task_id=qa_task_id,
                global_dedup=global_dedup,
            )

            if new_questions_by_chunk:
                from service.dedup import deduplicate_questions_by_chunk

                async def on_dedup_progress(done: int, total: int):
                    async with get_db() as db:
                        await db.execute(
                            "UPDATE loop_round SET dedup_progress=? WHERE id=?",
                            (f"{done}/{total}", round_id),
                        )
                        await db.commit()

                # 按切片并行查重（正则归一化 + 向量余弦相似度）
                dup_results = await deduplicate_questions_by_chunk(
                    new_questions_by_chunk,
                    existing_by_chunk,
                    embed_client,
                    embed_model,
                    similarity_threshold=0.85,
                    max_parallel_chunks=5,
                    stop_check=stop_check,
                    pause_check=async_pause_check,
                    on_progress=on_dedup_progress,
                )

                if stop_check():
                    return

                async with get_db() as db:
                    for qid, (dup_of, sim) in dup_results.items():
                        if dup_of:
                            await db.execute(
                                """UPDATE qa_gen_question
                                   SET dup_of=?, dup_similarity=?, status='rejected'
                                   WHERE id=?""",
                                (dup_of, sim, qid),
                            )
                    await db.commit()

        # 阶段边界：去重完成后检查暂停
        if await check_pause_between_stages():
            return

        # 统计本轮数据
        async with get_db() as db:
            counts = await db.execute_fetchall(
                """SELECT
                    COUNT(*) as generated,
                    SUM(CASE WHEN status='approved' THEN 1 ELSE 0 END) as approved,
                    SUM(CASE WHEN dup_of IS NOT NULL THEN 1 ELSE 0 END) as duplicates
                FROM qa_gen_question WHERE task_id=?""",
                (qa_task_id,),
            )
            gen_count = counts[0]["generated"] if counts else 0
            app_count = counts[0]["approved"] if counts else 0
            dup_count = counts[0]["duplicates"] if counts else 0
            # SUM 在没有匹配行时返回 NULL，统一成 0 避免后续 None 比较
            gen_count = gen_count or 0
            app_count = app_count or 0
            dup_count = dup_count or 0

        async with get_db() as db:
            await db.execute(
                """UPDATE loop_round
                   SET generated=?, approved=?, duplicates=?, status='testing'
                   WHERE id=?""",
                (gen_count, app_count, dup_count, round_id),
            )
            await db.commit()

        # 收敛检测
        if app_count == 0:
            consecutive_empty_rounds += 1
            if consecutive_empty_rounds >= 2:
                break
        else:
            consecutive_empty_rounds = 0

        # 4. 召回测试阶段
        if app_count > 0:
            await _run_single_jump_for_round(
                loop_task_id, loop_task_name, round_number, round_id, qa_task_id,
                env_url, org_id, d_user_id, agent_id, top_k, recall_top_k, concurrency, cross_chunk
            )

        # 阶段边界：召回测试完成后检查暂停
        if await check_pause_between_stages():
            return

        # 5. 更新累计统计
        await _update_loop_stats(loop_task_id)

        async with get_db() as db:
            await db.execute(
                "UPDATE loop_round SET status='done', finished_at=? WHERE id=?",
                (_now(), round_id),
            )
            await db.commit()

    # Loop finished normally
    async with get_db() as db:
        await db.execute(
            "UPDATE loop_task SET status='done', finished_at=? WHERE id=?",
            (_now(), loop_task_id),
        )
        await db.commit()


async def _get_existing_questions(loop_task_id: str, global_dedup: bool = False) -> dict[str, list[str]]:
    """Get all approved questions, grouped by section_path.

    Args:
        loop_task_id: Current loop task ID
        global_dedup: If True, get all approved questions from database (cross-task dedup)
                     If False, only get questions from this loop task (default)
    """
    async with get_db() as db:
        if global_dedup:
            # 全局去重：获取所有已批准的问题（跨任务）
            rows = await db.execute_fetchall(
                """SELECT q.id, q.section_path, q.question
                   FROM qa_gen_question q
                   WHERE q.status = 'approved'
                   ORDER BY q.created_at""",
            )
        else:
            # 任务内去重：只获取当前循环任务的问题
            rows = await db.execute_fetchall(
                """SELECT q.id, q.section_path, q.question
                   FROM qa_gen_question q
                   JOIN loop_round lr ON q.task_id = lr.qa_gen_task_id
                   WHERE lr.loop_task_id = ? AND q.status = 'approved'
                   ORDER BY q.created_at""",
                (loop_task_id,),
            )

    result: dict[str, list] = {}
    for row in rows:
        sp = row["section_path"]
        if sp not in result:
            result[sp] = []
        result[sp].append({"id": row["id"], "question": row["question"]})

    return result


async def _get_new_questions(qa_task_id: str) -> list[dict]:
    """Get all questions from a QA task."""
    async with get_db() as db:
        rows = await db.execute_fetchall(
            "SELECT id, question FROM qa_gen_question WHERE task_id=?",
            (qa_task_id,),
        )
    return [{"id": r["id"], "question": r["question"]} for r in rows]


async def _get_new_questions_by_chunk(qa_task_id: str) -> dict[str, list[dict]]:
    """按切片分组获取新问题。

    Returns:
        {chunk_id: [{id, question, ...}]}
    """
    async with get_db() as db:
        rows = await db.execute_fetchall(
            """SELECT id, question, chunk_id, section_path
               FROM qa_gen_question
               WHERE task_id=?""",
            (qa_task_id,),
        )

    result: dict[str, list] = {}
    for row in rows:
        chunk_id = row["chunk_id"] or row["section_path"] or "default"
        if chunk_id not in result:
            result[chunk_id] = []
        result[chunk_id].append({
            "id": row["id"],
            "question": row["question"],
            "chunk_id": row["chunk_id"],
            "section_path": row["section_path"],
        })

    return result


async def _get_existing_questions_by_chunk(
    loop_task_id: str,
    exclude_qa_task_id: str | None = None,
    global_dedup: bool = False,
) -> dict[str, list[tuple[str, str]]]:
    """按切片分组获取已有问题（用于查重）。

    Args:
        loop_task_id: 当前循环任务ID
        exclude_qa_task_id: 排除的 qa_gen_task_id（即本轮刚生成的一批，避免自己查自己）
        global_dedup: 是否全局去重（跨任务）

    Returns:
        {chunk_id: [(id, question)]}
    """
    async with get_db() as db:
        if global_dedup:
            # 全局去重：获取所有已批准的问题，但排除本轮 qa_task
            if exclude_qa_task_id:
                rows = await db.execute_fetchall(
                    """SELECT id, chunk_id, section_path, question
                       FROM qa_gen_question
                       WHERE status = 'approved' AND task_id != ?
                       ORDER BY created_at""",
                    (exclude_qa_task_id,),
                )
            else:
                rows = await db.execute_fetchall(
                    """SELECT id, chunk_id, section_path, question
                       FROM qa_gen_question
                       WHERE status = 'approved'
                       ORDER BY created_at""",
                )
        else:
            # 任务内去重：只获取当前循环任务的问题，但排除本轮 qa_task
            if exclude_qa_task_id:
                rows = await db.execute_fetchall(
                    """SELECT q.id, q.chunk_id, q.section_path, q.question
                       FROM qa_gen_question q
                       JOIN loop_round lr ON q.task_id = lr.qa_gen_task_id
                       WHERE lr.loop_task_id = ?
                         AND q.status = 'approved'
                         AND q.task_id != ?
                       ORDER BY q.created_at""",
                    (loop_task_id, exclude_qa_task_id),
                )
            else:
                rows = await db.execute_fetchall(
                    """SELECT q.id, q.chunk_id, q.section_path, q.question
                       FROM qa_gen_question q
                       JOIN loop_round lr ON q.task_id = lr.qa_gen_task_id
                       WHERE lr.loop_task_id = ? AND q.status = 'approved'
                       ORDER BY q.created_at""",
                    (loop_task_id,),
                )

    result: dict[str, list] = {}
    for row in rows:
        chunk_id = row["chunk_id"] or row["section_path"] or "default"
        if chunk_id not in result:
            result[chunk_id] = []
        result[chunk_id].append((row["id"], row["question"]))

    return result


async def _run_single_jump_for_round(
    loop_task_id: str,
    loop_task_name: str,
    round_number: int,
    round_id: str,
    qa_task_id: str,
    env_url: str,
    org_id: str,
    d_user_id: str,
    agent_id: str,
    top_k: int,
    recall_top_k: int,
    concurrency: int,
    cross_chunk: bool,
):
    """Run single-jump test for a round's approved questions."""

    def stop_check():
        ctrl = _loop_controls.get(loop_task_id)
        return ctrl is None or ctrl.get("stop", False)

    # Check stop before starting
    if stop_check():
        return

    # Create single-jump task
    sj_task_id = _id()
    async with get_db() as db:
        await db.execute(
            """INSERT INTO single_jump_task
               (id,name,env_url,org_id,d_user_id,agent_id,top_k,recall_top_k,concurrency,cross_chunk,
                status,progress,total,created_at,hit_top_k)
               VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?,?,?)""",
            (sj_task_id, f"{loop_task_name}-单跳测试-第{round_number}轮", env_url, org_id, d_user_id,
             agent_id, top_k, recall_top_k, concurrency, int(cross_chunk), "pending", 0, 0, _now(), top_k),
        )
        await db.execute(
            "UPDATE loop_round SET single_jump_task_id=? WHERE id=?",
            (sj_task_id, round_id),
        )
        await db.commit()

    # Build MD content from approved questions
    # Query approved questions from this QA task
    async with get_db() as db:
        rows = await db.execute_fetchall(
            """SELECT section_path, file_name, file_id, question, reference_answer, chunk_id, chunk_headers
               FROM qa_gen_question
               WHERE task_id=? AND status='approved'
               ORDER BY chunk_headers, created_at""",
            (qa_task_id,),
        )

    if not rows:
        # No approved questions, skip test
        return

    # Check stop before running test
    if stop_check():
        return

    # Group by chunk_headers (use section_path as fallback)
    from collections import defaultdict
    sections_dict: dict[str, list] = defaultdict(list)
    question_chunk_map: dict[str, str] = {}  # question -> chunk_id
    # section_key -> {file_id, file_name} from qa_gen_question
    section_file_info: dict[str, dict] = {}

    for row in rows:
        # Use chunk_headers as the grouping key if available, otherwise use section_path
        section_key = row["chunk_headers"] if row["chunk_headers"] else row["section_path"]
        if not section_key:
            section_key = row["file_name"] or "default"
        sections_dict[section_key].append({
            "question": row["question"],
            "reference_answer": row["reference_answer"],
            "file_name": row["file_name"],
            "chunk_headers": row["chunk_headers"],
            "chunk_id": row["chunk_id"],
        })
        # Build question to chunk_id mapping
        if row["chunk_id"] and row["question"]:
            question_chunk_map[row["question"]] = row["chunk_id"]
        # Remember file info for this section_key (first non-empty file_id wins)
        if row["file_id"] and section_key not in section_file_info:
            section_file_info[section_key] = {
                "file_id": row["file_id"],
                "file_name": row["file_name"] or "",
            }

    # Generate MD（与 HTTP 导出、离线脚本共用 loop_recall_md）
    prebuilt_file_map: dict[str, dict] = {}
    md_lines: list[str] = []

    section_index = 0
    for section_key, items in sections_dict.items():
        section_index += 1
        file_name = (items[0].get("file_name") or "").strip()
        slice_title = (items[0].get("chunk_headers") or "").strip() or section_key

        parsed_section_path = append_recall_md_section(
            md_lines,
            section_index,
            file_name=file_name,
            slice_title=slice_title,
            qa_items=items,
            meta_lines=[DEFAULT_LLM_NOTE],
        )
        finfo = section_file_info.get(section_key)
        if finfo:
            prebuilt_file_map[parsed_section_path] = {
                "file_id": finfo["file_id"],
                "file_name": finfo["file_name"],
                "match_type": "exact",
            }

    md_content = "\n".join(md_lines)

    # Check stop before running test
    if stop_check():
        return

    # Run single-jump test
    from api.single_jump import _run_task

    # Import necessary modules
    import sys
    from pathlib import Path
    sys.path.insert(0, str(Path(__file__).parent.parent / "sdk"))

    await _run_task(
        task_id=sj_task_id,
        qa_text=md_content,
        env_url=env_url,
        org_id=org_id,
        d_user_id=d_user_id,
        agent_id=agent_id,
        hit_top_k=top_k,
        recall_top_k=recall_top_k,
        concurrency=concurrency,
        cross_chunk=cross_chunk,
        prebuilt_file_map=prebuilt_file_map if prebuilt_file_map else None,
        prebuilt_chunk_map=question_chunk_map if question_chunk_map else None,
    )

    # After test completes, aggregate stats from single_jump_result
    async with get_db() as db:
        # Wait a bit for the test to complete (polling)
        max_wait = 1800  # Max 30 minutes wait for large tasks
        waited = 0
        while waited < max_wait:
            # Check stop during polling
            if stop_check():
                return

            row = await db.execute_fetchall(
                "SELECT status FROM single_jump_task WHERE id=?",
                (sj_task_id,)
            )
            if row and row[0]["status"] in ("done", "failed"):
                break
            await asyncio.sleep(2)
            waited += 2

        # Aggregate stats
        stats_rows = await db.execute_fetchall(
            """SELECT
                COUNT(*) as tested,
                SUM(CASE WHEN error IS NULL AND COALESCE(json_array_length(retrieved), 0) > 0 THEN 1 ELSE 0 END) as recalled,
                SUM(CASE WHEN is_file_hit = 1 THEN 1 ELSE 0 END) as file_hit,
                SUM(CASE WHEN is_chunk_hit = 1 THEN 1 ELSE 0 END) as chunk_hit
            FROM single_jump_result
            WHERE task_id=?""",
            (sj_task_id,)
        )

        if stats_rows:
            stats = dict(stats_rows[0])
            await db.execute(
                """UPDATE loop_round
                   SET tested=?, recalled=?, file_hit=?, chunk_hit=?
                   WHERE id=?""",
                (stats.get("tested") or 0, stats.get("recalled") or 0,
                 stats.get("file_hit") or 0, stats.get("chunk_hit") or 0,
                 round_id),
            )
            await db.commit()


async def _update_loop_stats(loop_task_id: str):
    """Update cumulative stats from all rounds."""
    async with get_db() as db:
        # Aggregate from loop_round
        rows = await db.execute_fetchall(
            """SELECT
                SUM(generated) as total_generated,
                SUM(approved) as total_approved,
                SUM(duplicates) as total_duplicates,
                SUM(tested) as total_tested,
                SUM(recalled) as total_recalled,
                SUM(file_hit) as total_file_hit,
                SUM(chunk_hit) as total_chunk_hit
            FROM loop_round WHERE loop_task_id=?""",
            (loop_task_id,),
        )

        stats = dict(rows[0]) if rows else {}

        # Count file_miss and recall_failed from single_jump_result
        miss_rows = await db.execute_fetchall(
            """SELECT
                SUM(CASE WHEN r.is_file_hit=0 AND COALESCE(json_array_length(r.retrieved), 0)>0 THEN 1 ELSE 0 END) as file_miss,
                SUM(CASE WHEN COALESCE(json_array_length(r.retrieved), 0)=0 AND r.error IS NULL THEN 1 ELSE 0 END) as recall_failed
            FROM single_jump_result r
            JOIN loop_round lr ON r.task_id = lr.single_jump_task_id
            WHERE lr.loop_task_id=?""",
            (loop_task_id,),
        )

        miss_stats = dict(miss_rows[0]) if miss_rows else {}

        await db.execute(
            """UPDATE loop_task SET
                total_generated=?,
                total_approved=?,
                total_duplicates=?,
                total_tested=?,
                total_recalled=?,
                total_file_hit=?,
                total_file_miss=?,
                total_recall_failed=?,
                total_chunk_hit=?
            WHERE id=?""",
            (
                stats.get("total_generated") or 0,
                stats.get("total_approved") or 0,
                stats.get("total_duplicates") or 0,
                stats.get("total_tested") or 0,
                stats.get("total_recalled") or 0,
                stats.get("total_file_hit") or 0,
                miss_stats.get("file_miss") or 0,
                miss_stats.get("recall_failed") or 0,
                stats.get("total_chunk_hit") or 0,
                loop_task_id,
            ),
        )
        await db.commit()


async def recover_orphaned_loops():
    """On startup, set any 'running' loop tasks to 'paused'."""
    async with get_db() as db:
        rows = await db.execute_fetchall(
            "SELECT id FROM loop_task WHERE status='running'"
        )
        for row in rows:
            await db.execute(
                "UPDATE loop_task SET status='paused', paused_at=? WHERE id=?",
                (_now(), row["id"]),
            )
        await db.commit()