#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""Small OpenAI-compatible chat helper used by query data scripts."""

from __future__ import annotations

import http.client
import json
import os
import random
import time
from typing import Any, Dict, List
from urllib.parse import urlparse


DEFAULT_BASE_URL = "https://dashscope.aliyuncs.com/compatible-mode/v1"


def chat_completion(
    messages: List[Dict[str, str]],
    model: str,
    base_url: str = DEFAULT_BASE_URL,
    api_key_env: str = "DASHSCOPE_API_KEY",
    temperature: float = 0.2,
    enable_thinking: bool = True,
    stream: bool = True,
    timeout: int = 180,
) -> str:
    api_key = os.environ.get(api_key_env) or os.environ.get("OPENAI_API_KEY")
    if not api_key:
        raise RuntimeError(f"Missing API key. Set ${api_key_env} or $OPENAI_API_KEY.")

    max_retries = max(1, int(os.environ.get("LLM_API_MAX_RETRIES", "4")))
    base_sleep = max(0.0, float(os.environ.get("LLM_API_RETRY_BASE_SEC", "2.0")))
    stream_fallback = os.environ.get("LLM_API_STREAM_FALLBACK", "1").lower() not in {"0", "false", "no"}
    force_non_stream = os.environ.get("LLM_API_DISABLE_STREAM", "0").lower() in {"1", "true", "yes"}
    requested_stream = bool(stream and not force_non_stream)

    errors: list[str] = []
    for attempt in range(1, max_retries + 1):
        try:
            return _chat_completion_once(
                messages=messages,
                model=model,
                base_url=base_url,
                api_key=api_key,
                temperature=temperature,
                enable_thinking=enable_thinking,
                stream=requested_stream,
                timeout=timeout,
            )
        except Exception as exc:
            errors.append(f"attempt {attempt} stream={requested_stream}: {exc!r}")
            if requested_stream and stream_fallback:
                try:
                    return _chat_completion_once(
                        messages=messages,
                        model=model,
                        base_url=base_url,
                        api_key=api_key,
                        temperature=temperature,
                        enable_thinking=enable_thinking,
                        stream=False,
                        timeout=timeout,
                    )
                except Exception as fallback_exc:
                    errors.append(f"attempt {attempt} stream=False fallback: {fallback_exc!r}")
            if attempt >= max_retries:
                break
            sleep_sec = base_sleep * (2 ** (attempt - 1)) + random.uniform(0, min(1.0, base_sleep))
            time.sleep(sleep_sec)

    raise RuntimeError("LLM chat completion failed after retries:\n" + "\n".join(errors[-8:]))


def _chat_completion_once(
    messages: List[Dict[str, str]],
    model: str,
    base_url: str,
    api_key: str,
    temperature: float,
    enable_thinking: bool,
    stream: bool,
    timeout: int,
) -> str:
    force_stdlib = os.environ.get("LLM_API_FORCE_STDLIB", "0").lower() in {"1", "true", "yes"}

    try:
        if force_stdlib:
            raise ImportError("LLM_API_FORCE_STDLIB is set")
        from openai import OpenAI
    except Exception:
        return _chat_completion_http(
            messages=messages,
            model=model,
            base_url=base_url,
            api_key=api_key,
            temperature=temperature,
            enable_thinking=enable_thinking,
            stream=stream,
            timeout=timeout,
        )

    try:
        client = OpenAI(api_key=api_key, base_url=base_url, timeout=timeout, max_retries=0)
    except TypeError:
        client = OpenAI(api_key=api_key, base_url=base_url)
    completion = client.chat.completions.create(
        model=model,
        messages=messages,
        temperature=temperature,
        extra_body={"enable_thinking": enable_thinking},
        stream=stream,
    )
    if not stream:
        msg = completion.choices[0].message
        content = getattr(msg, "content", None)
        if content:
            return content.strip()
        # Some OpenAI-compatible providers expose provider-specific payloads.
        try:
            return msg.model_dump_json(ensure_ascii=False)
        except Exception:
            return str(msg)

    parts: List[str] = []
    for chunk in completion:
        if not chunk.choices:
            continue
        delta = chunk.choices[0].delta
        content = getattr(delta, "content", None)
        if content:
                parts.append(content)
    result = "".join(parts).strip()
    if not result:
        raise RuntimeError("Empty streaming content from chat completions.")
    return result


