from typing import Dict, List, Optional, Sequence, TypedDict, Union

from galtea.application.services.inference_result_service import InferenceResultService
from galtea.application.services.metric_service import MetricService
from galtea.application.services.session_service import SessionService
from galtea.application.services.test_case_service import TestCaseService
from galtea.application.services.test_service import TestService
from galtea.application.services.version_service import VersionService
from galtea.domain.exceptions.entity_not_found_exception import EntityNotFoundException
from galtea.domain.models.evaluation import Evaluation
from galtea.domain.models.inference_result import CostInfoProperties, UsageInfoProperties
from galtea.domain.models.session import Session
from galtea.domain.models.test_case import TestCase
from galtea.infrastructure.clients.http_client import Client
from galtea.utils.custom_score_metric import CustomScoreEvaluationMetric
from galtea.utils.string import build_query_params, is_valid_id


class MetricInput(TypedDict, total=False):
    """
    Structured input format for specifying metrics when creating evaluations.

    This format provides flexibility by allowing you to identify metrics by ID or name,
    and optionally provide pre-calculated scores or custom scoring functions for self-hosted metrics.

    Attributes:
        id (Optional[str]): The unique identifier of the metric. Optional if name is provided.
            Must NOT be provided when score is a CustomScoreEvaluationMetric (the metric
            identifier is specified in the CustomScoreEvaluationMetric instance itself).
        name (Optional[str]): The name of the metric. Optional if id is provided.
            Must NOT be provided when score is a CustomScoreEvaluationMetric (the metric
            identifier is specified in the CustomScoreEvaluationMetric instance itself).
        score (Optional[Union[float, CustomScoreEvaluationMetric]]): Score for self-hosted metrics.
            - If float: Pre-calculated score between 0 and 1. Requires either 'id' or 'name'.
            - If CustomScoreEvaluationMetric: Custom scoring function that will be executed
              to calculate the score dynamically. The CustomScoreEvaluationMetric instance
              must contain either 'name' or 'id' to identify the metric. Do NOT provide
              'id' or 'name' in the MetricInput dict when using CustomScoreEvaluationMetric.
    """

    id: str
    name: str
    score: Union[float, CustomScoreEvaluationMetric]


