#!/usr/bin/env python3
"""LLM 客户端（OpenAI 兼容接口）"""
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 _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, temperature=None, max_tokens=None, timeout=None):
    """调用 OpenAI 兼容 chat/completions，返回 content 字符串"""
    url = f'{LLM_BASE_URL}/chat/completions'
    body = {
        'model': LLM_MODEL,
        'messages': messages,
        'temperature': temperature if temperature is not None else LLM_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 {LLM_API_KEY}',
        },
        method='POST',
    )
    try:
        with urllib.request.urlopen(req, timeout=timeout or LLM_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:
        return data['choices'][0]['message']['content']
    except (KeyError, IndexError, TypeError):
        raise LLMError(f'LLM 响应异常: {str(data)[:300]}')


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