from __future__ import annotations
import asyncio
import json
import logging
import re
import time
from typing import Any

from openai import AsyncOpenAI

logger = logging.getLogger(__name__)


class LLMClient:
    def __init__(self, api_key: str, base_url: str, model: str, max_concurrent: int = 10):
        self.client = AsyncOpenAI(api_key=api_key, base_url=base_url)
        self.model = model
        self._is_minimax = "minimax" in (base_url or "").lower()
        self.semaphore = asyncio.Semaphore(max_concurrent)
        self.total_tokens = 0
        self.total_calls = 0
        self._lock = asyncio.Lock()
        # Per-purpose call stats — useful when a sim stalls and you need to
        # know which LLM call category is slow. Reset by the engine at the
        # start of each sim if you want per-sim stats.
        self._purpose_stats: dict[str, dict[str, float]] = {}

    async def _track(self, usage):
        if usage:
            async with self._lock:
                self.total_tokens += usage.total_tokens or 0
                self.total_calls += 1

    async def _record_purpose(self, purpose: str | None, elapsed: float, ok: bool):
        if not purpose:
            return
        async with self._lock:
            st = self._purpose_stats.setdefault(
                purpose, {"calls": 0, "errors": 0, "total_ms": 0.0}
            )
            st["calls"] += 1
            st["total_ms"] += elapsed * 1000
            if not ok:
                st["errors"] += 1

    async def generate(
        self,
        system: str,
        user: str,
        json_mode: bool = False,
        max_tokens: int = 4096,
        retries: int = 3,
        purpose: str | None = None,
    ) -> str:
        kwargs: dict = {
            "model": self.model,
            "messages": [
                {"role": "system", "content": system},
                {"role": "user", "content": user},
            ],
            "temperature": 0.5 if json_mode else 1.0,
        }
        if self._is_minimax:
            # MiniMax M3 enables thinking by default.  In structured-output
            # calls that can consume the completion budget before a visible
            # message.content is produced.  MiniMax-specific fields must be
            # sent through extra_body when using the OpenAI Python SDK.
            kwargs["max_completion_tokens"] = max_tokens
            kwargs["extra_body"] = {"thinking": {"type": "disabled"}}
        else:
            kwargs["max_tokens"] = max_tokens

        if json_mode and not self._is_minimax:
            # MiniMax's newer M-series models do not reliably implement the
            # OpenAI json_object response_format.  They still receive the
            # explicit JSON schema in the prompt and are parsed defensively.
            kwargs["response_format"] = {"type": "json_object"}

        t0 = time.monotonic()
        for attempt in range(retries):
            try:
                async with self.semaphore:
                    response = await self.client.chat.completions.create(**kwargs)
                await self._track(response.usage)
                if not response.choices:
                    logger.warning(
                        "LLM 返回空 choices（用途=%s，响应ID=%s，模型=%s）",
                        purpose or "-", getattr(response, "id", "-"), self.model,
                    )
                    content = ""
                    finish = "no_choices"
                    message = None
                else:
                    choice = response.choices[0]
                    message = choice.message
                    content = getattr(message, "content", None) or ""
                    finish = getattr(choice, "finish_reason", "unknown")

                    # A few OpenAI-compatible providers return content parts
                    # instead of a plain string. Normalize text parts here.
                    if isinstance(content, list):
                        content = "".join(
                            part.get("text", "") if isinstance(part, dict)
                            else getattr(part, "text", "") or ""
                            for part in content
                        )

                if not content:
                    reasoning = getattr(message, "reasoning_content", None) if message else None
                    reasoning_details = getattr(message, "reasoning_details", None) if message else None
                    refusal = getattr(message, "refusal", None) if message else None
                    usage = getattr(response, "usage", None)
                    logger.warning(
                        "LLM 返回空内容（用途=%s，结束原因=%s，完成令牌=%s，推理内容长度=%s，推理详情=%s，拒答=%s）",
                        purpose or "-",
                        finish,
                        getattr(usage, "completion_tokens", "?"),
                        len(reasoning or ""),
                        bool(reasoning_details),
                        bool(refusal),
                    )
                    if attempt < retries - 1:
                        wait = 2 ** attempt
                        logger.warning(
                            "LLM 返回空内容，将在%d秒后重试（用途=%s，第%d/%d次）",
                            wait, purpose or "-", attempt + 1, retries,
                        )
                        await asyncio.sleep(wait)
                        continue
                    await self._record_purpose(purpose, time.monotonic() - t0, ok=False)
                    return ""
                await self._record_purpose(purpose, time.monotonic() - t0, ok=True)
                if purpose:
                    logger.debug(
                        "LLM 调用成功（用途=%s，令牌数=%s，耗时=%.2f秒）",
                        purpose,
                        getattr(response.usage, "total_tokens", "?"),
                        time.monotonic() - t0,
                    )
                return content.strip()
            except Exception as e:
                if attempt == retries - 1:
                    await self._record_purpose(purpose, time.monotonic() - t0, ok=False)
                    logger.error(
                        "LLM 调用重试 %d 次后仍失败（用途=%s，耗时=%.2f秒）：%s",
                        retries, purpose or "-", time.monotonic() - t0, e,
                    )
                    raise
                wait = 2 ** attempt
                logger.warning(
                    "LLM 调用失败（用途=%s，第%d次尝试，%d秒后重试，已耗时=%.2f秒）：%s",
                    purpose or "-", attempt + 1, wait, time.monotonic() - t0, e,
                )
                await asyncio.sleep(wait)
        return ""

    async def generate_batch(
        self,
        prompts: list[tuple[str, str]],
        json_mode: bool = False,
        max_tokens: int = 1000,
        purpose: str | None = None,
        sub_purpose: str | None = None,
    ) -> list[str]:
        # `purpose` is the batch label (e.g. "decision"); `sub_purpose` lets
        # us split stats by an inner dimension (e.g. "active/react").
        tasks = [
            self.generate(
                system, user,
                json_mode=json_mode, max_tokens=max_tokens,
                purpose=sub_purpose or purpose,
            )
            for system, user in prompts
        ]
        results = await asyncio.gather(*tasks, return_exceptions=True)
        return [r if isinstance(r, str) else "" for r in results]

    async def generate_light(self, system: str, user: str, max_tokens: int = 100) -> str:
        return await self.generate(system, user, json_mode=False, max_tokens=max_tokens)

    def get_stats(self) -> dict:
        return {
            "total_tokens": self.total_tokens,
            "total_calls": self.total_calls,
            "by_purpose": dict(self._purpose_stats),
        }


