# -*- coding: utf-8 -*-
"""
LLM 重排序（借鉴 zhikuncode-main MemoryRerankService.java）

对 BM25 召回的候选做语义重排：
- 调用 DeepSeekClient 让 LLM 对候选打分（0~1）
- 严格 JSON 输出 schema：{"results": [{"index": 0, "relevanceScore": 0.95}, ...]}
- 失败优雅降级：返回原 BM25 顺序
- 默认 3s 超时，512 max_tokens
"""
import json
from typing import Any, List, Optional

from .bm25 import ScoredResult
from .store import MemoryEntry


class LLMReranker:
    """对 BM25 召回候选做 LLM 语义重排。

    使用方式：
        reranker = LLMReranker(client, model="deepseek-v4-flash")
        reranked_indices = reranker.rerank(query, candidate_entries, top_k=3)
        # 失败时返回原顺序，不抛异常
    """

    # 严格 JSON 输出 schema（在 prompt 中明确要求）
    OUTPUT_SCHEMA_HINT = (
        '严格输出 JSON 格式：{"results": [{"index": 0, "relevanceScore": 0.95}, ...]}'
        '，index 是候选索引（从 0 开始），relevanceScore 是 0~1 的相关性分数。'
        '不要输出任何其他内容。'
    )

    def __init__(self, client: Any, model: str = "deepseek-v4-flash",
                 timeout: float = 3.0, max_tokens: int = 512,
                 enabled: bool = True):
        """
        Args:
            client: DeepSeekClient 实例（含 chat 方法）
            model: 用哪个模型做 rerank（默认 flash 便宜快）
            timeout: 单次 rerank 超时（秒）
            max_tokens: LLM 输出最大 token
            enabled: 是否启用 rerank，False 则直接返回原顺序
        """
        self._client = client
        self._model = model
        self._timeout = timeout
        self._max_tokens = max_tokens
        self._enabled = enabled

    @property
    def enabled(self) -> bool:
        return self._enabled

    @enabled.setter
    def enabled(self, value: bool) -> None:
        self._enabled = bool(value)

    def rerank(self, query: str, candidates: List[MemoryEntry],
               top_k: int = 3) -> List[int]:
        """对 candidates 做重排，返回按相关性降序的索引列表（指向 candidates）。

        失败或未启用时直接返回 [0, 1, ..., min(top_k, len)-1]（即原顺序前 top_k）。
        """
        if not self._enabled or not candidates or not query:
            return list(range(min(top_k, len(candidates))))

        if len(candidates) == 1:
            return [0]

        # 构造 prompt
        prompt = self._build_prompt(query, candidates)
        try:
            raw = self._call_llm(prompt)
            scored = self._parse_response(raw, len(candidates))
            if not scored:
                # 解析失败，降级
                return list(range(min(top_k, len(candidates))))
            # 按 score 降序，取 top_k
            scored.sort(key=lambda x: x[1], reverse=True)
            return [idx for idx, _ in scored[:top_k]]
        except Exception:
            # 任何异常都优雅降级
            return list(range(min(top_k, len(candidates))))

    def _build_prompt(self, query: str, candidates: List[MemoryEntry]) -> str:
        """构造 rerank prompt"""
        lines = [f"用户查询：{query}", "", "候选记忆："]
        for i, e in enumerate(candidates):
            # 截断过长内容，避免 prompt 膨胀
            content = e.content[:200] + "..." if len(e.content) > 200 else e.content
            title = e.title or "(无标题)"
            lines.append(f"[{i}] 标题：{title}")
            lines.append(f"    内容：{content}")
        lines.append("")
        lines.append("请对每个候选与查询的相关性打分（0~1）。" + self.OUTPUT_SCHEMA_HINT)
        return "\n".join(lines)

    def _call_llm(self, prompt: str) -> str:
        """调用 LLM 返回文本输出。

        非流式调用，使用最简单形式。失败抛异常由 rerank() 捕获。
        """
        if self._client is None:
            raise RuntimeError("LLM client 未配置")

        # 兼容两种 client：
        # 1. DeepSeekClient：chat(messages, ...) 返回流式迭代器
        # 2. openai.OpenAI：client.chat.completions.create(...)
        messages = [
            {"role": "system", "content": "你是相关性打分助手，只输出 JSON。"},
            {"role": "user", "content": prompt},
        ]

        # 尝试 DeepSeekClient.chat 接口（流式）
        # 注意：DeepSeekClient.chat 是流式接口，需要累积 content
        try:
            chunks = self._client.chat(
                messages=messages,
                model=self._model,
                stream=True,
                thinking=False,
                max_tokens=self._max_tokens,
            )
            content_parts: List[str] = []
            for chunk in chunks:
                # 兼容多种 chunk 格式
                if isinstance(chunk, dict):
                    delta = chunk.get("choices", [{}])[0].get("delta", {})
                    if isinstance(delta, dict):
                        c = delta.get("content")
                        if c:
                            content_parts.append(c)
                else:
                    # SDK 对象
                    try:
                        c = chunk.choices[0].delta.content
                        if c:
                            content_parts.append(c)
                    except (AttributeError, IndexError):
                        continue
            return "".join(content_parts)
        except (TypeError, ValueError):
            # chat 方法签名不匹配，尝试回退到原生 openai 接口
            pass

        # 回退：直接调用 openai 风格
        native = getattr(self._client, "_client", None) or self._client
        if native is None:
            raise RuntimeError("无法找到可用的 LLM client")
        resp = native.chat.completions.create(
            model=self._model,
            messages=messages,
            stream=False,
            max_tokens=self._max_tokens,
        )
        return resp.choices[0].message.content or ""

    def _parse_response(self, raw: str, expected_count: int) -> List[tuple]:
        """解析 LLM 输出，返回 [(index, score), ...] 列表"""
        if not raw:
            return []
        # 尝试从 raw 中提取 JSON（容忍前后多余文本/markdown 代码块）
        text = raw.strip()
        # 去掉 markdown 代码块标记
        if text.startswith("```"):
            lines = text.split("\n")
            # 去掉首尾的 ``` 行
            if lines and lines[0].startswith("```"):
                lines = lines[1:]
            if lines and lines[-1].startswith("```"):
                lines = lines[:-1]
            text = "\n".join(lines).strip()

        try:
            data = json.loads(text)
        except json.JSONDecodeError:
            # 尝试提取 {...} 子串
            start = text.find("{")
            end = text.rfind("}")
            if start >= 0 and end > start:
                try:
                    data = json.loads(text[start: end + 1])
                except json.JSONDecodeError:
                    return []
            else:
                return []

        results = data.get("results") if isinstance(data, dict) else None
        if not isinstance(results, list):
            return []

        scored: List[tuple] = []
        for item in results:
            if not isinstance(item, dict):
                continue
            idx = item.get("index")
            score = item.get("relevanceScore")
            if not isinstance(idx, int) or idx < 0 or idx >= expected_count:
                continue
            if not isinstance(score, (int, float)):
                continue
            scored.append((idx, float(score)))
        return scored
