"""LLM 翻译 + 投资事件抽取主流程。

输入: ProcessedArticle（M2/M3 输出）
输出: EnTranslatedArticle（含中英文双语内容 + 抽取事件 + 元信息）

设计:
    1. 加载 prompts/translation_and_extraction.md，分离 system/user 模板
    2. 单次 LLM 调用完成翻译 + 事件抽取（节省 token）
    3. 解析 JSON → Pydantic LLMTranslationOutput 强校验 + 重试
    4. 限制正文长度避免触顶 context window
"""

import asyncio
import json
import logging
import re
import time
from pathlib import Path

from openai import AsyncOpenAI, OpenAI

from extractor.models import ProcessedArticle
from llm.client import LLMConfig
from llm.models import (
    EnTranslatedArticle,
    LLMCallError,
    LLMTranslationOutput,
)

logger = logging.getLogger(__name__)

# Prompt 模板路径
DEFAULT_PROMPT_PATH = Path("prompts/translation_and_extraction.md")

# 文章正文截断长度（保守取 8K 字符，避免超出上下文窗口）
MAX_CONTENT_CHARS = 8000

# 重试设置
DEFAULT_MAX_ATTEMPTS = 3
RETRY_BASE_WAIT_SEC = 1.0
RETRY_MAX_WAIT_SEC = 8.0

# Prompt 中 system / user 分隔标记
_SYSTEM_SECTION_START = "## System Prompt"
_USER_SECTION_START = "## User Input"


class PromptTemplate:
    """Prompt 模板加载器。

    从 prompts/translation_and_extraction.md 读取模板，
    分离 System Prompt 和 User Input 两部分。
    User Input 支持 {title} / {source_name} / {publish_time} / {content} 占位符。
    """

    def __init__(self, template_path: str | Path = DEFAULT_PROMPT_PATH) -> None:
        self._path = Path(template_path)
        raw = self._path.read_text(encoding="utf-8")

        # 分离 system 和 user 两部分
        self._system_prompt, self._user_template = self._parse_template(raw)

    @staticmethod
    def _parse_template(raw: str) -> tuple[str, str]:
        """解析模板文件，返回 (system_prompt, user_template)。"""
        # 找到 ## System Prompt 之后的内容直到 ## User Input
        sys_match = re.search(
            r"## System Prompt\s*\n(.*?)(?=---\s*\n## User Input)",
            raw, re.DOTALL
        )
        user_match = re.search(
            r"## User Input\s*\n(.*)",
            raw, re.DOTALL
        )

        system = sys_match.group(1).strip() if sys_match else ""
        user = user_match.group(1).strip() if user_match else raw

        return system, user

    def render(self, article: ProcessedArticle) -> tuple[str, str]:
        """渲染 Prompt，返回 (system_prompt, user_prompt)。

        对正文做截断处理。
        """
        content = article.content
        if len(content) > MAX_CONTENT_CHARS:
            logger.debug(
                "文章 %s 超长截断: %d → %d",
                article.url_hash, len(content), MAX_CONTENT_CHARS,
            )
            content = content[:MAX_CONTENT_CHARS] + "\n\n[正文过长已截断]"

        user_prompt = (
            self._user_template
            .replace("{title}", article.title)
            .replace("{source_name}", article.source_name)
            .replace("{publish_time}", article.publish_time or "未知")
            .replace("{content}", content)
        )

        return self._system_prompt, user_prompt


# --------------------------------------------------------------------------- #
# JSON 提取（LLM 偶尔会包 ```json 围栏）
# --------------------------------------------------------------------------- #


def _extract_json_object(text: str) -> str:
    """从 LLM 输出中提取首个 JSON 对象字符串（去围栏 / 取首个 {...}）。"""
    s = text.strip()
    if s.startswith("```"):
        s = s.strip("`")
        if s.lower().startswith("json"):
            s = s[4:].lstrip("\n").lstrip()
        if s.endswith("```"):
            s = s[:-3]
    start = s.find("{")
    if start < 0:
        return s
    depth = 0
    for i in range(start, len(s)):
        if s[i] == "{":
            depth += 1
        elif s[i] == "}":
            depth -= 1
            if depth == 0:
                return s[start : i + 1]
    return s[start:]


