# -*- coding: utf-8 -*-
"""
统一模型网关：多供应商 OpenAI 兼容协议调用 + 计量 + 计价
"""
import time
import requests
import config


class LLMError(Exception):
    pass


def get_provider_cfg(provider):
    cfg = config.PROVIDERS.get(provider)
    if not cfg:
        raise LLMError(f'未知供应商: {provider}')
    return cfg


def model_price(model):
    p = config.MODEL_PRICING.get(model, config.DEFAULT_PRICE)
    return p['input'], p['output']


def calc_cost(model, prompt_tokens, completion_tokens):
    pin, pout = model_price(model)
    return round(prompt_tokens / 1e6 * pin + completion_tokens / 1e6 * pout, 6)


def chat(provider, model, messages, temperature=0.7, max_tokens=None,
         base_url=None, api_key=None, timeout=None, retries=None):
    """调用 OpenAI 兼容 chat/completions，返回 {text, usage, cost, model}"""
    cfg = get_provider_cfg(provider)
    url = (base_url or cfg['base_url']).rstrip('/') + '/chat/completions'
    key = api_key or cfg['api_key']
    if not key:
        raise LLMError(f'供应商 {provider} 未配置 API Key')
    headers = {
        'Authorization': f'Bearer {key}',
        'Content-Type': 'application/json',
    }
    payload = {
        'model': model,
        'messages': messages,
        'temperature': temperature,
    }
    if max_tokens:
        payload['max_tokens'] = max_tokens

    timeout = timeout or cfg.get('timeout', 300)
    retries = config.MAX_RETRY if retries is None else retries
    last_err = None
    for attempt in range(retries + 1):
        try:
            resp = requests.post(url, json=payload, headers=headers, timeout=timeout)
            if resp.status_code == 200:
                data = resp.json()
                text = data['choices'][0]['message']['content'] or ''
                usage = data.get('usage', {})
                pt = usage.get('prompt_tokens', 0)
                ct = usage.get('completion_tokens', 0)
                return {
                    'text': text,
                    'model': data.get('model', model),
                    'prompt_tokens': pt,
                    'completion_tokens': ct,
                    'total_tokens': pt + ct,
                    'cost': calc_cost(model, pt, ct),
                }
            if resp.status_code == 429:
                last_err = LLMError(f'模型限流(429): {resp.text[:200]}')
                time.sleep(2 * (attempt + 1))
                continue
            if resp.status_code >= 500:
                last_err = LLMError(f'服务端错误({resp.status_code}): {resp.text[:200]}')
                time.sleep(1)
                continue
            raise LLMError(f'调用失败({resp.status_code}): {resp.text[:300]}')
        except requests.exceptions.Timeout:
            last_err = LLMError(f'请求超时({timeout}s)')
        except requests.exceptions.ConnectionError as e:
            last_err = LLMError(f'连接失败: {e}')
    raise last_err or LLMError('未知错误')


def test_connection(provider, model, base_url=None, api_key=None):
    """连通性测试：发一条最小请求"""
    t0 = time.time()
    r = chat(provider, model,
             [{'role': 'user', 'content': '请回复"OK"两个字'}],
             temperature=0, max_tokens=16,
             base_url=base_url, api_key=api_key, timeout=30, retries=0)
    return {'ok': True, 'latency_ms': int((time.time() - t0) * 1000),
            'reply': r['text'][:50], 'cost': r['cost']}