"""
 - 

 KlynxAgent  .invoke(messages) .
 LiteLLM , (OpenAI, Anthropic, Google, DeepSeek ).
 httpx.Client .
"""

import logging
import os
import random
import time
from typing import Any, List, Optional

import httpx
os.environ.setdefault("LITELLM_LOG", "WARNING")
import litellm
from langchain_core.messages import AIMessage, BaseMessage, HumanMessage, SystemMessage

#  SSL 
litellm.suppress_debug_info = True
for _litellm_logger_name in ("LiteLLM", "LiteLLM Router", "LiteLLM Proxy"):
    logging.getLogger(_litellm_logger_name).setLevel(logging.WARNING)
logger = logging.getLogger(__name__)


def _coerce_usage_int(value: Any) -> int:
    try:
        return max(int(value), 0)
    except Exception:
        return 0


def normalize_usage_payload(payload: Any) -> dict:
    """Normalize provider usage payloads into prompt/completion/total token counts."""
    if payload is None:
        return {
            "prompt_tokens": 0,
            "completion_tokens": 0,
            "total_tokens": 0,
        }

    if isinstance(payload, dict):
        for key in ("usage", "usage_metadata", "token_usage"):
            nested = payload.get(key)
            if nested is not None and nested is not payload:
                normalized = normalize_usage_payload(nested)
                if normalized["total_tokens"] > 0:
                    return normalized

        response_metadata = payload.get("response_metadata")
        if response_metadata is not None and response_metadata is not payload:
            normalized = normalize_usage_payload(response_metadata)
            if normalized["total_tokens"] > 0:
                return normalized

        prompt_tokens = _coerce_usage_int(
            payload.get("prompt_tokens", payload.get("input_tokens", 0))
        )
        completion_tokens = _coerce_usage_int(
            payload.get("completion_tokens", payload.get("output_tokens", 0))
        )
        total_tokens = _coerce_usage_int(payload.get("total_tokens", 0))
    else:
        for attr_name in ("usage", "usage_metadata", "token_usage"):
            nested = getattr(payload, attr_name, None)
            if nested is not None and nested is not payload:
                normalized = normalize_usage_payload(nested)
                if normalized["total_tokens"] > 0:
                    return normalized

        response_metadata = getattr(payload, "response_metadata", None)
        if response_metadata is not None and response_metadata is not payload:
            normalized = normalize_usage_payload(response_metadata)
            if normalized["total_tokens"] > 0:
                return normalized

        prompt_tokens = _coerce_usage_int(
            getattr(payload, "prompt_tokens", getattr(payload, "input_tokens", 0))
        )
        completion_tokens = _coerce_usage_int(
            getattr(payload, "completion_tokens", getattr(payload, "output_tokens", 0))
        )
        total_tokens = _coerce_usage_int(getattr(payload, "total_tokens", 0))

    if total_tokens <= 0 and (prompt_tokens > 0 or completion_tokens > 0):
        total_tokens = prompt_tokens + completion_tokens
    if prompt_tokens <= 0 and total_tokens > 0 and completion_tokens > 0:
        prompt_tokens = max(total_tokens - completion_tokens, 0)
    if completion_tokens <= 0 and total_tokens > 0 and prompt_tokens > 0:
        completion_tokens = max(total_tokens - prompt_tokens, 0)

    return {
        "prompt_tokens": prompt_tokens,
        "completion_tokens": completion_tokens,
        "total_tokens": total_tokens,
    }


class _InferenceBackend:
    """."""

    def completion(self, call_kwargs: dict):
        raise NotImplementedError

    def stream_completion(self, call_kwargs: dict):
        raise NotImplementedError


class _LiteLLMBackend(_InferenceBackend):
    """:LiteLLM completion."""

    def completion(self, call_kwargs: dict):
        return litellm.completion(**call_kwargs)

    def stream_completion(self, call_kwargs: dict):
        return litellm.completion(**call_kwargs)