def parse_translation_json(raw: str) -> LLMTranslationOutput:
    """把 LLM 输出文本解析为 LLMTranslationOutput（可能抛 LLMCallError）。"""
    payload = _extract_json_object(raw)
    try:
        obj = json.loads(payload)
    except json.JSONDecodeError as e:
        raise LLMCallError(f"JSON 解析失败: {e}") from e
    if not isinstance(obj, dict):
        raise LLMCallError(f"JSON 顶层非对象: {type(obj).__name__}")
    try:
        return LLMTranslationOutput.model_validate(obj)
    except Exception as e:
        raise LLMCallError(f"翻译输出 schema 校验失败: {e}") from e


# --------------------------------------------------------------------------- #
# 词数计算
# --------------------------------------------------------------------------- #


def _count_zh_chars(text: str) -> int:
    """统计中文文本字数（汉字 + 数字 + 字母混合）。"""
    # 简单统计：去除空白后长度
    return len(text.replace(" ", "").replace("\n", "").replace("\r", ""))


# --------------------------------------------------------------------------- #
# 同步 / 异步 LLM 调用
# --------------------------------------------------------------------------- #


def _call_llm_sync(
    client: OpenAI,
    config: LLMConfig,
    system_prompt: str,
    user_prompt: str,
) -> tuple[str, dict[str, int | None]]:
    """同步单次 LLM 调用，返回 (raw_text, usage)。"""
    resp = client.chat.completions.create(
        model=config.model,
        messages=[
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_prompt},
        ],
        temperature=config.temperature,
        max_tokens=config.max_tokens,
        response_format={"type": "json_object"},
    )
    text = resp.choices[0].message.content or ""
    usage = {
        "prompt_tokens": getattr(resp.usage, "prompt_tokens", None) if resp.usage else None,
        "completion_tokens": (
            getattr(resp.usage, "completion_tokens", None) if resp.usage else None
        ),
    }
    return text, usage


async def _call_llm_async(
    client: AsyncOpenAI,
    config: LLMConfig,
    system_prompt: str,
    user_prompt: str,
) -> tuple[str, dict[str, int | None]]:
    """异步单次 LLM 调用。"""
    resp = await client.chat.completions.create(
        model=config.model,
        messages=[
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_prompt},
        ],
        temperature=config.temperature,
        max_tokens=config.max_tokens,
        response_format={"type": "json_object"},
    )
    text = resp.choices[0].message.content or ""
    usage = {
        "prompt_tokens": getattr(resp.usage, "prompt_tokens", None) if resp.usage else None,
        "completion_tokens": (
            getattr(resp.usage, "completion_tokens", None) if resp.usage else None
        ),
    }
    return text, usage


# --------------------------------------------------------------------------- #
# 主入口：翻译 + 事件抽取
# --------------------------------------------------------------------------- #


