import hashlib
import json
from pathlib import Path
from typing import Any

import litellm
from litellm import CustomStreamWrapper, Message
from litellm.exceptions import (
    AuthenticationError as LiteLLMAuthenticationError,
)
from litellm.exceptions import (
    BadRequestError as LiteLLMBadRequestError,
)
from litellm.exceptions import (
    ContextWindowExceededError as LiteLLMContextWindowExceededError,
)
from litellm.litellm_core_utils.get_supported_openai_params import (
    get_supported_openai_params,
)
from litellm.utils import get_model_info
from pydantic import BaseModel
from tenacity import (
    retry,
    retry_if_exception_type,
    retry_if_not_exception_type,
    stop_after_attempt,
    wait_exponential,
)

from harbor.llms.base import (
    BaseLLM,
    ContextLengthExceededError,
    LLMResponse,
    OutputLengthExceededError,
)
from harbor.llms.utils import (
    add_anthropic_caching,
    split_provider_model_name,
    validate_hosted_vllm_model_config,
)
from harbor.models.metric import UsageInfo
from harbor.utils.logger import logger

# This is used as a fallback for when the model does not support response_format

STRUCTURED_RESPONSE_PROMPT_TEMPLATE = """
You must respond in the following JSON format.

Here is the json schema:

```json
{schema}
```

Here is the prompt:

{prompt}
"""


