"""
召回测试执行器：对每条问答对调用 dagent 语义召回接口，记录结果。
"""
import asyncio
import json
import sys
import time
from dataclasses import dataclass, field
import aiohttp

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

from .parser import Section, QAPair


@dataclass
class RecallResult:
    section_path: str
    doc_name: str
    file_id: str | None
    match_type: str | None        # exact / contains / fuzzy / unmatched
    qid: str
    question: str
    reference_answer: str
    top_k: int                    # 用于判断命中的top_k值
    hit_top_k: int                # 用于判断切片是否命中的top_k阈值（可能不同于召回时的top_k）
    retrieved: list[dict] = field(default_factory=list)   # 召回的切片列表（全部，不截断）
    latency_ms: int = 0
    error: str | None = None
    expected_chunk_id: str | None = None  # 期望命中的切片ID
    raw_chunk_headers: str | None = None  # 原始切片标题（从元数据解析）

    # 计算属性
    @property
    def best_cosine_sim(self) -> float | None:
        sims = [1.0 - r.get("cosine_distance_1", 1.0) for r in self.retrieved if r.get("cosine_distance_1") is not None]
        return round(max(sims), 4) if sims else None

    @property
    def avg_cosine_sim(self) -> float | None:
        sims = [1.0 - r.get("cosine_distance_1", 1.0) for r in self.retrieved if r.get("cosine_distance_1") is not None]
        return round(sum(sims) / len(sims), 4) if sims else None

    @property
    def is_empty(self) -> bool:
        return len(self.retrieved) == 0

    @property
    def retrieved_file_ids(self) -> list[str]:
        return list({r.get("file_id", "") for r in self.retrieved if r.get("file_id")})

    @property
    def retrieved_chunk_ids(self) -> list[str]:
        """获取召回的所有切片ID"""
        chunk_ids = []
        for r in self.retrieved:
            chunk_id = r.get("knowledge_md_header_split_id") or r.get("id") or r.get("chunk_id")
            if chunk_id:
                chunk_ids.append(chunk_id)
        return chunk_ids

    @property
    def is_chunk_hit(self) -> bool:
        """检查期望切片是否在召回结果的前hit_top_k个结果中"""
        if not self.expected_chunk_id:
            return False
        return self.expected_chunk_id in self.retrieved_chunk_ids[:self.hit_top_k]

    @property
    def chunk_hit_rank(self) -> int | None:
        """返回期望切片在召回结果中的排名（1-based），未命中返回None

        只在hit_top_k范围内查找，超出范围视为未命中
        """
        if not self.expected_chunk_id:
            return None
        try:
            idx = self.retrieved_chunk_ids[:self.hit_top_k].index(self.expected_chunk_id)
            return idx + 1
        except ValueError:
            return None

    @property
    def is_file_hit(self) -> bool:
        """检查期望文件是否在召回结果的前hit_top_k个结果中"""
        if not self.file_id:
            return False
        # 获取前hit_top_k个结果的file_ids
        top_file_ids = []
        for r in self.retrieved[:self.hit_top_k]:
            fid = r.get("file_id")
            if fid:
                top_file_ids.append(fid)
        return self.file_id in top_file_ids


