from typing import Dict, List, Optional, Union

from galtea.application.services.session_service import SessionService
from galtea.domain.models.inference_result import InferenceResult, InferenceResultBase
from galtea.infrastructure.clients.http_client import Client
from galtea.utils.string import build_query_params, is_valid_id


class InferenceResultService:
    """
    Service for managing Inference Results.
    An Inference Result is a single turn in a conversation between a user and the AI, including the input and output.
    """

    def __init__(self, client: Client, session_service: SessionService):
        """Initialize the InferenceResultService with the provided HTTP client.

        Args:
            client (Client): The HTTP client for making API requests.
        """
        self.__client: Client = client
        self.__session_service: SessionService = session_service

    def create(
        self,
        session_id: str,
        input: str,
        output: Optional[str] = None,
        retrieval_context: Optional[str] = None,
        latency: Optional[float] = None,
        usage_info: Optional[Dict[str, int]] = None,
        cost_info: Optional[Dict[str, float]] = None,
        conversation_simulator_version: Optional[str] = None,
    ) -> InferenceResult:
        """Create a new inference result log in a session.

        Args:
            session_id (str): The session ID to log the inference result to
            input (str): The input text/prompt
            output (str, optional): The generated output/response
            retrieval_context (Optional[str]): Context retrieved for RAG systems
            latency (float, optional): Latency in milliseconds since the model was called until
                the response was received.
            usage_info (dict[str, int], optional): Information about token usage during the
                model call.
                Possible keys include:
                - 'input_tokens': Number of input tokens sent to the model.
                - 'output_tokens': Number of output tokens generated by the model.
                - 'cache_read_input_tokens': Number of input tokens read from the cache.
            cost_info (dict[str, float], optional): Information about the cost per token during
                the model call.
                Possible keys include:
                - 'cost_per_input_token': Cost per input token sent to the model.
                - 'cost_per_output_token': Cost per output token generated by the model.
                - 'cost_per_cache_read_input_token': Cost per input token read from the cache.
            conversation_simulator_version (str, optional): The version of Galtea's conversation simulator
                used to generate the user message (input).
                This should only be provided if the input was generated using the conversation_simulator_service.

        Returns:
            InferenceResult: The created inference result log object
        """
        # Construct InferenceResultBase payload
        inference_result_base: InferenceResultBase = InferenceResultBase(
            session_id=session_id,
            actual_input=input,
            actual_output=output,
            retrieval_context=retrieval_context,
            latency=latency,
            **(cost_info or {}),
            **(usage_info or {}),
            conversation_simulator_version=conversation_simulator_version,
        )

        # Validate the payload
        inference_result_base.model_validate(inference_result_base.model_dump())

        # Send the request - could be to /sessions/{session_id}/inference_results or /inference_results
        response = self.__client.post(
            "inferenceResults",
            json=inference_result_base.model_dump(by_alias=True, exclude_none=True),
        )

        return InferenceResult(**response.json())

    def create_batch(
        self,
        session_id: str,
        conversation_turns: List[Dict[str, str]],
        conversation_simulator_version: Optional[str] = None,
    ) -> List[InferenceResult]:
        """Create a batch of inference result in a session in a single http call.

        Args:
            session_id (str): The session ID to log the inference results to
            conversation_turns (list[dict[str, str]], optional): Historic of the past chat
                conversation turns from the user and the model. Each turn is a dictionary with
                "role" and "content" keys following the standard conversation format.
                For instance:
                - [
                    {"role": "user", "content": "what is the capital of France?"},
                    {"role": "assistant", "content": "Paris"},
                    {"role": "user", "content": "what is the population of that city?"},
                    {"role": "assistant", "content": "2M"}
                ]
            conversation_simulator_version (str, optional): The version of Galtea's conversation simulator
                used to generate the user messages (inputs).
                This should only be provided if using the conversation_simulator_service to generate user messages.

        Returns:
            List[InferenceResult]: List of created inference result objects
        """
        if not is_valid_id(session_id):
            raise ValueError("A valid session ID must be provided.")

        if not conversation_turns or not isinstance(conversation_turns, list):
            raise ValueError("conversation_turns must be a non-empty list of dictionaries.")

        response = self.__client.post(
            "inferenceResults/batch",
            json={
                "conversationTurns": conversation_turns,
                "sessionId": session_id,
                "conversationSimulatorVersion": conversation_simulator_version,
            },
        )
        inference_results = [InferenceResult(**result) for result in response.json()]
        return inference_results

    def list(
        self,
        session_id: Union[str, list[str]],
        sort_by_created_at: Optional[str] = None,
        offset: Optional[int] = None,
        limit: Optional[int] = None,
    ) -> List[InferenceResult]:
        """List inference result logs for a session.

        Args:
            session_id (str | list[str]): The session ID or list of session IDs to get inference results from.
            sort_by_created_at (str, optional): Sort by created at. Valid values are 'asc' and 'desc'.
            offset (int, optional): Offset for pagination.
                This refers to the number of items to skip before starting to collect the result set.
                The default value is 0.
            limit (int, optional): Limit for pagination.
                This refers to the maximum number of items to collect in the result set.

        Returns:
            List[InferenceResult]: List of inference result log objects
        """
        # 1. Validate IDs filter parameters
        session_ids = [session_id] if isinstance(session_id, str) else session_id
        if not session_ids or not all(is_valid_id(session_id) for session_id in session_ids):
            raise ValueError("A valid session ID must be provided.")

        # 2. Validate sort parameters
        if sort_by_created_at is not None and sort_by_created_at not in ["asc", "desc"]:
            raise ValueError("Sort by created at must be 'asc' or 'desc'.")

        query_params = build_query_params(
            sessionIds=session_ids,
            offset=offset,
            limit=limit,
            sort=["createdAt", sort_by_created_at] if sort_by_created_at else None,
        )
        response = self.__client.get(f"inferenceResults?{query_params}")
        inference_results = [InferenceResult(**result) for result in response.json()]

        if not inference_results:
            for session_id in session_ids:
                session = self.__session_service.get(session_id)
                if not session:
                    raise ValueError(f"Session with ID {session_id} does not exist.")

        return inference_results

    def get(self, inference_result_id: str) -> InferenceResult:
        """Retrieve an inference result log by its ID.

        Args:
            inference_result_id (str): The ID of the inference result log to retrieve

        Returns:
            InferenceResult: The retrieved inference result log object
        """
        if not is_valid_id(inference_result_id):
            raise ValueError("A valid inference result ID must be provided.")

        response = self.__client.get(f"inferenceResults/{inference_result_id}")
        return InferenceResult(**response.json())

    def delete(self, inference_result_id: str) -> None:
        """Delete an inference result log by its ID.

        Args:
            inference_result_id (str): The ID of the inference result log to delete
        """
        if not is_valid_id(inference_result_id):
            raise ValueError("A valid inference result ID must be provided.")

        self.__client.delete(f"inferenceResults/{inference_result_id}")
