# -*- coding: utf-8 -*-
"""
LLM 客户端（OpenAI 兼容 /chat/completions）。
- base_url / model / key 由环境变量配置（默认智谱 GLM）
- 出站安全：仅 http/https、显式拒绝 localhost、解析 IP 拒绝环回/私有/保留段、
  禁用重定向（防 DNS rebinding 绕过），值全部走 JSON 序列化，密钥不落日志
"""
import ipaddress
import json
import os
from urllib.parse import urlparse

import requests


class LlmError(Exception):
    pass


def _validated_url(base_url: str) -> str:
    u = urlparse(base_url)
    if u.scheme not in ('http', 'https'):
        raise LlmError('LLM base_url 仅允许 http/https')
    host = u.hostname or ''
    if not host or host.lower() in ('localhost', 'localhost.localdomain'):
        raise LlmError('LLM base_url 拒绝 localhost')
    port = u.port or (443 if u.scheme == 'https' else 80)
    try:
        infos = socket.getaddrinfo(host, port)
    except socket.gaierror as e:
        raise LlmError('LLM base_url 域名解析失败: {}'.format(host))
    for info in infos:
        ip = ipaddress.ip_address(info[4][0])
        if (ip.is_loopback or ip.is_private or ip.is_link_local or ip.is_reserved
                or ip.is_multicast or ip.is_unspecified):
            raise LlmError('LLM base_url 拒绝非公网地址: {}'.format(ip))
    return '{}://{}{}'.format(u.scheme, u.netloc, u.path)


class LlmClient:

    def __init__(self):
        self.api_key = os.environ.get('JQUANT_LLM_API_KEY', '').strip()
        self.base_url = (os.environ.get('JQUANT_LLM_BASE_URL', '').strip()
                         or 'https://open.bigmodel.cn/api/paas/v4')
        self.model = os.environ.get('JQUANT_LLM_MODEL', '').strip() or 'glm-4-flash'
        self.temperature = float(os.environ.get('JQUANT_LLM_TEMPERATURE', '0.2'))

    @property
    def enabled(self) -> bool:
        return bool(self.api_key)

    def chat(self, system_prompt: str, user_prompt: str) -> str:
        if not self.enabled:
            raise LlmError('未配置 LLM API Key（JQUANT_LLM_API_KEY），请在服务环境变量中设置')
        url = _validated_url(self.base_url.rstrip('/')) + '/chat/completions'
        body = {
            'model': self.model,
            'temperature': self.temperature,
            'messages': [
                {'role': 'system', 'content': system_prompt},
                {'role': 'user', 'content': user_prompt},
            ],
        }
        # 校验与请求紧邻；禁重定向防 DNS rebinding 绕过 IP 校验
        r = requests.post(url, json=body, timeout=90, allow_redirects=False,
                          headers={'Authorization': 'Bearer ' + self.api_key,
                                   'Content-Type': 'application/json'})
        if r.status_code in (301, 302, 303, 307, 308):
            raise LlmError('LLM 端点发生重定向，已拒绝（防 SSRF 绕过）')
        if r.status_code != 200:
            raise LlmError('LLM HTTP {}: {}'.format(r.status_code, r.text[:200]))
        content = r.json().get('choices', [{}])[0].get('message', {}).get('content')
        if not content:
            raise LlmError('LLM 响应缺少 content')
        return content