#!/usr/bin/env python3
"""LLM 客户端（OpenAI 兼容接口）

支持按模型配置调用（base_url / api_key / model / temperature / timeout），
支持多模态消息（content 为 [{type:text},{type:image_url}] 列表，用于视觉模型分析截图）。
"""
import json
import re
import urllib.request
import urllib.error

from config import LLM_BASE_URL, LLM_API_KEY, LLM_MODEL, LLM_TEMPERATURE, LLM_TIMEOUT


class LLMError(Exception):
    pass


def _cfg(cfg):
    """补齐默认值，返回完整配置 dict"""
    base = {
        'base_url': LLM_BASE_URL,
        'api_key': LLM_API_KEY,
        'model': LLM_MODEL,
        'temperature': LLM_TEMPERATURE,
        'timeout': LLM_TIMEOUT,
        'vision': False,
    }
    if cfg:
        base.update({k: v for k, v in cfg.items() if v is not None})
    return base


def _extract_json(text):
    """从 LLM 输出中提取 JSON（容忍 markdown 代码块等包裹）"""
    if not text:
        return None
    text = text.strip()
    # 去掉 markdown 代码块
    fence = re.search(r'```(?:json)?\s*(.*?)```', text, re.S)
    if fence:
        text = fence.group(1).strip()
    try:
        return json.loads(text)
    except json.JSONDecodeError:
        pass
    # 尝试截取第一个 { 到最后一个 }
    s, e = text.find('{'), text.rfind('}')
    if s != -1 and e > s:
        try:
            return json.loads(text[s:e + 1])
        except json.JSONDecodeError:
            return None
    return None


def chat(messages, cfg=None, temperature=None, max_tokens=None, timeout=None):
    """调用 OpenAI 兼容 chat/completions，返回 content 字符串

    messages 中的 content 可以是字符串，也可以是多模态列表：
    [{"type":"text","text":"..."},
     {"type":"image_url","image_url":{"url":"data:image/png;base64,..."}}]
    """
    c = _cfg(cfg)
    url = f"{c['base_url'].rstrip('/')}/chat/completions"
    body = {
        'model': c['model'],
        'messages': messages,
        'temperature': temperature if temperature is not None else c['temperature'],
    }
    if max_tokens:
        body['max_tokens'] = max_tokens
    req = urllib.request.Request(
        url,
        data=json.dumps(body).encode('utf-8'),
        headers={
            'Content-Type': 'application/json',
            'Authorization': f'Bearer {c["api_key"]}',
        },
        method='POST',
    )
    try:
        with urllib.request.urlopen(req, timeout=timeout or c['timeout']) as resp:
            data = json.loads(resp.read().decode('utf-8'))
    except urllib.error.HTTPError as e:
        detail = e.read().decode('utf-8', 'ignore')[:300]
        raise LLMError(f'LLM HTTP {e.code}: {detail}')
    except Exception as e:
        raise LLMError(f'LLM 调用失败: {e}')

    try:
        msg = data['choices'][0]['message']
        content = msg.get('content') or ''
        if not content.strip():
            # thinking 模型（如 qwen3/deepseek-r1 系）：最终答案可能为空，
            # 思考过程放在 reasoning_content 里（常包含最终 JSON），用它兜底
            reasoning = msg.get('reasoning_content') or ''
            if reasoning.strip():
                return reasoning
        return content
    except (KeyError, IndexError, TypeError):
        raise LLMError(f'LLM 响应异常: {str(data)[:300]}')


def _trim_content(text, limit=300):
    """截断过长的模型输出（避免超长思考内容反复进入上下文）"""
    text = (text or '').strip()
    return text if len(text) <= limit else text[:limit] + '...(已截断)'


def chat_json(messages, cfg=None, temperature=None, max_tokens=None, retries=2):
    """调用 LLM 并强制解析 JSON，失败重试"""
    last_err = None
    last_content = ''
    for i in range(retries + 1):
        try:
            last_content = chat(messages, cfg=cfg, temperature=temperature,
                                max_tokens=max_tokens)
            obj = _extract_json(last_content)
            if obj is not None:
                return obj
            last_err = f'无法从输出解析 JSON: {last_content[:200]}'
        except LLMError as e:
            last_err = str(e)
        if i < retries:
            messages = messages + [
                {'role': 'assistant', 'content': _trim_content(last_content)},
                {'role': 'user',
                 'content': f'刚才的输出不是合法 JSON，请只输出严格的 JSON 对象，不要输出思考过程和代码块。错误: {last_err}'},
            ]
    raise LLMError(f'LLM JSON 解析失败: {last_err}')