# Copyright 2025 Thinking Machines Lab
# Licensed under the Apache License, Version 2.0
#
# Modifications:
# - Adapted for HPC-AI cloud fine-tuning workflow
# Copyright © 2025 HPC-AI.COM
"""
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 hpcai
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, modelapi
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.tool import ToolChoice as InspectAIToolChoice
from inspect_ai.tool import ToolInfo as InspectAIToolInfo
from termcolor import colored
from hpcai.cookbook import renderers
from hpcai.cookbook.tokenizer_utils import get_tokenizer

logger = logging.getLogger(__name__)


def get_model_usage(
    tokenized_prompt: Sequence[int], responses: Sequence[hpcai.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
    ]


@modelapi(name="hpcai-sampling")
class InspectAPIFromTinkerSampling(InspectAIModelAPI):
    """
    A model API wrapper that adapts hpcai 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: hpcai.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 = hpcai.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 = hpcai.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,
        )

        start_time = time.time()
        assert num_responses == 1
        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:
            logger.info(
                colored(self.renderer.tokenizer.decode(prompt.to_ints()), "green")
                + colored(self.renderer.tokenizer.decode(sampled_token_sequences[0].tokens), "red")
            )

        end_time = time.time()

        parsed_responses = [
            self.renderer.parse_response(r.tokens)[0] for r in sampled_token_sequences
        ]
        responses_text: list[str] = []
        for r in parsed_responses:
            content = r["content"]
            assert isinstance(content, str), "Expected string content from parser"
            responses_text.append(content)
        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
        )