def extract_json_object(text: str) -> Dict[str, Any]:
    cleaned = text.strip()
    cleaned = cleaned.replace("```json", "").replace("```", "")
    try:
        return json.loads(cleaned)
    except json.JSONDecodeError:
        start = cleaned.find("{")
        end = cleaned.rfind("}")
        if start == -1 or end == -1 or end <= start:
            raise
        return json.loads(cleaned[start : end + 1])


def _chat_completion_http(
    messages: List[Dict[str, str]],
    model: str,
    base_url: str,
    api_key: str,
    temperature: float,
    enable_thinking: bool,
    stream: bool,
    timeout: int,
) -> str:
    parsed_url = urlparse(base_url.rstrip("/"))
    base_path = parsed_url.path.rstrip("/")
    if base_path.endswith("/chat/completions"):
        path = base_path
    else:
        path = base_path + "/chat/completions"
    body = json.dumps(
        {
            "model": model,
            "messages": messages,
            "temperature": temperature,
            "enable_thinking": enable_thinking,
            "stream": stream,
        },
        ensure_ascii=False,
    ).encode("utf-8")
    conn = _make_https_connection(parsed_url.netloc, timeout)
    try:
        conn.request(
            "POST",
            path,
            body=body,
            headers={
                "Authorization": f"Bearer {api_key}",
                "Content-Type": "application/json",
                "Accept": "text/event-stream" if stream else "application/json",
            },
        )
        resp = conn.getresponse()
        if resp.status >= 400:
            error_body = resp.read().decode("utf-8", errors="ignore")
            raise RuntimeError(f"HTTP {resp.status} from chat completions: {error_body[:1000]}")
        if not stream:
            raw_body = resp.read().decode("utf-8", errors="ignore")
            obj = json.loads(raw_body)
            content = _extract_message_content(obj)
            if content:
                return content.strip()
            raise RuntimeError(f"Empty message.content from chat completions. Raw response: {raw_body[:2000]}")

        parts: List[str] = []
        pending = b""
        while True:
            data = resp.read(1024)
            if not data:
                break
            pending += data
            while b"\n" in pending:
                line, pending = pending.split(b"\n", 1)
                line_text = line.decode("utf-8", errors="ignore").strip()
                if not line_text.startswith("data:"):
                    continue
                payload = line_text[5:].strip()
                if payload == "[DONE]":
                    return "".join(parts).strip()
                try:
                    obj = json.loads(payload)
                except json.JSONDecodeError:
                    continue
                for choice in obj.get("choices", []):
                    content = (choice.get("delta") or {}).get("content")
                    if content:
                        parts.append(content)
        result = "".join(parts).strip()
        if not result:
            raise RuntimeError("Empty streaming content from fallback HTTP client.")
        return result
    finally:
        conn.close()


def _extract_message_content(obj: Dict[str, Any]) -> str:
    choices = obj.get("choices") or []
    if not choices:
        return ""
    message = choices[0].get("message") or {}
    content = message.get("content")
    if isinstance(content, str):
        return content
    if isinstance(content, list):
        parts = []
        for part in content:
            if isinstance(part, dict):
                text = part.get("text") or part.get("content")
                if text:
                    parts.append(str(text))
            elif isinstance(part, str):
                parts.append(part)
        return "\n".join(parts)
    return ""


def _make_https_connection(target_netloc: str, timeout: int) -> http.client.HTTPSConnection:
    proxy = (
        os.environ.get("HTTPS_PROXY")
        or os.environ.get("https_proxy")
        or os.environ.get("ALL_PROXY")
        or os.environ.get("all_proxy")
    )
    if not proxy:
        return http.client.HTTPSConnection(target_netloc, timeout=timeout)
    proxy_url = proxy if "://" in proxy else f"http://{proxy}"
    parsed_proxy = urlparse(proxy_url)
    if parsed_proxy.scheme not in {"http", "https"} or not parsed_proxy.hostname:
        raise RuntimeError(f"Unsupported proxy URL: {proxy}")
    proxy_port = parsed_proxy.port or (443 if parsed_proxy.scheme == "https" else 8080)
    conn = http.client.HTTPSConnection(parsed_proxy.hostname, proxy_port, timeout=timeout)
    conn.set_tunnel(target_netloc)
    return conn