def parse_json(text: str, purpose: str | None = None) -> Any:
    """Parse a model response that was requested as JSON.

    Providers occasionally return prose, truncated JSON, or a markdown code
    block even when ``response_format=json_object`` was requested.  Keep the
    permissive recovery logic, but make a final failure actionable: include
    the call purpose, response length, and a bounded raw preview in the
    simulation log.
    """
    text = (text or "").strip()
    label = purpose or "未标注用途"

    if not text:
        logger.error("JSON 解析失败（用途=%s）：LLM 返回内容为空", label)
        return {}

    # Strip modern LLM "thinking" blocks (Anthropic, DeepSeek, etc.) that may
    # appear either before or around the JSON payload.
    text = re.sub(r"<think>.*?</think>", "", text, flags=re.DOTALL)

    # Markdown code fence: ```json ... ``` or ``` ... ```
    if text.startswith("```"):
        match = re.search(r"```(?:json)?\s*\n?(.*?)```", text, re.DOTALL)
        if match:
            text = match.group(1).strip()

    # Try the whole text first.
    try:
        return json.loads(text)
    except json.JSONDecodeError:
        pass

    # Otherwise find the first balanced { ... } or [ ... ] substring that parses.
    def _balanced_at(start: int) -> int | None:
        open_ch, other_ch = ("{", "}") if text[start] == "{" else (("[", "]") if text[start] == "[" else (None, None))
        if open_ch is None:
            return None
        depth = 0
        in_str = False
        esc = False
        for j in range(start, len(text)):
            c = text[j]
            if in_str:
                if esc:
                    esc = False
                elif c == "\\":
                    esc = True
                elif c == '"':
                    in_str = False
                continue
            if c == '"':
                in_str = True
            elif c == open_ch:
                depth += 1
            elif c == other_ch:
                depth -= 1
                if depth == 0:
                    return j
        return None

    for opener in ("{", "["):
        i = 0
        while True:
            idx = text.find(opener, i)
            if idx == -1:
                break
            end = _balanced_at(idx)
            if end is not None:
                try:
                    return json.loads(text[idx : end + 1])
                except json.JSONDecodeError:
                    pass
            i = idx + 1

    # Last-ditch: trim from the end while we can find a parseable prefix.
    for i in range(len(text) - 1, -1, -1):
        if text[i] in "]}":
            try:
                return json.loads(text[: i + 1])
            except json.JSONDecodeError:
                continue

    open_braces = text.count("{") - text.count("}")
    open_brackets = text.count("[") - text.count("]")
    repaired = text + ("}" * max(0, open_braces)) + ("]" * max(0, open_brackets))
    try:
        return json.loads(repaired)
    except json.JSONDecodeError:
        pass

    preview = text[:1000].replace("\r", "\\r").replace("\n", "\\n")
    logger.error(
        "JSON 解析失败（用途=%s，返回长度=%d）：内容不是有效 JSON 或已被截断；原始返回片段=%s%s",
        label,
        len(text),
        preview,
        "……" if len(text) > 1000 else "",
    )
    return {}