class _ResponsesCompatBackend(_InferenceBackend):
    """
    Responses API .

    , LiteLLM completion.
     Responses API .
    """

    def completion(self, call_kwargs: dict):
        return litellm.completion(**call_kwargs)

    def stream_completion(self, call_kwargs: dict):
        return litellm.completion(**call_kwargs)


class LiteLLMResponse:
    """
    
     graph.py  response.content / response.reasoning_content / response.usage 
     Tool Calling( Function Calling) tool_calls
    """
    def __init__(self, content: str, reasoning_content: str = ""):
        self.content = content
        self.reasoning_content = reasoning_content
        self.additional_kwargs = {}
        self.response_metadata = {}
        self.usage = None  # 
        self.tool_calls = []  #  Tool Calling 


class LiteLLMChat:
    """
    
    
     LiteLLM ,
     .invoke(messages)  .stream(messages) .
     Tool Calling( Function Calling,tools ).
    """
    
    def __init__(self, model: str, api_key: str = None, api_base: str = None,
                 temperature: float = 0.1, backend: str = "litellm", **kwargs):
        """
        Args:
            model:  ( "deepseek/deepseek-reasoner", "gpt-4o")
                    litellm provider/model 
            api_key: API Key
            api_base:  API  ()
            temperature: 
            **kwargs: 
        """
        self.model = model
        self.api_key = api_key
        self.api_base = api_base
        self.temperature = temperature
        self.max_context_tokens = 128000
        self.backend_name = str(backend or "litellm").strip().lower()
        self.backend = self._create_backend(self.backend_name)

        # ()
        self.retry_enabled = self._parse_bool_option(
            kwargs.pop("retry_enabled", os.getenv("KLYNX_MODEL_RETRY_ENABLED", "1")),
            default=True,
        )
        self.retry_max_attempts = self._parse_int_option(
            kwargs.pop("retry_max_attempts", os.getenv("KLYNX_MODEL_RETRY_MAX_ATTEMPTS", "3")),
            default=3,
            minimum=1,
        )
        self.retry_initial_delay = self._parse_float_option(
            kwargs.pop("retry_initial_delay", os.getenv("KLYNX_MODEL_RETRY_INITIAL_DELAY", "1.0")),
            default=1.0,
            minimum=0.0,
        )
        self.retry_backoff_factor = self._parse_float_option(
            kwargs.pop("retry_backoff_factor", os.getenv("KLYNX_MODEL_RETRY_BACKOFF_FACTOR", "2.0")),
            default=2.0,
            minimum=1.0,
        )
        self.retry_max_delay = self._parse_float_option(
            kwargs.pop("retry_max_delay", os.getenv("KLYNX_MODEL_RETRY_MAX_DELAY", "8.0")),
            default=8.0,
            minimum=0.0,
        )
        self.retry_jitter = self._parse_float_option(
            kwargs.pop("retry_jitter", os.getenv("KLYNX_MODEL_RETRY_JITTER", "0.2")),
            default=0.2,
            minimum=0.0,
        )
        fallback_models_raw = kwargs.pop("fallback_models", os.getenv("KLYNX_FALLBACK_MODELS", ""))
        self.fallback_models = self._parse_fallback_models(fallback_models_raw)

        self.extra_kwargs = kwargs

        # (, SSL )
        self.http_client = httpx.Client(trust_env=False)
    
    def invoke(self, messages, tools: Optional[list] = None):
        """
        , LiteLLMResponse
        
        Args:
            messages: LangChain BaseMessage   OpenAI 
            tools: ,OpenAI  tools JSON Schema ( Tool Calling)
            
        Returns:
            LiteLLMResponse 
        """
        # 
        openai_messages = self._sanitize_jsonable(self._convert_messages(messages))
        
        # 
        call_kwargs = {
            "model": self.model,
            "messages": openai_messages,
            "api_key": self.api_key,
        }
        
        if self.api_base:
            call_kwargs["api_base"] = self.api_base
            
        #  temperature 
        if not any(k in self.model for k in ["reasoner", "o1", "kimi-k2.5"]):
            call_kwargs["temperature"] = self.temperature
        
        #  Tool Calling:  tools
        if tools:
            call_kwargs["tools"] = self._sanitize_jsonable(tools)
            call_kwargs["tool_choice"] = "auto"
        
        # 
        call_kwargs.update(self.extra_kwargs)
        call_kwargs = self._sanitize_jsonable(call_kwargs)
        
        #  API()
        response = self._completion_with_retry(call_kwargs)
        
        # 
        message = response.choices[0].message
        content = message.content or ""
        #  reasoning_content 
        reasoning_content = getattr(message, 'reasoning_content', '') or ""
        
        # 
        resp = LiteLLMResponse(content=content, reasoning_content=reasoning_content)
        
        #  tool_calls()
        if hasattr(message, 'tool_calls') and message.tool_calls:
            resp.tool_calls = self._parse_native_tool_calls(message.tool_calls)
        
        #  usage
        usage = normalize_usage_payload(getattr(response, "usage", None) or response)
        if usage["total_tokens"] > 0:
            resp.usage = usage
        
        return resp
    
    def stream(self, messages, tools: Optional[list] = None):
        """
        
        
        Args:
            messages: LangChain 
            tools: ,OpenAI  tools JSON Schema ( Tool Calling)
            
        Yields:
            chunks with 'content', 'reasoning_content', 'usage', and 'tool_calls'
        """
        openai_messages = self._sanitize_jsonable(self._convert_messages(messages))
        
        call_kwargs = {
            "model": self.model,
            "messages": openai_messages,
            "stream": True,
            "stream_options": {"include_usage": True},
            "api_key": self.api_key,
        }
        
        if self.api_base:
            call_kwargs["api_base"] = self.api_base
            
        if not any(k in self.model for k in ["reasoner", "o1", "kimi-k2.5"]):
            call_kwargs["temperature"] = self.temperature
        
        #  Tool Calling:  tools
        if tools:
            call_kwargs["tools"] = self._sanitize_jsonable(tools)
            call_kwargs["tool_choice"] = "auto"
            
        call_kwargs.update(self.extra_kwargs)
        call_kwargs = self._sanitize_jsonable(call_kwargs)
            
        response = self._stream_with_retry(call_kwargs)

        #  tool_calls 
        tool_call_accumulators = {}

        for chunk in response:
            #  Usage  ( chunk)
            usage = normalize_usage_payload(getattr(chunk, "usage", None) or chunk)
            if usage["total_tokens"] > 0:
                yield {
                    "usage": usage
                }

            if not chunk.choices:
                continue

            delta = chunk.choices[0].delta
            content = delta.content or ""
            reasoning_content = getattr(delta, 'reasoning_content', '') or ""

            #  tool_calls 
            if hasattr(delta, 'tool_calls') and delta.tool_calls:
                for tc_delta in delta.tool_calls:
                    idx = tc_delta.index
                    if idx not in tool_call_accumulators:
                        tool_call_accumulators[idx] = {
                            "id": "",
                            "function": {"name": "", "arguments": ""}
                        }
                    acc = tool_call_accumulators[idx]
                    if tc_delta.id:
                        acc["id"] = tc_delta.id
                    if tc_delta.function:
                        if tc_delta.function.name:
                            acc["function"]["name"] += tc_delta.function.name
                        if tc_delta.function.arguments:
                            acc["function"]["arguments"] += tc_delta.function.arguments

            if content or reasoning_content or (hasattr(delta, 'tool_calls') and delta.tool_calls):
                yield {
                    "content": content,
                    "reasoning_content": reasoning_content
                }

        # , tool_calls, yield 
        if tool_call_accumulators:
            assembled = []
            for idx in sorted(tool_call_accumulators.keys()):
                acc = tool_call_accumulators[idx]
                assembled.append(self._parse_single_tool_call(acc))
            yield {"tool_calls": assembled}
    
    def _parse_native_tool_calls(self, raw_tool_calls) -> list:
        """
         LiteLLM  tool_calls 
        
        Returns:
            [{"tool": "read_file", "params": {"path": "..."}}, ...]
        """
        import json
        results = []
        for tc in raw_tool_calls:
            func = tc.function
            name = func.name
            try:
                args = json.loads(func.arguments) if isinstance(func.arguments, str) else func.arguments
            except json.JSONDecodeError:
                args = {"raw_arguments": func.arguments}
            results.append({
                "tool": name,
                "params": args
            })
        return results

    def _create_backend(self, backend: str) -> _InferenceBackend:
        if backend == "responses":
            return _ResponsesCompatBackend()
        return _LiteLLMBackend()
    
    def _parse_single_tool_call(self, acc: dict) -> dict:
        """
         tool_call 
        """
        import json
        name = acc["function"]["name"]
        raw_args = acc["function"]["arguments"]
        try:
            args = json.loads(raw_args) if raw_args else {}
        except json.JSONDecodeError:
            args = {"raw_arguments": raw_args}
        return {
            "tool": name,
            "params": args
        }
    
    def _convert_messages(self, messages):
        """ LangChain  OpenAI """
        openai_messages = []
        for msg in messages:
            if isinstance(msg, dict):
                openai_messages.append(msg)
            elif isinstance(msg, SystemMessage):
                openai_messages.append({"role": "system", "content": msg.content})
            elif isinstance(msg, HumanMessage):
                openai_messages.append({"role": "user", "content": msg.content})
            elif isinstance(msg, AIMessage):
                openai_messages.append({"role": "assistant", "content": msg.content})
            else:
                role = getattr(msg, 'type', 'user')
                openai_messages.append({"role": role, "content": str(msg.content)})
        return openai_messages

    def _sanitize_jsonable(self, value: Any) -> Any:
        """ JSON  Unicode ,."""
        if isinstance(value, str):
            #  read_tui  lone surrogate(\ud800-\udfff)
            # /JSON  UTF-8 , U+FFFD.
            return value.encode("utf-8", errors="replace").decode("utf-8")
        if isinstance(value, list):
            return [self._sanitize_jsonable(item) for item in value]
        if isinstance(value, tuple):
            return tuple(self._sanitize_jsonable(item) for item in value)
        if isinstance(value, dict):
            return {
                self._sanitize_jsonable(key): self._sanitize_jsonable(item)
                for key, item in value.items()
            }
        return value

    def _completion_with_retry(self, call_kwargs: dict):
        model_candidates = self._get_model_candidates()
        for model_idx, model_name in enumerate(model_candidates):
            model_kwargs = dict(call_kwargs)
            model_kwargs["model"] = model_name
            for attempt in range(1, self.retry_max_attempts + 1):
                try:
                    return self.backend.completion(model_kwargs)
                except Exception as exc:
                    if (not self.retry_enabled) or (not self._is_retriable_error(exc)):
                        raise

                    has_next_attempt = attempt < self.retry_max_attempts
                    has_next_model = model_idx < len(model_candidates) - 1
                    if not has_next_attempt and not has_next_model:
                        raise

                    if has_next_attempt:
                        delay = self._retry_delay_seconds(attempt)
                        logger.warning(
                            "Model completion failed, retrying in %.2fs (model=%s, attempt=%s/%s): %s",
                            delay,
                            model_name,
                            attempt,
                            self.retry_max_attempts,
                            exc,
                        )
                        time.sleep(delay)
                    else:
                        logger.warning(
                            "Model completion failed after %s attempts, switching to fallback model: %s -> %s; error=%s",
                            self.retry_max_attempts,
                            model_name,
                            model_candidates[model_idx + 1],
                            exc,
                        )
                    continue

        raise RuntimeError("Model completion failed unexpectedly without captured exception")

    def _stream_with_retry(self, call_kwargs: dict):
        model_candidates = self._get_model_candidates()
        for model_idx, model_name in enumerate(model_candidates):
            model_kwargs = dict(call_kwargs)
            model_kwargs["model"] = model_name
            for attempt in range(1, self.retry_max_attempts + 1):
                emitted_any = False
                try:
                    response = self.backend.stream_completion(model_kwargs)
                    for chunk in response:
                        emitted_any = True
                        yield chunk
                    return
                except Exception as exc:
                    # ,
                    if emitted_any:
                        raise
                    if (not self.retry_enabled) or (not self._is_retriable_error(exc)):
                        raise

                    has_next_attempt = attempt < self.retry_max_attempts
                    has_next_model = model_idx < len(model_candidates) - 1
                    if not has_next_attempt and not has_next_model:
                        raise

                    if has_next_attempt:
                        delay = self._retry_delay_seconds(attempt)
                        logger.warning(
                            "Model stream failed, retrying in %.2fs (model=%s, attempt=%s/%s): %s",
                            delay,
                            model_name,
                            attempt,
                            self.retry_max_attempts,
                            exc,
                        )
                        time.sleep(delay)
                    else:
                        logger.warning(
                            "Model stream failed after %s attempts, switching to fallback model: %s -> %s; error=%s",
                            self.retry_max_attempts,
                            model_name,
                            model_candidates[model_idx + 1],
                            exc,
                        )
                    continue

        raise RuntimeError("Model stream failed unexpectedly without captured exception")

    def _get_model_candidates(self) -> list:
        names = [str(self.model).strip()]
        for item in self.fallback_models:
            candidate = str(item or "").strip()
            if candidate and candidate not in names:
                names.append(candidate)
        return names

    def _retry_delay_seconds(self, attempt: int) -> float:
        # attempt  1 ,
        base_delay = self.retry_initial_delay * (self.retry_backoff_factor ** max(0, attempt - 1))
        jitter = random.uniform(0.0, self.retry_jitter) if self.retry_jitter > 0 else 0.0
        return min(self.retry_max_delay, base_delay + jitter)

    def _is_retriable_error(self, exc: Exception) -> bool:
        if isinstance(exc, (httpx.TimeoutException, httpx.TransportError, TimeoutError, ConnectionError)):
            return True

        text = f"{type(exc).__name__}: {exc}".lower()
        non_retriable_markers = (
            "invalid api key",
            "incorrect api key",
            "authentication",
            "unauthorized",
            "forbidden",
            "permission",
            "invalid_request",
            "badrequest",
            "context_length_exceeded",
            "max context length",
            "unsupported",
            "not found",
            "does not exist",
        )
        if any(marker in text for marker in non_retriable_markers):
            return False

        retriable_markers = (
            "connection error",
            "connection aborted",
            "connection reset",
            "remote protocol error",
            "temporarily unavailable",
            "service unavailable",
            "timeout",
            "timed out",
            "eof occurred in violation of protocol",
            "ssl",
            "429",
            "502",
            "503",
            "504",
            "rate limit",
            "try again",
        )
        return any(marker in text for marker in retriable_markers)

    @staticmethod
    def _parse_bool_option(value: Any, default: bool) -> bool:
        if value is None:
            return default
        if isinstance(value, bool):
            return value
        text = str(value).strip().lower()
        if text in {"1", "true", "yes", "y", "on"}:
            return True
        if text in {"0", "false", "no", "n", "off"}:
            return False
        return default

    @staticmethod
    def _parse_int_option(value: Any, default: int, minimum: int) -> int:
        try:
            parsed = int(value)
        except Exception:
            return default
        return max(minimum, parsed)

    @staticmethod
    def _parse_float_option(value: Any, default: float, minimum: float) -> float:
        try:
            parsed = float(value)
        except Exception:
            return default
        return max(minimum, parsed)

    @staticmethod
    def _parse_fallback_models(value: Any) -> list:
        if value is None:
            return []
        if isinstance(value, (list, tuple)):
            return [str(item).strip() for item in value if str(item or "").strip()]
        text = str(value).strip()
        if not text:
            return []
        return [segment.strip() for segment in text.split(",") if segment.strip()]

    def __repr__(self):
        return f"LiteLLMChat(model={self.model!r}, max_context_tokens={self.max_context_tokens})"


# ,, LiteLLMChat
DeepSeekReasonerChat = LiteLLMChat
DeepSeekReasonerResponse = LiteLLMResponse