class RecallTester:
    def __init__(self, env_url: str, org_id: str, d_user_id: str = "test"):
        self.env_url = env_url.rstrip("/")
        self.org_id = org_id
        self.headers = {
            "Content-Type": "application/json",
            "d-user-id": d_user_id,
            "org-id": org_id,
        }

    async def _recall_one(
        self,
        session: aiohttp.ClientSession,
        question: str,
        file_id_list: list[str] | None,
        recall_top_k: int,  # 用于API调用时的top_k，可以设置较大值获取所有结果
        agent_id: str = "",  # 用于召回测试的 agent ID
    ) -> tuple[list[dict], int]:
        # 如果提供了 agent_id，使用 agent chat API 进行召回
        if agent_id:
            return await self._recall_via_agent(session, question, agent_id, recall_top_k)

        # 否则直接使用知识库搜索 API
        url = f"{self.env_url}/dagent/knowledge/hub/semantic_search_knowledge/detail"
        payload: dict = {
            "query": question,
            "org_id": self.org_id,
            "top_k": recall_top_k,
        }
        if file_id_list:
            payload["file_id_list"] = file_id_list

        start = time.monotonic()
        # 增加超时时间到60秒，并添加重试逻辑
        max_retries = 3
        last_error = None

        for attempt in range(max_retries):
            try:
                async with session.post(url, json=payload, timeout=aiohttp.ClientTimeout(total=60)) as resp:
                    resp.raise_for_status()
                    data = await resp.json()
                break  # 成功则跳出重试循环
            except asyncio.TimeoutError as e:
                last_error = e
                print(f"[DEBUG] Recall timeout (attempt {attempt+1}/{max_retries}) for: {question[:50]}...")
                if attempt < max_retries - 1:
                    await asyncio.sleep(2 ** attempt)  # 指数退避: 1s, 2s, 4s
                else:
                    raise  # 最后一次重试失败，抛出异常
            except Exception as e:
                raise  # 其他异常直接抛出

        latency_ms = int((time.monotonic() - start) * 1000)

        # 检查 API 返回的业务错误码
        code = data.get("code")
        if code is not None and code != 0:
            msg = data.get("msg", "Unknown error")
            raise Exception(f"API error: code={code}, msg={msg}")

        result_data = data.get("data", {}) or {}

        # 调试：如果结果为空，打印调试信息
        if not result_data or (not result_data.get("standard_answer_results") and not result_data.get("related_knowledge_rerank_results_top")):
            print(f"[DEBUG] Empty/No results for question: {question[:50]}...")
            print(f"[DEBUG] Response code: {data.get('code')}, msg: {data.get('msg')}")
            print(f"[DEBUG] org_id used: {self.org_id}")
            print(f"[DEBUG] Request payload: {payload}")
            print(f"[DEBUG] Response data keys: {list(data.keys())}")
            if result_data:
                print(f"[DEBUG] result_data keys: {list(result_data.keys())}")

        standard = result_data.get("standard_answer_results") or []
        rerank_top = result_data.get("related_knowledge_rerank_results_top") or []
        all_items = standard + rerank_top

        # 调试：记录召回结果数量
        if len(all_items) == 0:
            print(f"[DEBUG] No recall results for: {question[:50]}... (standard={len(standard)}, rerank={len(rerank_top)})")

        return all_items, latency_ms

    async def _recall_via_agent(
        self,
        session: aiohttp.ClientSession,
        question: str,
        agent_id: str,
        recall_top_k: int,
    ) -> tuple[list[dict], int]:
        """通过 Agent chat SSE 接口获取召回结果。

        解析策略：
        - 逐行读取 SSE（服务端单 `\n` 分隔，不是双换行）
        - 每个 EVENT.event_name == "TOOL_END" 的 event_data.items 里有一批 chunk
        - Agent 可能多轮工具调用，每次 TOOL_END 都累加；按 (file_id, paragraph_chunk_id) 去重
        - 顺序保留首次出现位置（作为伪 rank），用于命中排名统计
        """
        import uuid
        payload = {
            "chat_id": uuid.uuid4().hex,
            "task": question,
            "agent_id": agent_id,
            "llm_type": "deepseek_v3",
        }

        start = time.monotonic()
        items: list[dict] = []
        seen: set[tuple[str, str]] = set()

        try:
            async with session.post(
                f"{self.env_url}/dagent/agent/chat",
                json=payload,
                headers={"Accept": "text/event-stream"},
                timeout=aiohttp.ClientTimeout(total=300),
            ) as resp:
                resp.raise_for_status()
                line_buf = ""
                async for raw in resp.content:
                    line_buf += raw.decode("utf-8", errors="replace")
                    while "\n" in line_buf:
                        line, line_buf = line_buf.split("\n", 1)
                        line = line.rstrip("\r")
                        if not line.startswith("data:"):
                            continue
                        data_str = line[5:].strip()
                        if not data_str or data_str == "[DONE]":
                            continue
                        try:
                            chunk = json.loads(data_str)
                        except json.JSONDecodeError:
                            continue
                        if chunk.get("message_type") != "EVENT" or chunk.get("is_chunk_data"):
                            continue
                        event_data_raw = chunk.get("data")
                        if isinstance(event_data_raw, str):
                            try:
                                event_data = json.loads(event_data_raw)
                            except json.JSONDecodeError:
                                continue
                        else:
                            event_data = event_data_raw
                        if not isinstance(event_data, dict):
                            continue
                        if event_data.get("event_name") != "TOOL_END":
                            continue
                        tool_event_data = event_data.get("event_data")
                        if not isinstance(tool_event_data, dict):
                            continue
                        reference_items = tool_event_data.get("items") or []
                        if not isinstance(reference_items, list):
                            continue
                        for item in reference_items:
                            if not isinstance(item, dict):
                                continue
                            file_id = str(item.get("file_id") or "")
                            chunk_id = str(
                                item.get("paragraph_chunk_id")
                                or item.get("knowledge_md_header_split_id")
                                or ""
                            )
                            # 跳过不带 file_id/chunk_id 的外链类条目（只有 file_name+url）
                            if not file_id and not chunk_id:
                                continue
                            key = (file_id, chunk_id)
                            if key in seen:
                                continue
                            seen.add(key)
                            items.append({
                                "file_id": file_id,
                                "file_name": "",
                                "headers": str(item.get("headers") or ""),
                                "content": item.get("active_paragraph_context")
                                    or item.get("active_context") or "",
                                "knowledge_md_header_split_id": chunk_id,
                                "id": chunk_id,
                                "paragraph_md5": str(item.get("paragraph_md5") or ""),
                                "cosine_distance_1": None,
                            })
        except Exception as e:
            print(f"[DEBUG] Agent recall error: {e}")

        latency_ms = int((time.monotonic() - start) * 1000)
        return items[:recall_top_k], latency_ms

    async def run(
        self,
        sections: list[Section],
        file_map: dict[str, dict | None],
        top_k: int = 5,           # 用于判断命中的top_k阈值
        recall_top_k: int = 100,  # 用于API调用时的top_k，默认100获取更多结果
        concurrency: int = 20,     # 增加默认并发数到20
        cross_chunk: bool = False,  # 保留参数兼容旧调用，但不再控制搜索范围
        result_cb=None,
        progress_cb=None,  # 保留兼容旧调用
        chunk_map: dict[str, str] | None = None,  # question -> expected_chunk_id
        agent_id: str = "",  # 用于召回测试的 agent ID
    ) -> list[RecallResult]:
        results: list[RecallResult] = []
        sem = asyncio.Semaphore(concurrency)
        total = sum(len(s.qa_pairs) for s in sections)
        done = 0

        async with aiohttp.ClientSession(headers=self.headers) as session:
            async def _test_one(section: Section, qa: QAPair) -> RecallResult:
                nonlocal done
                mapping = file_map.get(section.section_path)
                file_id = mapping["file_id"] if mapping else None
                match_type = mapping["match_type"] if mapping else "unmatched"

                # 优先使用 QAPair 上已注入的 chunk_id，其次从 chunk_map 查找
                expected_chunk_id = qa.expected_chunk_id or (
                    chunk_map.get(qa.question) if chunk_map else None
                )

                result = RecallResult(
                    section_path=section.section_path,
                    doc_name=section.doc_name,
                    file_id=file_id,
                    match_type=match_type,
                    qid=qa.qid,
                    question=qa.question,
                    reference_answer=qa.answer,
                    top_k=top_k,
                    hit_top_k=top_k,  # 用于判断命中的阈值
                    expected_chunk_id=expected_chunk_id,
                    raw_chunk_headers=section.raw_chunk_headers,
                )

                # 始终全库搜索（不传 file_id_list），以切片命中为主要指标
                # 使用较大的 recall_top_k 获取所有召回结果
                async with sem:
                    try:
                        chunks, latency = await self._recall_one(session, qa.question, None, recall_top_k, agent_id)
                        result.retrieved = chunks
                        result.latency_ms = latency
                        # 调试：记录召回结果数量
                        if len(chunks) == 0:
                            print(f"[DEBUG] Empty recall for question: {qa.question[:60]}... (section: {section.section_path[:40]}...)")
                    except Exception as e:
                        result.error = str(e)
                        print(f"[DEBUG] Recall error for question: {qa.question[:60]}... Error: {e}")

                done += 1
                if result_cb:
                    await result_cb(result, done, total)
                elif progress_cb and (done % 10 == 0 or done == total):
                    await progress_cb(done, total)
                return result

            tasks = [
                _test_one(section, qa)
                for section in sections
                for qa in section.qa_pairs
            ]
            for coro in asyncio.as_completed(tasks):
                results.append(await coro)

        return results