def translate_and_extract(
    client: OpenAI,
    config: LLMConfig,
    article: ProcessedArticle,
    *,
    template: PromptTemplate | None = None,
    max_attempts: int = DEFAULT_MAX_ATTEMPTS,
) -> EnTranslatedArticle:
    """同步翻译 + 事件抽取（单篇文章，带重试）。

    Args:
        client: OpenAI 同步客户端
        config: LLM 配置
        article: 待处理的英文新闻
        template: Prompt 模板，默认加载 prompts/translation_and_extraction.md
        max_attempts: 最大重试次数

    Returns:
        EnTranslatedArticle 含双语内容 + 事件

    Raises:
        LLMCallError: 所有重试均失败
    """
    tpl = template or PromptTemplate()
    system_prompt, user_prompt = tpl.render(article)

    last_err: Exception | None = None
    for attempt in range(1, max_attempts + 1):
        try:
            raw, usage = _call_llm_sync(client, config, system_prompt, user_prompt)
            output = parse_translation_json(raw)

            # 检查 content_zh 非空
            if not output.content_zh.strip():
                raise LLMCallError("LLM 返回的 content_zh 为空")

            return EnTranslatedArticle(
                source_id=article.source_id,
                source_name=article.source_name,
                url=article.url,
                url_hash=article.url_hash,
                title=article.title,
                title_zh=output.title_zh,
                content_en=article.content,
                content_zh=output.content_zh,
                publish_time=article.publish_time,
                word_count=article.word_count,
                word_count_zh=_count_zh_chars(output.content_zh),
                events=output.events,
                provider=config.provider,
                model=config.model,
                attempts=attempt,
                prompt_tokens=usage.get("prompt_tokens"),
                completion_tokens=usage.get("completion_tokens"),
            )

        except LLMCallError as e:
            last_err = e
            logger.warning(
                "LLM 翻译抽取失败 url=%s 尝试 %d/%d: %s",
                article.url, attempt, max_attempts, e.reason,
            )
        except Exception as e:
            last_err = e
            logger.warning(
                "LLM 调用异常 url=%s 尝试 %d/%d: %s: %s",
                article.url, attempt, max_attempts, type(e).__name__, e,
            )

        if attempt < max_attempts:
            wait = min(RETRY_BASE_WAIT_SEC * (2 ** (attempt - 1)), RETRY_MAX_WAIT_SEC)
            time.sleep(wait)

    raise LLMCallError(
        f"LLM 翻译抽取放弃，共 {max_attempts} 次尝试: {last_err}",
        attempts=max_attempts,
    )


async def translate_and_extract_async(
    client: AsyncOpenAI,
    config: LLMConfig,
    article: ProcessedArticle,
    *,
    template: PromptTemplate | None = None,
    max_attempts: int = DEFAULT_MAX_ATTEMPTS,
    semaphore: asyncio.Semaphore | None = None,
) -> EnTranslatedArticle:
    """异步翻译 + 事件抽取（批处理用），与同步版逻辑等价。"""
    tpl = template or PromptTemplate()
    system_prompt, user_prompt = tpl.render(article)

    async def _run() -> EnTranslatedArticle:
        last_err: Exception | None = None
        for attempt in range(1, max_attempts + 1):
            try:
                raw, usage = await _call_llm_async(client, config, system_prompt, user_prompt)
                output = parse_translation_json(raw)

                if not output.content_zh.strip():
                    raise LLMCallError("LLM 返回的 content_zh 为空")

                return EnTranslatedArticle(
                    source_id=article.source_id,
                    source_name=article.source_name,
                    url=article.url,
                    url_hash=article.url_hash,
                    title=article.title,
                    title_zh=output.title_zh,
                    content_en=article.content,
                    content_zh=output.content_zh,
                    publish_time=article.publish_time,
                    word_count=article.word_count,
                    word_count_zh=_count_zh_chars(output.content_zh),
                    events=output.events,
                    provider=config.provider,
                    model=config.model,
                    attempts=attempt,
                    prompt_tokens=usage.get("prompt_tokens"),
                    completion_tokens=usage.get("completion_tokens"),
                )

            except LLMCallError as e:
                last_err = e
                logger.warning(
                    "LLM 翻译抽取失败 url=%s 尝试 %d/%d: %s",
                    article.url, attempt, max_attempts, e.reason,
                )
            except Exception as e:
                last_err = e
                logger.warning(
                    "LLM 调用异常 url=%s 尝试 %d/%d: %s: %s",
                    article.url, attempt, max_attempts, type(e).__name__, e,
                )

            if attempt < max_attempts:
                wait = min(RETRY_BASE_WAIT_SEC * (2 ** (attempt - 1)), RETRY_MAX_WAIT_SEC)
                await asyncio.sleep(wait)

        raise LLMCallError(
            f"LLM 翻译抽取放弃，共 {max_attempts} 次尝试: {last_err}",
            attempts=max_attempts,
        )

    if semaphore is None:
        return await _run()
    async with semaphore:
        return await _run()