"""Tinker LLM implementation for sampling with Terminus 2.

This module provides a TinkerLLM class that integrates with Thinking Machines Lab's
Tinker API for inference/sampling. It can be used as a drop-in replacement for LiteLLM
when running Terminus 2 agent evaluations.

Requirements:
    Install the tinker optional dependencies:
    ```bash
    uv pip install harbor[tinker]
    ```
"""

from __future__ import annotations

from typing import TYPE_CHECKING, Any

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

# Tinker is an optional dependency
try:
    import tinker
    from tinker_cookbook.model_info import get_recommended_renderer_name
    from tinker_cookbook.renderers import Renderer, get_renderer
    from tinker_cookbook.tokenizer_utils import get_tokenizer

    TINKER_AVAILABLE = True
except ImportError:
    TINKER_AVAILABLE = False

if TYPE_CHECKING:
    import tinker
    from tinker_cookbook.model_info import get_recommended_renderer_name
    from tinker_cookbook.renderers import Renderer, get_renderer
    from tinker_cookbook.tokenizer_utils import get_tokenizer

DEFAULT_CONTEXT_LIMIT = 32000
DEFAULT_OUTPUT_LIMIT = 8192


class TinkerLLM(BaseLLM):
    """LLM implementation using Tinker API for sampling.

    This class provides an interface to Tinker's sampling capabilities that is compatible
    with Harbor's BaseLLM interface, enabling it to be used with Terminus 2 and other
    agents that accept custom LLM backends.

    Use ``llm_backend="tinker"`` when configuring Terminus 2 to use this backend.
    """

    def __init__(
        self,
        model_name: str,
        model_path: str | None = None,
        temperature: float = 1.0,
        max_tokens: int = DEFAULT_OUTPUT_LIMIT,
        renderer_name: str | None = None,
        context_limit: int | None = None,
        output_limit: int | None = None,
        collect_rollout_details: bool = True,
        **kwargs,
    ):
        """Initialize TinkerLLM.

        Args:
            model_name: Name of the base model (e.g., "Qwen/Qwen3-8B").
                Used to automatically select the renderer and tokenizer via
                ``tinker_cookbook.model_info``.
            model_path: Optional Tinker path to saved weights
                (e.g., "tinker://run-id/weights/checkpoint-001").
            temperature: Sampling temperature (default: 1.0).
            max_tokens: Maximum tokens to generate per response (default: 8192).
            renderer_name: Name of the renderer to use for message formatting.
                If not provided, automatically discovered via
                ``tinker_cookbook.model_info.get_recommended_renderer_name``.
            context_limit: Override for model context limit (default: 32000).
            output_limit: Override for model output limit.
            collect_rollout_details: Whether to collect token IDs and logprobs (default: True).
            **kwargs: Additional arguments passed to BaseLLM.

        Raises:
            ImportError: If tinker or tinker-cookbook packages are not installed.
        """
        if not TINKER_AVAILABLE:
            raise ImportError(
                "TinkerLLM requires the 'tinker' and 'tinker-cookbook' packages. "
                "Install them with: uv pip install harbor[tinker]"
            )

        super().__init__(**kwargs)
        self._model_name = model_name
        self._model_path = model_path
        self._temperature = temperature
        self._max_tokens = max_tokens
        self._collect_rollout_details = collect_rollout_details
        self._logger = logger.getChild(__name__)

        self._context_limit = context_limit or DEFAULT_CONTEXT_LIMIT
        self._output_limit = output_limit or max_tokens

        # Discover renderer name from tinker_cookbook if not explicitly provided
        resolved_renderer_name = renderer_name
        if resolved_renderer_name is None:
            try:
                resolved_renderer_name = get_recommended_renderer_name(model_name)
            except (ValueError, KeyError):
                raise ValueError(
                    f"Could not discover renderer for model '{model_name}'. "
                    f"Pass renderer_name= explicitly."
                )
        self._renderer_name = resolved_renderer_name

        # Load tokenizer and initialize renderer
        tokenizer = get_tokenizer(model_name)
        self._renderer: Renderer = get_renderer(self._renderer_name, tokenizer)

        # Lazily initialized clients
        self._service_client: tinker.ServiceClient | None = None
        self._sampling_client: tinker.SamplingClient | None = None

        self._logger.info(
            f"TinkerLLM initialized with model={model_name}, "
            f"renderer={self._renderer_name}, "
            f"context_limit={self._context_limit}, "
            f"collect_rollout_details={collect_rollout_details}"
        )

    async def _ensure_client(self) -> tinker.SamplingClient:
        """Ensure the Tinker sampling client is initialized."""
        if self._sampling_client is not None:
            return self._sampling_client

        self._logger.debug("Initializing Tinker service client...")
        self._service_client = tinker.ServiceClient()

        if self._model_path:
            self._logger.info(
                f"Creating sampling client from saved weights: {self._model_path}"
            )
            self._sampling_client = (
                await self._service_client.create_sampling_client_async(
                    model_path=self._model_path
                )
            )
        else:
            self._logger.info(
                f"Creating sampling client for base model: {self._model_name}"
            )
            self._sampling_client = (
                await self._service_client.create_sampling_client_async(
                    base_model=self._model_name
                )
            )

        return self._sampling_client

    async def call(
        self,
        prompt: str,
        message_history: list[dict[str, Any]] = [],
        **kwargs,
    ) -> LLMResponse:
        """Make a sampling call to the Tinker API.

        Args:
            prompt: The user prompt for this turn.
            message_history: Previous messages in the conversation.
            **kwargs: Additional arguments (ignored for compatibility).

        Returns:
            LLMResponse containing the generated content and metadata.

        Raises:
            ContextLengthExceededError: If the prompt exceeds the model's context limit.
            OutputLengthExceededError: If the response was truncated due to max_tokens.
        """
        sampling_client = await self._ensure_client()

        # Convert messages to renderer format
        messages = []
        for msg in message_history:
            role = msg.get("role", "user")
            content = msg.get("content", "")
            messages.append({"role": role, "content": content})
        messages.append({"role": "user", "content": prompt})

        # Build the generation prompt using the renderer
        model_input = self._renderer.build_generation_prompt(messages)

        # Get prompt token count for context checking
        prompt_tokens = model_input.to_ints()
        prompt_token_count = len(prompt_tokens)

        if prompt_token_count > self._context_limit:
            raise ContextLengthExceededError(
                f"Prompt length ({prompt_token_count} tokens) exceeds "
                f"model context limit ({self._context_limit} tokens)"
            )

        # Get stop sequences from renderer
        stop_sequences = self._renderer.get_stop_sequences()

        # Build sampling parameters
        sampling_params = tinker.SamplingParams(
            max_tokens=self._max_tokens,
            temperature=self._temperature,
            stop=stop_sequences,
        )

        try:
            sample_response = await sampling_client.sample_async(
                prompt=model_input,
                num_samples=1,
                sampling_params=sampling_params,
            )

            sequence = sample_response.sequences[0]
            completion_tokens = sequence.tokens
            completion_logprobs = sequence.logprobs

            # Parse the response using the renderer
            parsed_message, parse_success = self._renderer.parse_response(
                completion_tokens
            )
            content = parsed_message.get("content", "")

            # Check if response was truncated (hit max_tokens without stop token)
            if not parse_success and len(completion_tokens) >= self._max_tokens:
                raise OutputLengthExceededError(
                    f"Response was truncated at max_tokens={self._max_tokens}",
                    truncated_response=content,
                )

            usage = UsageInfo(
                prompt_tokens=prompt_token_count,
                completion_tokens=len(completion_tokens),
                cache_tokens=0,
                cost_usd=0.0,
            )

            response_kwargs: dict[str, Any] = {
                "content": content,
                "reasoning_content": None,
                "usage": usage,
            }

            if self._collect_rollout_details:
                response_kwargs["prompt_token_ids"] = prompt_tokens
                response_kwargs["completion_token_ids"] = completion_tokens
                if completion_logprobs is not None:
                    response_kwargs["logprobs"] = list(completion_logprobs)

            return LLMResponse(**response_kwargs)

        except (ContextLengthExceededError, OutputLengthExceededError):
            raise
        except Exception as e:
            error_msg = str(e).lower()
            if any(
                phrase in error_msg
                for phrase in [
                    "context length",
                    "context_length",
                    "maximum context",
                    "token limit",
                    "too long",
                ]
            ):
                raise ContextLengthExceededError(str(e)) from e
            raise

    def get_model_context_limit(self) -> int:
        return self._context_limit

    def get_model_output_limit(self) -> int | None:
        return self._output_limit
