"""
LLM 客户端模块 (LLM Client Module) - 简化 ID 锚点匹配版

该模块负责发送翻译请求，并使用简化后的 ID 作为锚点解析 LLM 的响应。
"""

import asyncio
from openai import AsyncOpenAI
from typing import List, Dict, Optional, Any
from loguru import logger
import time
import re
from .manifest_manager import ManifestItem


class RateLimiter:
    """并发与 RPM 速率限制器。"""
    def __init__(self, requests_per_minute: int, concurrent_requests: int):
        self.semaphore = asyncio.Semaphore(concurrent_requests)
        self.min_interval = 60.0 / requests_per_minute if requests_per_minute > 0 else 0
        self.last_request_time = 0

    async def acquire(self):
        await self.semaphore.acquire()
        current_time = time.time()
        wait_time = self.min_interval - (current_time - self.last_request_time)
        if wait_time > 0:
            await asyncio.sleep(wait_time)
        self.last_request_time = time.time()

    def release(self):
        self.semaphore.release()


class OpenRouterClient:
    """基于简化 ID 锚点匹配逻辑的 LLM 客户端。"""
    
    def __init__(self, config: Dict):
        self.config = config
        or_config = config["openrouter"]
        api_key = or_config.get("api_key")
        if not api_key or api_key == "YOUR_OPENROUTER_API_KEY":
            raise ValueError("请设置有效的 OpenRouter API Key")
            
        self.client = AsyncOpenAI(
            base_url=or_config["base_url"],
            api_key=api_key,
            default_headers={"HTTP-Referer": "https://github.com/epub-translator", "X-Title": "EPUB Translator"}
        )
        self.models = or_config["models"]
        self.rate_limiter = RateLimiter(
            or_config["rate_limits"]["requests_per_minute"],
            or_config["rate_limits"]["concurrent_requests"]
        )

    async def translate_chunk(self, items: List[ManifestItem], model_type: str = "production") -> Dict[str, str]:
        """
        翻译一个段落块。
        """
        if not items: return {}
        
        prompt = self._build_prompt(items)
        model = self.models.get(model_type, self.models["production"])
        
        try:
            raw_response = await self._make_request(prompt, model)
            if not raw_response:
                return {item.global_id: f"[翻译失败 - API空响应]" for item in items}
            
            return self._parse_with_anchors(raw_response, items)
            
        except Exception as e:
            logger.error(f"翻译请求异常: {e}")
            return {item.global_id: f"[翻译失败 - {str(e)}]" for item in items}

    def _build_prompt(self, items: List[ManifestItem]) -> str:
        lines = ["请将以下编号的英文段落翻译成中文。每段翻译前必须带上原编号，格式为 p_xxxxx。", "原文：", ""]
        for item in items:
            lines.append(f"{item.global_id} {item.clean_text}")
        lines.extend(["", "要求：只返回翻译，不要解释。必须保留原编号。"])
        return "\n".join(lines)

    def _parse_with_anchors(self, response: str, items: List[ManifestItem]) -> Dict[str, str]:
        results = {}
        
        for i, item in enumerate(items):
            current_id = item.global_id
            
            # 使用简单的字符串拼接构造正则
            start_pat = r"\b" + current_id + r"\b"
            start_match = re.search(start_pat, response)
            if not start_match:
                continue
            
            start_pos = start_match.start()
            end_pos = len(response)
            
            if i + 1 < len(items):
                next_id = items[i+1].global_id
                next_pat = r"\b" + next_id + r"\b"
                next_match = re.search(next_pat, response[start_pos + 1:])
                if next_match:
                    end_pos = start_pos + 1 + next_match.start()
            
            segment = response[start_pos:end_pos].strip()
            
            # 清洗 ID
            # 注意：这里的正则使用了更稳健的拼接方式
            id_clean_pat = r"^(\[?" + current_id + r"\]?[:：\s]*)+"
            cleaned = re.sub(id_clean_pat, "", segment).strip()
            
            if cleaned:
                results[current_id] = cleaned

        if len(results) < len(items):
            expected_ids = {item.global_id for item in items}
            missing_ids = expected_ids - set(results.keys())
            logger.warning(f"锚点匹配缺失 {len(missing_ids)} 个。尝试降级解析...")
            fallback_results = self._fallback_parse(response, items)
            for pid, trans in fallback_results.items():
                if pid not in results:
                    results[pid] = trans

        return results

    def _fallback_parse(self, response: str, items: List[ManifestItem]) -> Dict[str, str]:
        results = {}
        lines = [l.strip() for l in response.split("\n") if l.strip()]
        for line in lines:
            for item in items:
                if item.global_id in line:
                    # 同样的稳健拼接
                    fallback_clean_pat = r"^\(?" + item.global_id + r"\)?[:：\s]*"
                    clean_line = re.sub(fallback_clean_pat, "", line).strip()
                    if clean_line:
                        results[item.global_id] = clean_line
        return results

    async def _make_request(self, prompt: str, model: str) -> str:
        await self.rate_limiter.acquire()
        try:
            resp = await self.client.chat.completions.create(
                model=model,
                messages=[
                    {"role": "system", "content": "你是一位专业的翻译。请严格按编号输出翻译。"},
                    {"role": "user", "content": prompt}
                ],
                temperature=0.2,
                max_tokens=8000
            )
            return resp.choices[0].message.content.strip()
        finally:
            self.rate_limiter.release()

    async def close(self):
        await self.client.close()