from __future__ import annotations

import json
import os
import threading
from typing import Any, AsyncGenerator, Dict, Optional

_LITELLM_ACOMPLETION = None
_LITELLM_PREWARM_STARTED = False
_LITELLM_PREWARM_LOCK = threading.Lock()
_LITELLM_PREWARM_DONE = threading.Event()
_LITELLM_PREWARM_ERROR: str | None = None


def _obj_get(obj: Any, key: str, default: Any = None) -> Any:
    if isinstance(obj, dict):
        return obj.get(key, default)
    return getattr(obj, key, default)


def _get_litellm_acompletion():
    global _LITELLM_ACOMPLETION, _LITELLM_PREWARM_ERROR
    if _LITELLM_ACOMPLETION is not None:
        return _LITELLM_ACOMPLETION
    os.environ.setdefault("LITELLM_LOCAL_MODEL_COST_MAP", "true")
    try:
        import litellm as _litellm
        from litellm.main import acompletion as _acompletion

        _litellm.suppress_debug_info = True
    except Exception as exc:
        _LITELLM_PREWARM_ERROR = f"{type(exc).__name__}: {exc}"
        raise RuntimeError(
            "litellm is required for model calls but is not available. Install with: uv pip install litellm"
        ) from exc
    _LITELLM_ACOMPLETION = _acompletion
    _LITELLM_PREWARM_ERROR = None
    _LITELLM_PREWARM_DONE.set()
    return _LITELLM_ACOMPLETION


def start_global_litellm_prewarm(*, blocking: bool = False, timeout: float | None = None) -> bool:
    global _LITELLM_PREWARM_STARTED
    if _LITELLM_ACOMPLETION is not None:
        return True
    with _LITELLM_PREWARM_LOCK:
        if not _LITELLM_PREWARM_STARTED:
            _LITELLM_PREWARM_STARTED = True

            def _runner() -> None:
                try:
                    _get_litellm_acompletion()
                except Exception:
                    return
                finally:
                    _LITELLM_PREWARM_DONE.set()

            t = threading.Thread(target=_runner, name="oturn-litellm-prewarm", daemon=True)
            t.start()
    if blocking:
        _LITELLM_PREWARM_DONE.wait(timeout=timeout)
    return _LITELLM_ACOMPLETION is not None


def get_global_litellm_prewarm_error() -> str | None:
    return _LITELLM_PREWARM_ERROR


class LiteLLMTransport:
    def __init__(
        self,
        *,
        model: str,
        timeout: float = 60.0,
        max_tokens: Optional[int] = None,
        temperature: Optional[float] = None,
        extra_body: Optional[dict[str, Any]] = None,
        api_base: str = "",
        api_key: Optional[str] = None,
    ) -> None:
        self.model = model
        self.timeout = timeout
        self.max_tokens = max_tokens
        self.temperature = temperature
        self.extra_body = extra_body
        self.api_base = api_base
        self.api_key = api_key

    async def acreate(
        self,
        api_messages: list[dict[str, Any]],
        *,
        stream: bool = False,
        tools: list | None = None,
        tool_choice: str | dict | None = None,
        tool_stream: bool = False,
        debug: bool = False,
        sse_debug: bool = False,
        **extra: Any,
    ) -> dict | AsyncGenerator[object, None]:
        params: Dict[str, Any] = {
            "model": self.model,
            "messages": api_messages,
            "timeout": self.timeout,
            "drop_params": True,
            **({"tools": tools} if tools is not None else {}),
            **({"tool_choice": tool_choice} if tool_choice is not None else {}),
            **extra,
        }
        if self.max_tokens is not None:
            params["max_tokens"] = int(self.max_tokens)
        if self.temperature is not None:
            params["temperature"] = float(self.temperature)
        if self.extra_body:
            params.update(self.extra_body)
        if self.api_base:
            params["api_base"] = self.api_base
        if self.api_key:
            params["api_key"] = self.api_key

        if debug:
            import sys as _sys

            print("[debug-litellm] acompletion params:", file=_sys.stderr)
            print(json.dumps(params, ensure_ascii=False, indent=2), file=_sys.stderr)

        acompletion = _get_litellm_acompletion()

        if not stream:
            resp = await acompletion(stream=False, **params)
            if hasattr(resp, "model_dump"):
                return resp.model_dump(mode="json")
            if isinstance(resp, dict):
                return resp
            return {"raw": str(resp)}

        async def _agen() -> AsyncGenerator[object, None]:
            stream_resp = await acompletion(
                stream=True,
                stream_options={"include_usage": True},
                **params,
            )
            usage_obj = None
            resp_id = None
            model_name_from_resp = None

            async for chunk in stream_resp:
                if sse_debug:
                    try:
                        import sys as _sys

                        body = chunk.model_dump(mode="json") if hasattr(chunk, "model_dump") else chunk
                        print("[sse-debug] chunk:", json.dumps(body, ensure_ascii=False), file=_sys.stderr)
                    except Exception:
                        pass

                try:
                    chunk_id = _obj_get(chunk, "id")
                    chunk_model = _obj_get(chunk, "model")
                    if chunk_id and resp_id is None:
                        resp_id = chunk_id
                    if chunk_model and model_name_from_resp is None:
                        model_name_from_resp = chunk_model

                    choices = _obj_get(chunk, "choices", []) or []
                    for i, choice in enumerate(choices):
                        finish_reason = _obj_get(choice, "finish_reason")
                        delta = _obj_get(choice, "delta")
                        if not delta:
                            if finish_reason == "tool_calls" and tool_stream:
                                yield {"type": "tool_calls_done"}
                            continue

                        raw_delta = delta.model_dump(exclude_none=True) if hasattr(delta, "model_dump") else delta

                        tc_list = _obj_get(delta, "tool_calls")
                        if tool_stream and tc_list:
                            for j, tc in enumerate(tc_list):
                                tc_index = _obj_get(tc, "index")
                                ev = {
                                    "type": "tool_call_delta",
                                    "output_index": tc_index if isinstance(tc_index, int) else j,
                                    "model": model_name_from_resp,
                                }
                                tc_id = _obj_get(tc, "id")
                                if tc_id:
                                    ev["tool_call_id"] = tc_id
                                func = _obj_get(tc, "function")
                                if func is not None:
                                    name = _obj_get(func, "name")
                                    if name:
                                        ev["name"] = name
                                    arguments = _obj_get(func, "arguments")
                                    if arguments:
                                        ev["arguments"] = arguments
                                yield ev
                            if finish_reason == "tool_calls":
                                yield {"type": "tool_calls_done"}
                            continue

                        if isinstance(raw_delta, dict):
                            rc = raw_delta.get("reasoning_content") or raw_delta.get("reasoning")
                            if rc:
                                yield {"type": "assistant_reasoning_delta", "text": str(rc), "model": model_name_from_resp}

                        content = _obj_get(delta, "content")
                        if content:
                            yield {"type": "assistant_delta", "text": str(content), "model": model_name_from_resp}

                        if finish_reason == "tool_calls" and tool_stream:
                            yield {"type": "tool_calls_done"}

                    u = _obj_get(chunk, "usage")
                    if u is not None:
                        if hasattr(u, "model_dump"):
                            usage_obj = u.model_dump(exclude_none=True)
                        elif isinstance(u, dict):
                            usage_obj = u
                except Exception:
                    continue

            if usage_obj is not None or resp_id is not None or model_name_from_resp is not None:
                yield {"type": "assistant_done", "usage": usage_obj, "id": resp_id, "model": model_name_from_resp}

        return _agen()