class EvaluationService:
    """
    Service for managing Evaluations.
    An Evaluation is the result of an evaluation against a specific metric and its criteria.
    """

    def __init__(
        self,
        client: Client,
        metric_service: MetricService,
        session_service: SessionService,
        test_case_service: TestCaseService,
        test_service: TestService,
        version_service: VersionService,
        inference_result_service: InferenceResultService,
    ):
        self.__client = client
        self.__metric_service = metric_service
        self.__session_service = session_service
        self.__test_case_service = test_case_service
        self.__test_service = test_service
        self.__version_service = version_service
        self.__inference_result_service = inference_result_service

    def __process_metrics_with_score_calculation(
        self,
        metrics: Sequence[Union[str, CustomScoreEvaluationMetric, MetricInput]],
        input_data: Optional[str],
        actual_output: Optional[str],
        expected_output: Optional[str],
        retrieval_context: Optional[str],
        context: Optional[str],
    ) -> List[Dict[str, Union[str, float]]]:
        """
        Processes metrics from various input formats and calculates scores for CustomScoreEvaluationMetric.

        This method ensures backward compatibility by accepting multiple input formats:
        - Strings (metric names) - deprecated format, maintained for legacy support
        - CustomScoreEvaluationMetric objects - deprecated format with dynamic score calculation
        - MetricInput dicts - current recommended format

        Args:
            metrics: List of metrics in any supported format.
            input_data: The input text or prompt.
            actual_output: The actual output generated by the system.
            expected_output: The expected output for comparison.
            retrieval_context: Context retrieved by a RAG system, if applicable.
            context: Additional context for the evaluation.

        Returns:
            List of metric dictionaries in standardized API format: [{ id?, name?, score? }, ...]
        """
        api_metrics: List[Dict[str, Union[str, float]]] = []

        for metric_obj in metrics:
            if isinstance(metric_obj, str):
                # Deprecated format: string metric name
                api_metrics.append({"name": metric_obj})
            elif isinstance(metric_obj, CustomScoreEvaluationMetric):
                # Deprecated format: CustomScoreEvaluationMetric - need to calculate score
                try:
                    score = metric_obj(
                        input=input_data,
                        actual_output=actual_output,
                        expected_output=expected_output,
                        retrieval_context=retrieval_context,
                        context=context,
                    )
                    # CustomScoreEvaluationMetric can have either name or id
                    metric_identifier = {}
                    if metric_obj.name is not None:
                        metric_identifier["name"] = metric_obj.name
                    elif metric_obj.id is not None:
                        metric_identifier["id"] = metric_obj.id

                    api_metrics.append({**metric_identifier, "score": score})
                except Exception as e:
                    identifier = metric_obj.name or metric_obj.id or "unknown"
                    raise ValueError(f"Failed to generate score for custom metric '{identifier}': {e}") from e
            elif isinstance(metric_obj, dict):
                # Current format: MetricInput dict
                metric_dict: Dict[str, Union[str, float]] = {}

                # Validate that when using CustomScoreEvaluationMetric as score,
                # id and name must NOT be provided (they're in the CustomScoreEvaluationMetric itself)
                if "score" in metric_obj and isinstance(metric_obj["score"], CustomScoreEvaluationMetric):
                    if "id" in metric_obj or "name" in metric_obj:
                        raise ValueError(
                            "When using CustomScoreEvaluationMetric as score, do NOT provide 'id' or 'name' "
                            "in the MetricInput dict. The metric identifier should be specified in the "
                            "CustomScoreEvaluationMetric instance itself (via its 'name' or 'id' parameter)."
                        )

                if "id" in metric_obj:
                    metric_dict["id"] = metric_obj["id"]
                if "name" in metric_obj:
                    metric_dict["name"] = metric_obj["name"]

                if "score" in metric_obj:
                    score_value = metric_obj["score"]
                    if isinstance(score_value, CustomScoreEvaluationMetric):
                        # Need to calculate the score dynamically
                        try:
                            calculated_score = score_value(
                                input=input_data,
                                actual_output=actual_output,
                                expected_output=expected_output,
                                retrieval_context=retrieval_context,
                                context=context,
                            )
                            # Add metric identifier from CustomScoreEvaluationMetric
                            if score_value.name is not None:
                                metric_dict["name"] = score_value.name
                            elif score_value.id is not None:
                                metric_dict["id"] = score_value.id

                            metric_dict["score"] = calculated_score
                        except Exception as e:
                            identifier = score_value.name or score_value.id or "unknown"
                            raise ValueError(f"Failed to generate score for custom metric '{identifier}': {e}") from e
                    else:
                        # Pre-calculated score (float)
                        metric_dict["score"] = score_value

                api_metrics.append(metric_dict)
            else:
                raise TypeError(
                    f"Invalid metric type: {type(metric_obj)}. Expected str, CustomScoreEvaluationMetric, or dict."
                )

        return api_metrics

    def create_single_turn(
        self,
        version_id: str,
        actual_output: str,
        metrics: Optional[Sequence[Union[str, CustomScoreEvaluationMetric, MetricInput]]] = None,
        test_case_id: Optional[str] = None,
        input: Optional[str] = None,
        is_production: Optional[bool] = None,
        retrieval_context: Optional[str] = None,
        latency: Optional[float] = None,
        usage_info: Optional[Dict[str, float]] = None,
        cost_info: Optional[Dict[str, float]] = None,
        conversation_simulator_version: Optional[str] = None,
    ) -> Optional[List[Evaluation]]:
        """
        Creates evaluation for a single-turn interaction, assessing system performance based on
        specified metrics. A separate evaluation is generated for each metric provided.

        This method supports both current and deprecated input formats for metrics, ensuring
        backward compatibility with existing integrations while encouraging migration to the
        recommended MetricInput format.

        Args:
            version_id (str): The unique identifier of the version being evaluated.
            actual_output (str): The actual output generated by the system under evaluation.
            metrics (Optional[List[Union[str, CustomScoreEvaluationMetric, MetricInput]]]):
                List of metrics to evaluate against. Supports multiple formats for flexibility:
                - Strings: Metric names (deprecated, maintained for backward compatibility)
                - CustomScoreEvaluationMetric: Objects with dynamic score calculation (deprecated,
                  maintained for backward compatibility). Must be initialized with either 'name'
                  or 'id' parameter.
                - MetricInput dicts: Recommended format with optional id, name, and score.
                  The score field can be either:
                  * float: Pre-calculated score. Requires 'id' or 'name' in the dict.
                    Example: {'name': 'accuracy', 'score': 0.95}
                  * CustomScoreEvaluationMetric: Object for dynamic score calculation.
                    The CustomScoreEvaluationMetric must be initialized with 'name' or 'id'.
                    Do NOT provide 'id' or 'name' in the MetricInput dict itself.
                    Example: {'score': CustomScoreEvaluationMetric(name='custom_metric')}
            test_case_id (Optional[str]): The ID of the test case used for evaluation. Links to
                predefined inputs, expected outputs, and context. Required for test-based evaluation
                (when is_production is False or omitted).
            input (Optional[str]): The input text or prompt for the evaluation. Only used when
                is_production is True to track production data without a predefined test case.
            is_production (Optional[bool]): Indicates whether this evaluation is from production data.
                Default is False. When True, test_case_id is not required but input becomes required.
            retrieval_context (Optional[str]): Context retrieved by a RAG system, if applicable.
                Useful for evaluations involving retrieval-augmented generation.
            latency (Optional[float]): The latency in milliseconds from model invocation to response
                receipt.
            usage_info (Optional[Dict[str, float]]): Token usage information from the model call.
                Supported 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 cached input tokens used
            cost_info (Optional[Dict[str, float]]): Cost breakdown for the model call.
                Supported keys include:
                - 'cost_per_input_token': Cost per input token
                - 'cost_per_output_token': Cost per output token
                - 'cost_per_cache_read_input_token': Cost per cached input token
            conversation_simulator_version (Optional[str]): Version of Galtea's conversation
                simulator used to generate the input. Include this when the input was produced
                by the conversation_simulator_service.

        Returns:
            Optional[List[Evaluation]]: A list of created Evaluation objects, one per metric.
                Returns None if an error occurs during evaluation creation.

        Raises:
            ValueError: If required parameters are missing or invalid.
            KeyError: If usage_info or cost_info contain unsupported keys.

        Example usage with deprecated format (still supported):
            >>> evaluations = galtea.evaluations.create_single_turn(
            ...     test_case_id="tc_123",
            ...     version_id="v_456",
            ...     actual_output="Model response...",
            ...     metrics=["accuracy", "precision"],  # Deprecated but still works
            ... )

        Example usage with recommended format (pre-calculated score):
            >>> evaluations = galtea.evaluations.create_single_turn(
            ...     test_case_id="tc_123",
            ...     version_id="v_456",
            ...     actual_output="Model response...",
            ...     metrics=[
            ...         {"name": "accuracy"},
            ...         {"id": "metric_xyz", "score": 0.95},  # Pre-calculated score
            ...     ],
            ... )

        Example usage with CustomScoreEvaluationMetric in MetricInput:
            >>> from galtea.utils.custom_score_metric import CustomScoreEvaluationMetric
            >>>
            >>> class MyCustomMetric(CustomScoreEvaluationMetric):
            ...     def measure(self, input, actual_output, expected_output, retrieval_context, context):
            ...         return 0.85  # Your custom logic here
            >>>
            >>> # Initialize with name or id - NOT both
            >>> custom_metric = MyCustomMetric(name="custom_metric")
            >>> # OR: custom_metric = MyCustomMetric(id="metric_xyz")
            >>>
            >>> evaluations = galtea.evaluations.create_single_turn(
            ...     test_case_id="tc_123",
            ...     version_id="v_456",
            ...     actual_output="Model response...",
            ...     metrics=[
            ...         {"name": "accuracy"},
            ...         # Note: No 'id' or 'name' in dict when using CustomScoreEvaluationMetric
            ...         {"score": custom_metric},  # Dynamic score calculation
            ...     ],
            ... )
        """

        metrics = metrics or []

        # Fetch test case if needed for score calculation
        test_case: Union[TestCase, None] = None
        if test_case_id is not None:
            test_case = self.__test_case_service.get(test_case_id)

        # Process metrics and calculate scores for CustomScoreEvaluationMetric
        api_metrics = self.__process_metrics_with_score_calculation(
            metrics=metrics,
            input_data=input,
            actual_output=actual_output,
            expected_output=test_case.expected_output if test_case else None,
            retrieval_context=retrieval_context,
            context=test_case.context if test_case else None,
        )

        if usage_info is not None:
            for key, _ in usage_info.items():
                if key not in UsageInfoProperties.model_fields:
                    raise KeyError(
                        f"Invalid key: {key}. Must be one of: {', '.join(UsageInfoProperties.model_fields.keys())}"
                    )

        if cost_info is not None:
            for key, _ in cost_info.items():
                if key not in CostInfoProperties.model_fields:
                    raise KeyError(
                        f"Invalid key: {key}. Must be one of: {', '.join(CostInfoProperties.model_fields.keys())}"
                    )

        request_body = {
            "metrics": api_metrics,
            "versionId": version_id,
            "testCaseId": test_case_id,
            "actualOutput": actual_output,
            "input": input,
            "retrievalContext": retrieval_context,
            "latency": latency,
            "usageInfo": usage_info,
            "costInfo": cost_info,
            "isProduction": is_production,
            "conversationSimulatorVersion": conversation_simulator_version,
        }

        response = self.__client.post("evaluations/singleTurn", json=request_body)
        evaluations = [Evaluation(**evaluation) for evaluation in response.json()]
        return evaluations

    def create(
        self,
        session_id: str,
        metrics: Optional[Sequence[Union[str, CustomScoreEvaluationMetric, MetricInput]]] = None,
    ) -> List[Evaluation]:
        """
        Creates evaluations for an existing session, assessing system performance using
        specified metrics. A separate evaluation is generated for each metric.

        This method supports both current and deprecated input formats for metrics, ensuring
        backward compatibility with existing integrations while encouraging migration to the
        recommended MetricInput format.

        Args:
            session_id (str): The unique identifier of the session to evaluate.
            metrics (Optional[List[Union[str, CustomScoreEvaluationMetric, MetricInput]]]):
                List of metrics to evaluate against. Supports multiple formats for flexibility:
                - Strings: Metric names (deprecated, maintained for backward compatibility)
                - CustomScoreEvaluationMetric: Objects with dynamic score calculation (deprecated,
                  maintained for backward compatibility). Must be initialized with either 'name'
                  or 'id' parameter.
                - MetricInput dicts: Recommended format with optional id, name, and score.
                  The score field can be either:
                  * float: Pre-calculated score. Requires 'id' or 'name' in the dict.
                    Example: {'name': 'accuracy', 'score': 0.95}
                  * CustomScoreEvaluationMetric: Object for dynamic score calculation.
                    The CustomScoreEvaluationMetric must be initialized with 'name' or 'id'.
                    Do NOT provide 'id' or 'name' in the MetricInput dict itself.
                    Example: {'score': CustomScoreEvaluationMetric(name='custom_metric')}

        Returns:
            List[Evaluation]: A list of created Evaluation objects, one per metric.

        Raises:
            ValueError: If session_id is invalid or other validation fails.

        Example usage with deprecated format (still supported):
            >>> evaluations = galtea.evaluations.create(
            ...     session_id="session_123",
            ...     metrics=["accuracy", "precision"],  # Deprecated but still works
            ... )

        Example usage with recommended format (pre-calculated score):
            >>> evaluations = galtea.evaluations.create(
            ...     session_id="session_123",
            ...     metrics=[
            ...         {"name": "accuracy"},
            ...         {"id": "metric_xyz", "score": 0.95},  # Pre-calculated score
            ...     ],
            ... )

        Example usage with CustomScoreEvaluationMetric in MetricInput:
            >>> from galtea.utils.custom_score_metric import CustomScoreEvaluationMetric
            >>>
            >>> class MyCustomMetric(CustomScoreEvaluationMetric):
            ...     def measure(self, input, actual_output, expected_output, retrieval_context, context):
            ...         return 0.85  # Your custom logic here
            >>>
            >>> # Initialize with name or id - NOT both
            >>> custom_metric = MyCustomMetric(name="custom_metric")
            >>> # OR: custom_metric = MyCustomMetric(id="metric_xyz")
            >>>
            >>> evaluations = galtea.evaluations.create(
            ...     session_id="session_123",
            ...     metrics=[
            ...         {"name": "accuracy"},
            ...         # Note: No 'id' or 'name' in dict when using CustomScoreEvaluationMetric
            ...         {"score": custom_metric},  # Dynamic score calculation
            ...     ],
            ... )
        """

        if not is_valid_id(session_id):
            raise ValueError("Session ID provided is not valid.")

        metrics = metrics or []

        # Fetch session and related data for score calculation
        session: Union[Session, None] = None
        test_case: Union[TestCase, None] = None
        latest_result: Optional[Dict] = None

        # Only fetch data if we have metrics that might need it
        if metrics:
            session = self.__session_service.get(session_id)

            # Get test case if available
            if session.test_case_id is not None:
                test_case = self.__test_case_service.get(session.test_case_id)

            # Get the latest inference result for this session
            inference_results = self.__inference_result_service.list(
                session_id=session_id,
                sort_by_created_at="desc",
                limit=1,
            )

            if not inference_results:
                raise ValueError(f"No inference results found for session '{session_id}'")

            latest_result_obj = inference_results[0]
            # Convert to dict format for compatibility with existing code
            latest_result = latest_result_obj.model_dump()

        # Process metrics and calculate scores for CustomScoreEvaluationMetric
        api_metrics = self.__process_metrics_with_score_calculation(
            metrics=metrics,
            input_data=latest_result.get("actualInput") if latest_result else None,
            actual_output=latest_result.get("actualOutput") if latest_result else None,
            expected_output=test_case.expected_output if test_case else None,
            retrieval_context=latest_result.get("retrievalContext") if latest_result else None,
            context=test_case.context if test_case else None,
        )

        request_body = {
            "metrics": api_metrics,
            "sessionId": session_id,
        }

        response = self.__client.post("evaluations/fromSession", json=request_body)
        evaluations = [Evaluation(**evaluation) for evaluation in response.json()]
        return evaluations

    def get(self, evaluation_id: str):
        """
        Retrieve an evaluation by its ID.

        Args:
            evaluation_id (str): ID of the evaluation to retrieve.

        Returns:
            Evaluation: The retrieved evaluation object.
        """
        if not is_valid_id(evaluation_id):
            raise ValueError("Evaluation ID provided is not valid.")

        response = self.__client.get(f"evaluations/{evaluation_id}")
        return Evaluation(**response.json())

    def list(
        self,
        session_id: Optional[Union[str, list[str]]] = None,
        version_id: Optional[str] = None,
        test_case_id: Optional[Union[str, list[str]]] = None,
        test_id: Optional[Union[str, list[str]]] = None,
        metric_id: Optional[Union[str, list[str]]] = None,
        sort_by_created_at: Optional[str] = None,
        sort_by_score: Optional[str] = None,
        offset: Optional[int] = None,
        limit: Optional[int] = None,
    ):
        """
        Get a list of evaluations for a given session.

        Args:
            session_id (str | list[str], optional): ID or IDs of the sessions associated with the
                evaluations to retrieve.
            version_id (str, optional): ID of the version associated with the evaluations to retrieve.
            test_case_id (str | list[str], optional): ID or IDs of the test cases associated with the
                evaluations to retrieve.
            test_id (str | list[str], optional): ID or IDs of the tests associated with the
                evaluations to retrieve.
            metric_id (str | list[str], optional): ID or IDs of the metrics associated with the
                evaluations to retrieve.
            sort_by_created_at (str, optional): Sort by created at. Valid values are 'asc' and 'desc'.
            sort_by_score (str, optional): Sort by score. 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[Evaluation]: List of evaluations.
        """
        # 1. Validate IDs filter parameters
        if version_id is not None and not is_valid_id(version_id):
            raise ValueError("version_id parameter has to be a string with a valid ID.")

        # 1.1 Normalize single string inputs to lists
        session_ids = [session_id] if isinstance(session_id, str) else session_id
        version_ids = [version_id] if isinstance(version_id, str) else version_id
        test_case_ids = [test_case_id] if isinstance(test_case_id, str) else test_case_id
        test_ids = [test_id] if isinstance(test_id, str) else test_id
        metric_ids = [metric_id] if isinstance(metric_id, str) else metric_id
        if not (session_ids or version_ids or test_case_ids or test_ids or metric_ids):
            raise ValueError("At least one ID filter parameter must be provided.")

        # 1.2 Validate IDs
        if session_ids is not None and not all(is_valid_id(session_id) for session_id in session_ids):
            raise ValueError("Session ID provided is not valid.")
        if version_ids is not None and not all(is_valid_id(version_id) for version_id in version_ids):
            raise ValueError("Version ID provided is not valid.")
        if test_case_ids is not None and not all(is_valid_id(test_case_id) for test_case_id in test_case_ids):
            raise ValueError("Test case ID provided is not valid.")
        if test_ids is not None and not all(is_valid_id(test_id) for test_id in test_ids):
            raise ValueError("Test ID provided is not valid.")
        if metric_ids is not None and not all(is_valid_id(metric_id) for metric_id in metric_ids):
            raise ValueError("Metric ID provided is not valid.")

        # 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'.")
        if sort_by_score is not None and sort_by_score not in ["asc", "desc"]:
            raise ValueError("Sort by score must be 'asc' or 'desc'.")

        query_params = build_query_params(
            sessionIds=session_ids,
            versionIds=version_ids,
            testCaseIds=test_case_ids,
            testIds=test_ids,
            metricIds=metric_ids,
            offset=offset,
            limit=limit,
            sort=[
                *(["createdAt", sort_by_created_at] if sort_by_created_at else []),
                *(["score", sort_by_score] if sort_by_score else []),
            ],
        )
        response = self.__client.get(f"evaluations?{query_params}")
        evaluations = [Evaluation(**evaluation) for evaluation in response.json()]

        if not evaluations:
            if session_ids:
                try:
                    for session_id in session_ids:
                        self.__session_service.get(session_id)
                except Exception:
                    raise EntityNotFoundException(f"Session with ID {session_id} does not exist.")
            if version_ids:
                try:
                    for version_id in version_ids:
                        self.__version_service.get(version_id)
                except Exception:
                    raise EntityNotFoundException(f"Version with ID {version_id} does not exist.")
            if test_case_ids:
                try:
                    for test_case_id in test_case_ids:
                        self.__test_case_service.get(test_case_id)
                except Exception:
                    raise EntityNotFoundException(f"Test case with ID {test_case_id} does not exist.")
            if test_ids:
                try:
                    for test_id in test_ids:
                        self.__test_service.get(test_id)
                except Exception:
                    raise EntityNotFoundException(f"Test with ID {test_id} does not exist.")
            if metric_ids:
                try:
                    for metric_id in metric_ids:
                        self.__metric_service.get(metric_id)
                except Exception:
                    raise EntityNotFoundException(f"Metric with ID {metric_id} does not exist.")

        return evaluations

    def delete(self, evaluation_id: str):
        """
        Delete an evaluation by its ID.

        Args:
            evaluation_id (str): ID of the evaluation to delete.

        Returns:
            None.
        """
        if not is_valid_id(evaluation_id):
            raise ValueError("Evaluation ID provided is not valid.")

        self.__client.delete(f"evaluations/{evaluation_id}")
