# Copyright 2025 - Oumi
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import copy
import dataclasses
import hashlib
import json
import time
import uuid
from abc import ABC, abstractmethod
from pathlib import Path
from typing import Optional

import jsonlines
from hdrh.histogram import HdrHistogram
from tqdm import tqdm

from oumi.core.configs import (
    GenerationParams,
    InferenceConfig,
    ModelParams,
)
from oumi.core.types.conversation import Conversation
from oumi.utils.logging import logger
from oumi.utils.math_utils import is_power_of_two



[docs]
class BaseInferenceEngine(ABC):
    """Base class for running model inference."""

    _model_params: ModelParams
    """The model parameters."""

    _generation_params: GenerationParams
    """The generation parameters."""

    def __init__(
        self,
        model_params: ModelParams,
        *,
        generation_params: Optional[GenerationParams] = None,
    ):
        """Initializes the inference engine.

        Args:
            model_params: The model parameters.
            generation_params: The generation parameters.
        """
        self._model_params = copy.deepcopy(model_params)
        if generation_params:
            self._check_unsupported_params(generation_params)
        else:
            generation_params = GenerationParams()
        self._generation_params = generation_params

        self._latency_histogram_online = HdrHistogram(1, 60 * 1000, 1)
        self._latency_histogram_from_file = HdrHistogram(20, 180 * 1000, 1)
        self._dataset_hash = None


