"""
Shared utilities for inspect evaluation.

This module contains the common classes and functions used by both
run_inspect_evals.py and inspect_evaluator.py to avoid code duplication.
"""

import logging
import time
from typing import Sequence

import tinker
from inspect_ai.model import ChatCompletionChoice as InspectAIModelOutputChoice
from inspect_ai.model import ChatMessage as InspectAIChatMessage
from inspect_ai.model import ChatMessageAssistant as InspectAIChatMessageAssistant
from inspect_ai.model import ChatMessageSystem, Content
from inspect_ai.model import GenerateConfig as InspectAIGenerateConfig
from inspect_ai.model import ModelAPI as InspectAIModelAPI
from inspect_ai.model import ModelOutput as InspectAIModelOutput
from inspect_ai.model import ModelUsage as InspectAIModelUsage
from inspect_ai.model._registry import modelapi_register
from inspect_ai.tool import ToolChoice as InspectAIToolChoice
from inspect_ai.tool import ToolInfo as InspectAIToolInfo
from termcolor import colored
from tinker_cookbook import renderers
from tinker_cookbook.tokenizer_utils import get_tokenizer

logger = logging.getLogger(__name__)


def get_model_usage(
    tokenized_prompt: Sequence[int], responses: Sequence[tinker.SampledSequence]
) -> InspectAIModelUsage:
    """
    Given a tokenized prompt and a list of responses, return the number of tokens used/generated by the model.
    """
    num_input_tokens = len(tokenized_prompt)
    num_output_tokens = sum(len(r.tokens) for r in responses)
    total_tokens = num_input_tokens + num_output_tokens
    usage = InspectAIModelUsage(
        input_tokens=num_input_tokens, output_tokens=num_output_tokens, total_tokens=total_tokens
    )
    return usage


def convert_inspect_messages(messages: list[InspectAIChatMessage]) -> list[renderers.Message]:
    def assert_string(content: str | list[Content]) -> str:
        if isinstance(content, str):
            return content
        else:
            raise ValueError(f"Invalid content: {content}")

    return [
        renderers.Message(role=m.role, content=assert_string(m.content).strip()) for m in messages
    ]


class InspectAPIFromTinkerSampling(InspectAIModelAPI):
    """
    A model API wrapper that adapts tinker sampling clients to the inspect API interface.

    This class can be initialized either with a model_path (for standalone use)
    or with a sampling_client (for use in evaluators).
    """

    def __init__(
        self,
        renderer_name: str,
        model_name: str,
        model_path: str | None = None,
        sampling_client: tinker.SamplingClient | None = None,
        base_url: str | None = None,
        api_key: str | None = None,
        api_key_vars: list[str] = [],
        config: InspectAIGenerateConfig = InspectAIGenerateConfig(),
        verbose: bool = False,
    ):
        super().__init__(
            model_name=model_name,
            base_url=base_url,
            api_key=api_key,
            api_key_vars=api_key_vars,
            config=config,
        )

        # Initialize sampling client
        if sampling_client is not None:
            self.sampling_client = sampling_client
        elif model_path is not None:
            service_client = tinker.ServiceClient(api_key=api_key)
            self.sampling_client = service_client.create_sampling_client(model_path=model_path)
        else:
            raise ValueError("Either model_path or sampling_client must be provided")

        # Initialize renderer and tokenizer
        tokenizer = get_tokenizer(model_name)
        self.renderer = renderers.get_renderer(name=renderer_name, tokenizer=tokenizer)
        self.verbose = verbose

    async def generate(
        self,
        input: list[InspectAIChatMessage],
        tools: list[InspectAIToolInfo],
        tool_choice: InspectAIToolChoice,
        config: InspectAIGenerateConfig,
    ) -> InspectAIModelOutput:
        """
        The main interface that needs to be implemented to test a new model.
        """
        if config.system_message:
            input = [ChatMessageSystem(content=config.system_message)] + input
        convo = convert_inspect_messages(input)
        prompt = self.renderer.build_generation_prompt(convo)
        num_responses = 1 if config.num_choices is None else config.num_choices
        sampling_params = tinker.SamplingParams(
            temperature=config.temperature if config.temperature is not None else 1.0,
            max_tokens=config.max_tokens or 128,
            stop=self.renderer.get_stop_sequences(),
            top_p=config.top_p if config.top_p is not None else 1.0,
            top_k=config.top_k if config.top_k is not None else -1,
            seed=config.seed,
        )

        start_time = time.time()
        sample_result = await self.sampling_client.sample_async(
            prompt=prompt, sampling_params=sampling_params, num_samples=num_responses
        )
        sampled_token_sequences = sample_result.sequences

        # Optional verbose output (only for standalone use)
        if self.verbose:
            prompt_text = colored(self.renderer.tokenizer.decode(prompt.to_ints()), "green")
            logger.info(f"[Prompt]\n{prompt_text}")
            for i, seq in enumerate(sampled_token_sequences):
                response_text = colored(self.renderer.tokenizer.decode(seq.tokens), "red")
                logger.info(f"[Response {i + 1}/{num_responses}]\n{response_text}")

        end_time = time.time()

        parsed_responses = [
            self.renderer.parse_response(r.tokens)[0] for r in sampled_token_sequences
        ]
        responses_text: list[str] = [renderers.get_text_content(r) for r in parsed_responses]
        all_choices = [
            InspectAIModelOutputChoice(
                message=InspectAIChatMessageAssistant(content=r, model=self.model_name),
                stop_reason="stop",
            )
            for r in responses_text
        ]
        usage = get_model_usage(prompt.to_ints(), sampled_token_sequences)

        return InspectAIModelOutput(
            model=self.model_name, choices=all_choices, time=end_time - start_time, usage=usage
        )


# Register with inspect_ai's model registry.
# Using modelapi_register instead of @modelapi decorator preserves the __init__ signature for pyright.
modelapi_register(InspectAPIFromTinkerSampling, "tinker-sampling")
