"""LLM 处理：调用 NVIDIA NIM（OpenAI 兼容 API）进行笑话提取、改写和分类。"""

import json
import os
from openai import OpenAI


class AiService:
    def __init__(self, api_base: str, api_key: str, model_name: str, temperature: float = 0.7, max_tokens: int = 2048):
        self.client = OpenAI(base_url=api_base, api_key=api_key)
        self.model = model_name
        self.temperature = temperature
        self.max_tokens = max_tokens

    def extract_jokes(self, page_content: str, known_types: list[str], known_crowds: list[str]) -> list[dict]:
        """从页面内容中提取笑话，返回结构化数据。"""
        from crawler.prompts import EXTRACTION_SYSTEM_PROMPT, EXTRACTION_USER_PROMPT

        user_prompt = EXTRACTION_USER_PROMPT.format(
            page_content=page_content[:8000],
            known_types=", ".join(known_types),
            known_crowds=", ".join(known_crowds),
        )

        response = self.client.chat.completions.create(
            model=self.model,
            messages=[
                {"role": "system", "content": EXTRACTION_SYSTEM_PROMPT},
                {"role": "user", "content": user_prompt},
            ],
            temperature=self.temperature,
            max_tokens=self.max_tokens,
        )

        raw = response.choices[0].message.content
        return self._parse_json(raw)

    def rewrite_joke(self, content: str) -> str:
        """润色单条笑话内容。"""
        from crawler.prompts import REWRITE_SYSTEM_PROMPT, REWRITE_USER_PROMPT

        response = self.client.chat.completions.create(
            model=self.model,
            messages=[
                {"role": "system", "content": REWRITE_SYSTEM_PROMPT},
                {"role": "user", "content": REWRITE_USER_PROMPT.format(content=content)},
            ],
            temperature=0.8,
            max_tokens=500,
        )
        return response.choices[0].message.content.strip()

    def _parse_json(self, raw: str) -> list[dict]:
        """安全解析 LLM 返回的 JSON。"""
        try:
            data = json.loads(raw)
            if isinstance(data, dict):
                jokes = data.get("jokes") or data.get("items") or [data]
            else:
                jokes = data
            return [j for j in jokes if isinstance(j, dict) and j.get("content")]
        except json.JSONDecodeError:
            # 尝试提取 markdown 代码块
            if "```json" in raw:
                raw = raw.split("```json")[1].split("```")[0]
            elif "```" in raw:
                raw = raw.split("```")[1].split("```")[0]
            try:
                return json.loads(raw.strip())
            except Exception:
                return []