[docs]
    def infer(
        self,
        input: Optional[list[Conversation]] = None,
        inference_config: Optional[InferenceConfig] = None,
    ) -> list[Conversation]:
        """Runs model inference.

        Args:
            input: A list of conversations to run inference on. Optional.
            inference_config: Parameters for inference.
                If not specified, a default config is inferred.

        Returns:
            List[Conversation]: Inference output.
        """
        if input is not None and (
            inference_config and inference_config.input_path is not None
        ):
            raise ValueError(
                "Only one of input or inference_config.input_path should be provided."
            )

        # Ensure the inference config has up-to-date generation parameters.
        if inference_config:
            if inference_config.generation:
                self._check_unsupported_params(inference_config.generation)
            elif self._generation_params:
                inference_config = copy.deepcopy(inference_config)
                inference_config.generation = self._generation_params

                # Warn the user: They provided an inference config without generation
                # params, so what was the point of providing it in the first place?
                logger.warning(
                    "No generation parameters provided in the inference config. Using "
                    "the generation parameters that the engine was initialized with."
                )

        output_path = inference_config.output_path if inference_config else None

        # Load input conversations either from file or use provided input
        conversations_to_process: list[Conversation] = []
        if input is not None:
            conversations_to_process = input
        elif inference_config and inference_config.input_path is not None:
            conversations_to_process = self._read_conversations(
                inference_config.input_path,
            )
        else:
            raise ValueError(
                "One of input or inference_config.input_path must be provided."
            )

        unique_conversation_ids = set()
        for i, conversation in enumerate(conversations_to_process):
            if conversation.conversation_id is not None:
                if conversation.conversation_id in unique_conversation_ids:
                    raise ValueError(
                        f"Conversation ID {conversation.conversation_id} is not unique."
                    )
                unique_conversation_ids.add(conversation.conversation_id)
                continue

            # Generate a deterministic conversation ID based on index if none exists
            messages = conversation.messages
            all_str_message_content = ",".join(
                [
                    message.content
                    for message in messages
                    if isinstance(message.content, str)
                ]
            )
            content_hash = hashlib.sha256(all_str_message_content.encode()).hexdigest()
            id_name = str(i) + "_" + content_hash

            conversation.conversation_id = str(
                uuid.uuid5(
                    uuid.NAMESPACE_DNS,
                    id_name,
                )
            )
            unique_conversation_ids.add(conversation.conversation_id)

        # Calculate the hash of the dataset to use for referencing the inference
        # filename between runs, given the same model and generation parameters.
        dataset_hash = ""
        if len(conversations_to_process) > 0:
            row_hashes = [
                hashlib.sha256(c.to_json().encode()).hexdigest()
                for c in conversations_to_process
            ]
            dataset_hash = hashlib.sha256(",".join(row_hashes).encode()).hexdigest()
        self._dataset_hash = dataset_hash

        completed_conversations = self._load_from_scratch(output_path)

        # Filter out already completed conversations
        remaining_conversations = self._filter_incomplete_conversations(
            conversations_to_process, completed_conversations
        )

        if len(remaining_conversations) < len(conversations_to_process):
            logger.info(
                f"Found {len(completed_conversations)} completed conversations. "
                f"Processing remaining {len(remaining_conversations)} conversations."
            )

        # Run inference only on remaining conversations
        start_time = time.perf_counter()
        histogram = self._latency_histogram_online
        inference_results = self._infer_online(
            remaining_conversations, inference_config
        )
        histogram.record_value((time.perf_counter() - start_time) * 1e3)
        self._maybe_log_latency_histogram(histogram)

        if len(inference_results) == len(conversations_to_process):
            final_results = inference_results
        else:
            # Incomplete inference results were saved to scratch file.
            # Load all results from scratch to get all results.
            final_results = self._load_from_scratch(output_path)
            if len(final_results) != len(conversations_to_process):
                raise ValueError(
                    f"Number of final results ({len(final_results)}) does not match "
                    f"number of conversations to process "
                    f"({len(conversations_to_process)})."
                )

        self._cleanup_scratch_file(output_path)

        sorted_conversations = {
            conv.conversation_id: conv for conv in conversations_to_process
        }

        for conv in final_results:
            sorted_conversations[conv.conversation_id] = conv

        final_results = list(sorted_conversations.values())

        if inference_config and inference_config.output_path:
            self._save_conversations(final_results, inference_config.output_path)

        return final_results


    def _maybe_log_latency_histogram(self, histogram: Optional[HdrHistogram]) -> None:
        """Logs the histogram if it is not None.

        Args:
            histogram: The histogram to log.
        """
        if histogram is None:
            return
        total_count = histogram.get_total_count()
        # TODO: Define a better way to enable/configure this logging.
        if not (
            isinstance(total_count, int)
            and total_count >= 2
            and is_power_of_two(total_count)
        ):
            return

        p50 = histogram.get_value_at_percentile(50)
        p90 = histogram.get_value_at_percentile(90)
        p99 = histogram.get_value_at_percentile(99)
        logger.debug(
            f"{self.__class__.__name__}: "
            f"Latency Histogram: {total_count} samples recorded:"
            f"\tp50: {p50:.1f}ms\tp90: {p90:.1f}ms\tp99: {p99:.1f}ms"
        )

    def _read_conversations(self, input_filepath: str) -> list[Conversation]:
        """Reads conversations from a file in Oumi chat format.

        Args:
            input_filepath: The path to the file containing the conversations.

        Returns:
            List[Conversation]: A list of conversations read from the file.
        """
        conversations = []
        with open(input_filepath) as f:
            for line in f:
                # Only parse non-empty lines.
                if line.strip():
                    conversation = Conversation.from_json(line)
                    conversations.append(conversation)
        return conversations

    def _load_from_scratch(self, output_filepath: Optional[str]) -> list[Conversation]:
        """Loads conversations from a scratch file.

        Args:
            output_filepath: The path to the output file. This is used to determine the
                location of the scratch file.

        Returns:
            list[Conversation]: A list of conversations loaded from the scratch file,
                or an empty list if the scratch file does not exist or is empty.
        """
        scratch_filepath = self._get_scratch_filepath(output_filepath)
        if not Path(scratch_filepath).exists():
            return []

        conversations = []
        with jsonlines.open(scratch_filepath, mode="r") as reader:
            for line in reader:
                conversations.append(Conversation.from_dict(line))
        return conversations

    def _filter_incomplete_conversations(
        self,
        input_conversations: list[Conversation],
        completed_conversations: list[Conversation],
    ) -> list[Conversation]:
        """Filters out conversations that have already been completed.

        Args:
            input_conversations: List of conversations to run inference on
            completed_conversations: List of conversations already completed

        Returns:
            list[Conversation]: List of conversations that still need inference results
        """
        completed_ids = {
            conv.conversation_id
            for conv in completed_conversations
            if conv.conversation_id is not None
        }
        return [
            conv
            for conv in input_conversations
            if conv.conversation_id not in completed_ids
        ]

    def _get_scratch_filepath(self, output_filepath: Optional[str]) -> str:
        """Returns a scratch filepath for the given output filepath.

        For example, if the output filepath is "/foo/bar/output.json", the scratch
        filepath will be "/foo/bar/scratch/output.json"

        If no output filepath is provided, a temporary file is used and placed in the
        current working directory under the name "tmp/temp_inference_output.jsonl".

        Args:
            output_filepath: The output filepath.

        Returns:
            str: The scratch filepath.
        """
        if output_filepath is not None:
            original_filepath = Path(output_filepath)
            return str(original_filepath.parent / "scratch" / original_filepath.name)

        model_params = self._model_params
        model_params_str = json.dumps(dataclasses.asdict(model_params))
        generation_params = self._generation_params
        generation_params_str = json.dumps(dataclasses.asdict(generation_params))
        inference_hash = hashlib.sha256(
            f"{model_params_str}_{generation_params_str}_{self._dataset_hash}".encode()
        ).hexdigest()

        path_prefix = Path.home() / ".cache" / "oumi" / "tmp"
        return str(path_prefix / f"temp_inference_output_{inference_hash}.jsonl")

    def _save_conversation_to_scratch(
        self, conversation: Conversation, output_filepath: Optional[str]
    ) -> None:
        """Appends a conversation to a file in Oumi chat format.

        Args:
            conversation: The conversation to save.
            output_filepath: The path to the file where the conversation should be
                saved.
        """
        # Make the directory if it doesn't exist.
        scratch_filepath = self._get_scratch_filepath(output_filepath)
        Path(scratch_filepath).parent.mkdir(parents=True, exist_ok=True)
        with jsonlines.open(scratch_filepath, mode="a") as writer:
            json_obj = conversation.to_dict()
            writer.write(json_obj)

    def _cleanup_scratch_file(self, output_filepath: Optional[str]) -> None:
        """Delete the scratch file from the file system if it exists.

        Args:
            output_filepath: The path to the output file. This is used to determine the
                location of the scratch file.
        """
        scratch_filepath = self._get_scratch_filepath(output_filepath)
        if Path(scratch_filepath).exists():
            Path(scratch_filepath).unlink()

    def _save_conversations(
        self, conversations: list[Conversation], output_filepath: str
    ) -> None:
        """Saves conversations to a file in Oumi chat format.

        Args:
            conversations: A list of conversations to save.
            output_filepath: The path to the file where the conversations should be
                saved.
        """
        # Make the directory if it doesn't exist.
        Path(output_filepath).parent.mkdir(parents=True, exist_ok=True)
        with jsonlines.open(output_filepath, mode="w") as writer:
            for conversation in tqdm(conversations, desc="Saving conversations"):
                json_obj = conversation.to_dict()
                writer.write(json_obj)

    def _check_unsupported_params(self, generation_params: GenerationParams):
        """Checks for unsupported parameters and logs warnings.

        If a parameter is not supported, and a non-default value is provided,
        a warning is logged.
        """
        supported_params = self.get_supported_params()
        default_generation_params = GenerationParams()

        for param_name, value in generation_params:
            if param_name not in supported_params:
                is_non_default_value = (
                    getattr(default_generation_params, param_name) != value
                )

                if is_non_default_value:
                    logger.warning(
                        f"{self.__class__.__name__} does not support {param_name}. "
                        f"Received value: {param_name}={value}. "
                        "This parameter will be ignored."
                    )


