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

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.utils import get_model_info
from litellm.litellm_core_utils.get_supported_openai_params import (
    get_supported_openai_params,
)
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
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: Literal["none", "minimal", "low", "medium", "high", "default"]
        | None = None,
        model_info: dict[str, Any] | None = None,
        **kwargs,
    ):
        super().__init__(**kwargs)
        self._model_name = model_name
        self._temperature = temperature
        self._model_info = model_info
        self._logger = logger.getChild(__name__)

        # Register custom model if model_info is provided
        if model_info is not None:
            try:
                litellm.register_model({model_name: model_info})
                self._logger.debug(
                    f"Registered custom model '{model_name}' with info: {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(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._structured_response_prompt_template = STRUCTURED_RESPONSE_PROMPT_TEMPLATE

    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._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._model_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._model_name}': {e}. "
                f"Using fallback context limit: {fallback_context_limit}"
            )

        return fallback_context_limit

    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 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)

        if logging_path is not None:
            logger_fn = self._init_logger_fn(logging_path)
        else:
            logger_fn = None

        try:
            # Build completion_kwargs with all parameters
            completion_kwargs = {
                "model": self._model_name,
                "messages": messages,
                "temperature": self._temperature,
                "response_format": response_format,
                "drop_params": True,
                "logger_fn": logger_fn,
                "api_base": self._api_base,
                "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
                completion_kwargs.setdefault("extra_body", {})["return_token_ids"] = (
                    True
                )

            # Add any additional kwargs
            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:
                completion_kwargs.setdefault("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."
                        )
                        for param in rejected_params:
                            if param in completion_kwargs["extra_body"]:
                                del completion_kwargs["extra_body"][param]

                        if not completion_kwargs["extra_body"]:
                            del completion_kwargs["extra_body"]

                        response = await litellm.acompletion(**completion_kwargs)
                    else:
                        raise e
                else:
                    raise e
        except Exception as e:
            # Return the terminal-bench exception
            if isinstance(e, LiteLLMContextWindowExceededError):
                raise ContextLengthExceededError
            if isinstance(e, LiteLLMAuthenticationError):
                raise e  # Re-raise as-is so QualityChecker can catch it
            if isinstance(e, LiteLLMBadRequestError):
                if self._is_context_length_error(e):
                    raise ContextLengthExceededError from e

            raise 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_usage_info(self, response) -> UsageInfo | None:
        """Extract token usage and cost from a 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
                    )

            # Get cost from _hidden_params or calculate it
            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

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

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

    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)