class LiteLLM(BaseLLM):
    def __init__(
        self,
        model_name: str,
        temperature: float = 1,
        api_base: str | None = None,
        session_id: str | None = None,
        collect_rollout_details: bool = False,
        max_thinking_tokens: int | None = None,
        reasoning_effort: str | None = None,
        model_info: dict[str, Any] | None = None,
        use_responses_api: bool = False,
        **kwargs,
    ):
        super().__init__()
        self._model_name = model_name
        self._llm_kwargs = kwargs
        self._temperature = temperature
        self._model_info = model_info
        self._logger = logger.getChild(__name__)

        hosted_vllm_validation: tuple[str, dict[str, Any]] | None = None
        if "hosted_vllm" in model_name.lower():
            hosted_vllm_validation = validate_hosted_vllm_model_config(
                model_name, self._model_info
            )

        (
            self._provider_prefix,
            self._canonical_model_name,
        ) = split_provider_model_name(model_name)

        if hosted_vllm_validation is not None:
            self._provider_prefix = "hosted_vllm"
            self._canonical_model_name, self._model_info = hosted_vllm_validation

        # Only use canonical model name for hosted_vllm; other providers like
        # openrouter need the full prefixed name for LiteLLM lookups
        self._litellm_model_name = (
            self._canonical_model_name
            if self._provider_prefix == "hosted_vllm"
            else model_name
        )

        # Register custom model if model_info is provided
        if self._model_info is not None:
            try:
                litellm.register_model({self._litellm_model_name: self._model_info})
                self._logger.debug(
                    f"Registered custom model '{model_name}' with info: {self._model_info}"
                )
            except Exception as e:
                self._logger.warning(
                    f"Failed to register custom model '{model_name}': {e}"
                )

        self._supported_params = get_supported_openai_params(self._litellm_model_name)
        self._api_base = api_base
        self._session_id = session_id
        self._collect_rollout_details = collect_rollout_details
        self._max_thinking_tokens = max_thinking_tokens
        self._reasoning_effort = reasoning_effort

        if self._supported_params is not None:
            self._supports_response_format = "response_format" in self._supported_params
            self._supports_temperature = "temperature" in self._supported_params
        else:
            self._supports_response_format = False
            self._supports_temperature = False

        self._use_responses_api = use_responses_api
        self._structured_response_prompt_template = STRUCTURED_RESPONSE_PROMPT_TEMPLATE

    @property
    def _lookup_model_name(self) -> str:
        """Get the model name to use for lookups in LiteLLM's model database.

        Returns the canonical model name (without provider prefix) if available,
        otherwise falls back to the original model name.
        """
        return self._litellm_model_name or self._model_name

    @property
    def _display_name(self) -> str:
        """Get a display name for logging that shows both canonical and original names if different."""
        lookup_name = self._lookup_model_name
        if lookup_name != self._model_name:
            return f"{lookup_name} (from '{self._model_name}')"
        return lookup_name

    def get_model_context_limit(self) -> int:
        """Get the context limit (max input tokens) for the current model.

        Returns:
            int: The maximum input tokens the model can accept, or a fallback value if unavailable.
        """
        fallback_context_limit = 1000000

        try:
            model_info = get_model_info(self._lookup_model_name)
            max_input_tokens = model_info.get("max_input_tokens")

            # Fallback to max_tokens if max_input_tokens not available
            if max_input_tokens is None:
                max_input_tokens = model_info.get("max_tokens")

            if max_input_tokens is not None:
                return max_input_tokens

            # Model info exists but doesn't have context limit info
            self._logger.warning(
                f"Model '{self._display_name}' info found but missing context limit fields. "
                f"Using fallback context limit: {fallback_context_limit}"
            )
        except Exception as e:
            self._logger.warning(
                f"Failed to retrieve model info for '{self._display_name}': {e}. "
                f"Using fallback context limit: {fallback_context_limit}"
            )

        return fallback_context_limit

    def get_model_output_limit(self) -> int | None:
        """Get the output limit (max output tokens) for the current model.

        Returns:
            int | None: The maximum output tokens the model can generate, or None if unavailable.
        """
        try:
            model_info = get_model_info(self._lookup_model_name)
            max_output_tokens = model_info.get("max_output_tokens")

            if max_output_tokens is None:
                # Model info exists but doesn't have max_output_tokens
                self._logger.debug(
                    f"Model '{self._display_name}' info found but missing max_output_tokens field."
                )

            return max_output_tokens
        except Exception as e:
            self._logger.debug(
                f"Failed to retrieve model info for '{self._display_name}': {e}."
            )
            return None

    def _clean_value(self, value):
        match value:
            case _ if callable(value):
                return None
            case dict():
                return {
                    k: v
                    for k, v in {
                        k: self._clean_value(v) for k, v in value.items()
                    }.items()
                    if v is not None
                }
            case list():
                return [
                    self._clean_value(v)
                    for v in value
                    if self._clean_value(v) is not None
                ]
            case str() | int() | float() | bool():
                return value
            case _:
                return str(value)

    def _init_logger_fn(self, logging_path: Path):
        def logger_fn(model_call_dict: dict):
            clean_dict = self._clean_value(model_call_dict)
            if isinstance(clean_dict, dict) and "api_key" in clean_dict:
                hash_key = hashlib.sha256(clean_dict["api_key"].encode()).hexdigest()
                clean_dict["api_key_sha256"] = hash_key
                del clean_dict["api_key"]

            if isinstance(clean_dict, dict) and "x-api-key" in clean_dict:
                hash_key = hashlib.sha256(clean_dict["x-api-key"].encode()).hexdigest()
                clean_dict["x-api-key_sha256"] = hash_key
                del clean_dict["x-api-key"]

            # Only save post_api_call (response) to preserve logprobs
            # logger_fn is called multiple times with different event types, and
            # write_text() overwrites the file each time. By filtering for
            # post_api_call only, we ensure the response with logprobs is saved.
            log_event_type = clean_dict.get("log_event_type", "unknown")
            if log_event_type == "post_api_call":
                logging_path.write_text(
                    json.dumps(
                        clean_dict,
                        indent=4,
                    )
                )

        return logger_fn

    @retry(
        stop=stop_after_attempt(3),
        wait=wait_exponential(multiplier=1, min=4, max=15),
        retry=(
            # To avoid asyncio.CancelledError retries which inherits from BaseException
            # rather than Exception
            retry_if_exception_type(Exception)
            & retry_if_not_exception_type(
                (
                    ContextLengthExceededError,
                    OutputLengthExceededError,
                    LiteLLMAuthenticationError,
                )
            )
        ),
        reraise=True,
    )
    async def call(
        self,
        prompt: str,
        message_history: list[dict[str, Any] | Message] = [],
        response_format: dict | type[BaseModel] | None = None,
        logging_path: Path | None = None,
        **kwargs,
    ) -> LLMResponse:
        if self._use_responses_api:
            return await self._call_responses(
                prompt, message_history, response_format, logging_path, **kwargs
            )

        if response_format is not None and not self._supports_response_format:
            if isinstance(response_format, dict):
                schema = json.dumps(response_format, indent=2)
            else:
                schema = json.dumps(response_format.model_json_schema(), indent=2)
            prompt = self._structured_response_prompt_template.format(
                schema=schema, prompt=prompt
            )
            response_format = None

        # Prepare messages with caching for Anthropic models
        messages = message_history + [{"role": "user", "content": prompt}]
        messages = add_anthropic_caching(messages, self._model_name)

        try:
            # Build completion_kwargs with all parameters
            completion_kwargs = {
                **self._build_base_kwargs(logging_path),
                "messages": messages,
                "temperature": self._temperature,
                "response_format": response_format,
                "reasoning_effort": self._reasoning_effort,
            }

            # Add logprobs and return_token_ids if rollout details collection is enabled
            if self._collect_rollout_details:
                completion_kwargs["logprobs"] = True
                # Request token IDs from provider (supported by vLLM)
                # Note: Some providers (e.g., OpenAI) will reject this parameter, but we'll catch and retry without it
                if "extra_body" not in completion_kwargs:
                    completion_kwargs["extra_body"] = {}
                extra_body: dict[str, Any] = completion_kwargs["extra_body"]  # type: ignore[assignment]
                extra_body["return_token_ids"] = True

            # Add any additional kwargs, deep-merging extra_body to preserve
            # internally-set fields (e.g., return_token_ids) when caller also
            # passes extra_body via llm_call_kwargs.
            if "extra_body" in completion_kwargs and "extra_body" in kwargs:
                existing_extra_body: dict[str, Any] = completion_kwargs["extra_body"]  # type: ignore[assignment]
                new_extra_body: dict[str, Any] = kwargs.pop("extra_body")
                completion_kwargs["extra_body"] = {
                    **existing_extra_body,
                    **new_extra_body,
                }
            elif "extra_body" in kwargs:
                kwargs["extra_body"] = {**kwargs["extra_body"]}
            completion_kwargs.update(kwargs)

            # Add thinking parameter for Anthropic models if max_thinking_tokens is set
            if self._max_thinking_tokens is not None and (
                "anthropic" in self._model_name.lower()
                or "claude" in self._model_name.lower()
            ):
                budget = self._max_thinking_tokens
                if budget < 1024:
                    self._logger.warning(
                        f"max_thinking_tokens={budget} is below minimum of 1024. "
                        "Using minimum value of 1024."
                    )
                    budget = 1024
                completion_kwargs["thinking"] = {
                    "type": "enabled",
                    "budget_tokens": budget,
                }

            # Add session_id to extra_body if available
            if self._session_id is not None:
                if "extra_body" not in completion_kwargs:
                    completion_kwargs["extra_body"] = {}
                extra_body: dict[str, Any] = completion_kwargs["extra_body"]  # type: ignore[assignment]
                extra_body["session_id"] = self._session_id

            try:
                response = await litellm.acompletion(**completion_kwargs)
            except LiteLLMBadRequestError as e:
                # If provider (e.g., OpenAI) rejects extra_body parameters, retry without them
                # Some providers reject custom parameters like: return_token_ids, session_id, etc.
                error_msg = str(e)
                if (
                    "Unrecognized request argument" in error_msg
                    and "extra_body" in completion_kwargs
                ):
                    rejected_params = []
                    if "return_token_ids" in error_msg:
                        rejected_params.append("return_token_ids")
                    if "session_id" in error_msg:
                        rejected_params.append("session_id")

                    if rejected_params:
                        self._logger.warning(
                            f"Provider {self._model_name} rejected extra_body parameters: {', '.join(rejected_params)}. "
                            f"Retrying without them. Token IDs will not be available."
                        )
                        extra_body_val = completion_kwargs.get("extra_body")
                        if extra_body_val and isinstance(extra_body_val, dict):
                            extra_body: dict[str, Any] = extra_body_val
                            for param in rejected_params:
                                if param in extra_body:
                                    del extra_body[param]

                            if not extra_body:
                                del completion_kwargs["extra_body"]

                        response = await litellm.acompletion(**completion_kwargs)
                    else:
                        raise e
                else:
                    raise e
        except Exception as e:
            self._handle_litellm_error(e)

        if isinstance(response, CustomStreamWrapper):
            raise NotImplementedError("Streaming is not supported for T bench yet")

        # Extract usage information
        usage_info = self._extract_usage_info(response)

        # Extract and process token IDs if rollout details collection is enabled
        prompt_token_ids = None
        completion_token_ids = None
        logprobs = None

        if self._collect_rollout_details:
            prompt_token_ids, completion_token_ids = self._extract_token_ids(response)
            logprobs = self._extract_logprobs(response)

        choice = response["choices"][0]
        message = choice["message"]
        content = message.get("content") or ""
        reasoning_content = message.get("reasoning_content")

        # Sometimes the LLM returns a response with a finish reason of "length"
        # This typically means we hit the max_tokens limit, not the context window
        if choice.get("finish_reason") == "length":
            # Create exception with truncated response attached
            exc = OutputLengthExceededError(
                f"Model {self._model_name} hit max_tokens limit. "
                f"Response was truncated. Consider increasing max_tokens if possible.",
                truncated_response=content,
            )
            raise exc

        return LLMResponse(
            content=content,
            reasoning_content=reasoning_content,
            usage=usage_info,
            prompt_token_ids=prompt_token_ids,
            completion_token_ids=completion_token_ids,
            logprobs=logprobs,
        )

    def _extract_token_ids(self, response) -> tuple[list[int] | None, list[int] | None]:
        """Extract token IDs from a response.

        Returns:
            Tuple of (prompt_token_ids, completion_token_ids) where:
            - prompt_token_ids: Full prompt token IDs including conversation history
            - completion_token_ids: Token IDs for the generated completion as provided by the provider
        """
        try:
            # Extract completion token IDs
            completion_token_ids = None
            choices = response.get("choices", [])
            if choices:
                choice = choices[0]
                if (
                    hasattr(choice, "provider_specific_fields")
                    and choice.provider_specific_fields
                ):
                    provider_token_ids = choice.provider_specific_fields.get(
                        "token_ids"
                    )
                    if provider_token_ids and isinstance(provider_token_ids, list):
                        completion_token_ids = provider_token_ids

            # Extract full prompt token IDs (including conversation history)
            prompt_token_ids = getattr(response, "prompt_token_ids", None)

            return prompt_token_ids, completion_token_ids

        except Exception as e:
            self._logger.debug(f"Error extracting token IDs: {e}")
            return None, None

    def _extract_logprobs(self, response) -> list[float] | None:
        """Extract logprobs from a response.

        Args:
            response: The LLM response object

        Returns:
            List of log probabilities for each token, or None if not available.
        """
        try:
            choices = response.get("choices", [])
            if not choices:
                return None

            logprobs_data = choices[0].get("logprobs")
            if not logprobs_data:
                return None

            content = logprobs_data.get("content", [])
            return [
                token_data["logprob"]
                for token_data in content
                if "logprob" in token_data
            ]
        except (KeyError, TypeError, IndexError):
            return None

    def _extract_cost(self, response) -> float:
        """Extract cost from a response's _hidden_params or compute via litellm.

        Args:
            response: The LLM response object

        Returns:
            The cost in USD, or 0.0 if unavailable.
        """
        cost = 0.0
        if hasattr(response, "_hidden_params"):
            hidden_params = response._hidden_params
            if isinstance(hidden_params, dict):
                cost = hidden_params.get("response_cost", 0.0) or 0.0

        if cost == 0.0:
            try:
                cost = litellm.completion_cost(completion_response=response) or 0.0
            except Exception:
                cost = 0.0

        return float(cost)

    def _extract_usage_info(self, response) -> UsageInfo | None:
        """Extract token usage and cost from a completion API response.

        Args:
            response: The LLM response object

        Returns:
            UsageInfo with token counts and cost, or None if not available.
        """
        try:
            # LiteLLM returns a ModelResponse object with a usage attribute
            # See: https://docs.litellm.ai/docs/completion/prompt_caching
            usage = response.usage

            prompt_tokens = getattr(usage, "prompt_tokens", 0)
            completion_tokens = getattr(usage, "completion_tokens", 0)

            # Get cache tokens from prompt_tokens_details.cached_tokens
            cache_tokens = 0
            if hasattr(usage, "prompt_tokens_details"):
                prompt_tokens_details = usage.prompt_tokens_details
                if prompt_tokens_details is not None:
                    cache_tokens = (
                        getattr(prompt_tokens_details, "cached_tokens", 0) or 0
                    )

            cost = self._extract_cost(response)

            return UsageInfo(
                prompt_tokens=prompt_tokens,
                completion_tokens=completion_tokens,
                cache_tokens=cache_tokens,
                cost_usd=cost,
            )
        except (AttributeError, TypeError):
            return None

    def _extract_responses_usage_info(self, response) -> UsageInfo | None:
        """Extract token usage and cost from a responses API response.

        Args:
            response: The responses API response object

        Returns:
            UsageInfo with token counts and cost, or None if not available.
        """
        if not hasattr(response, "usage") or response.usage is None:
            return None

        usage = response.usage
        prompt_tokens = getattr(usage, "input_tokens", 0) or 0
        completion_tokens = getattr(usage, "output_tokens", 0) or 0
        cost = self._extract_cost(response)

        return UsageInfo(
            prompt_tokens=prompt_tokens,
            completion_tokens=completion_tokens,
            cache_tokens=0,
            cost_usd=cost,
        )

    def _handle_litellm_error(self, e: Exception) -> None:
        """Translate litellm exceptions into harbor exceptions.

        Always re-raises; never returns normally.
        """
        if isinstance(e, LiteLLMContextWindowExceededError):
            raise ContextLengthExceededError
        if isinstance(e, LiteLLMAuthenticationError):
            raise e
        if isinstance(e, LiteLLMBadRequestError):
            if self._is_context_length_error(e):
                raise ContextLengthExceededError from e
        raise e

    def _build_base_kwargs(self, logging_path: Path | None = None) -> dict[str, Any]:
        """Build the base kwargs shared by both completion and responses API calls."""
        logger_fn = (
            self._init_logger_fn(logging_path) if logging_path is not None else None
        )
        return {
            **self._llm_kwargs,
            "model": self._model_name,
            "drop_params": True,
            "logger_fn": logger_fn,
            "api_base": self._api_base,
        }

    def _is_context_length_error(self, error: LiteLLMBadRequestError) -> bool:
        """Check provider error payloads for context-length overflow signals."""

        parts = [
            str(error),
            str(getattr(error, "body", "")),
            str(getattr(error, "message", "")),
            str(getattr(error, "error", "")),
        ]

        combined = " ".join(part.lower() for part in parts if part)
        phrases = (
            "context length exceeded",
            "context_length_exceeded",
            "maximum context length",
            "`inputs` tokens + `max_new_tokens`",
        )
        return any(phrase in combined for phrase in phrases)

    async def _call_responses(
        self,
        prompt: str,
        message_history: list[dict[str, Any] | Message] = [],
        response_format: dict | type[BaseModel] | None = None,
        logging_path: Path | None = None,
        **kwargs,
    ) -> LLMResponse:
        """Call the OpenAI Responses API via litellm.aresponses().

        When previous_response_id is provided (via kwargs), uses server-side
        state chaining — only the new user prompt is sent. Otherwise, builds
        the full input from message_history + prompt.
        """
        previous_response_id = kwargs.pop("previous_response_id", None)

        try:
            # Build responses_kwargs
            responses_kwargs: dict[str, Any] = self._build_base_kwargs(logging_path)

            if self._reasoning_effort is not None:
                responses_kwargs["reasoning"] = {
                    "effort": self._reasoning_effort,
                }
            elif self._supports_temperature:
                # Reasoning models don't support temperature, only set it
                # when reasoning_effort is not configured
                responses_kwargs["temperature"] = self._temperature

            # Get max_output_tokens from model info
            max_output_tokens = self.get_model_output_limit()
            if max_output_tokens is not None:
                responses_kwargs["max_output_tokens"] = max_output_tokens

            if response_format is not None:
                responses_kwargs["response_format"] = response_format

            if previous_response_id is not None:
                # Server-side state chaining: only send the new prompt
                responses_kwargs["previous_response_id"] = previous_response_id
                responses_kwargs["input"] = prompt
            else:
                # Build full input from message history + new prompt
                input_items = []
                for msg in message_history:
                    role = (
                        msg.get("role", "user")
                        if isinstance(msg, dict)
                        else getattr(msg, "role", "user")
                    )
                    content = (
                        msg.get("content", "")
                        if isinstance(msg, dict)
                        else getattr(msg, "content", "")
                    )
                    input_items.append({"role": role, "content": content})
                input_items.append({"role": "user", "content": prompt})
                responses_kwargs["input"] = input_items

            response = await litellm.aresponses(**responses_kwargs)

        except Exception as e:
            self._handle_litellm_error(e)

        # Extract text content from response.output
        content = ""
        reasoning_content = None
        for output_item in response.output:
            if getattr(output_item, "type", None) == "message":
                for content_part in getattr(output_item, "content", []):
                    if getattr(content_part, "type", None) == "output_text":
                        content += getattr(content_part, "text", "")

        # Extract usage information
        usage_info = self._extract_responses_usage_info(response)

        # Check for truncation via response status
        response_status = getattr(response, "status", None)
        if response_status == "incomplete":
            incomplete_details = getattr(response, "incomplete_details", None)
            reason = (
                getattr(incomplete_details, "reason", "unknown")
                if incomplete_details
                else "unknown"
            )
            if reason == "max_output_tokens":
                raise OutputLengthExceededError(
                    f"Model {self._model_name} hit max_tokens limit. "
                    f"Response was truncated.",
                    truncated_response=content,
                )

        response_id = getattr(response, "id", None)

        return LLMResponse(
            content=content,
            reasoning_content=reasoning_content,
            usage=usage_info,
            response_id=response_id,
        )