[docs]
    @abstractmethod
    def get_supported_params(self) -> set[str]:
        """Returns a set of supported generation parameters for this engine.

        Override this method in derived classes to specify which parameters
        are supported.

        Returns:
            Set[str]: A set of supported parameter names.
        """
        raise NotImplementedError


    @abstractmethod
    def _infer_online(
        self,
        input: list[Conversation],
        inference_config: Optional[InferenceConfig] = None,
    ) -> list[Conversation]:
        """Runs model inference online.

        Args:
            input: A list of conversations to run inference on.
            inference_config: Parameters for inference.

        Returns:
            List[Conversation]: Inference output.
        """
        raise NotImplementedError


[docs]
    def apply_chat_template(
        self, conversation: Conversation, **tokenizer_kwargs
    ) -> str:
        """Applies the chat template to the conversation.

        Args:
            conversation: The conversation to apply the chat template to.
            tokenizer_kwargs: Additional keyword arguments to pass to the tokenizer.

        Returns:
            str: The conversation with the chat template applied.
        """
        tokenizer = getattr(self, "_tokenizer", None)

        if tokenizer is None:
            raise ValueError("Tokenizer is not initialized.")

        if tokenizer.chat_template is None:
            raise ValueError("Tokenizer does not have a chat template.")

        if "tokenize" not in tokenizer_kwargs:
            tokenizer_kwargs["tokenize"] = False

        return tokenizer.apply_chat_template(conversation, **tokenizer_kwargs)