import logging
from abc import ABC, abstractmethod
from collections import Counter
from dataclasses import dataclass
from typing import Generic, Literal, TypeVar, cast

from evalassist.judges.utils import get_to_evaluate_text
from unitxt.inference import InferenceEngine

from .types import (
    TO_EVALUATE_FIELD_DEFAULT,
    Criteria,
    CriteriaOption,
    DirectInstanceResult,
    DirectPositionalBiasResult,
    Instance,
    MultiCriteria,
    MultiCriteriaDirectInstanceResult,
    PairwiseInstanceResult,
    PairwisePositionalBiasResult,
)

logger = logging.getLogger(__name__)


# ----------------------------------------------------------------------
# Core abstract judge definition
# ----------------------------------------------------------------------
InstanceTypeVar = TypeVar("InstanceTypeVar", bound=Instance)
ReturnVarType = TypeVar("ReturnVarType")


@dataclass
class JudgeDescriptor:
    name: str
    eval_type: Literal["direct", "pairwise"]
    inference_engine_id: str

    def __str__(self) -> str:
        return f"{self.name}-{self.eval_type}-{self.inference_engine_id}"


class UnitxtInferenceEngineMixin:
    inference_engine: InferenceEngine

    def __init__(
        self,
        inference_engine,
        *args,
        **kwargs,
    ):
        super().__init__(*args, **kwargs)
        self.inference_engine = inference_engine

    def get_inference_engine_id(self) -> str:
        """Return the identifier of the underlying inference engine."""
        return "_".join(self.inference_engine.get_engine_id().split("_")[:-1])


class BaseJudge(
    ABC,
    Generic[InstanceTypeVar, ReturnVarType],
):
    """
    Abstract base class for all judges.

    A *judge* evaluates one or more ``Instance`` objects against a set of
    ``Criteria`` and returns a result specific to the concrete implementation.
    """

    self_consistency: int | bool = False
    check_positional_bias: bool = False

    def __init__(
        self,
        self_consistency: bool | int = False,
        check_positional_bias: bool = False,
        *args,
        **kwargs,
    ):
        super().__init__(*args, **kwargs)
        if self_consistency < 0:
            raise ValueError(
                "self_consistency must have a boolean value or be an int greater than 0."
            )
        self.self_consistency = (
            self_consistency
            if isinstance(self_consistency, int)
            and not isinstance(self_consistency, bool)
            else (3 if self_consistency is True else 0)
        )

        self.check_positional_bias = check_positional_bias

    def get_ai_message_from_prompt(
        self, prompt: str, role: Literal["system", "user", "assistant"] = "user"
    ) -> dict[str, str]:
        return {
            "role": role,
            "content": prompt,
        }

    def _verify_fields(self, instances: list[Instance], criteria: list[Criteria]):
        for instance, criterion in zip(instances, criteria):
            if criterion.to_evaluate_field not in instance.fields:
                raise ValueError(
                    f"'{criterion.name}' criteria's to_evaluate_field '{criterion.to_evaluate_field}' is required in instance.field but was not found. Instance's fields are {list(instance.fields.keys())}"
                )
            if criterion.context_fields is not None:
                if not all(
                    field in instance.fields for field in criterion.context_fields
                ):
                    raise ValueError(
                        f"Instance doesn't have the fields that the criteria's context_fields expects. Criteria's context fields are: {list(criterion.context_fields)} while instance fields are has {list(instance.fields.keys())}"
                    )

    def evaluate(
        self,
        instances: list[InstanceTypeVar] | list[str] | list[list[str]],
        criteria: Criteria | list[Criteria] | str,
    ) -> list[ReturnVarType]:
        """Run the judge on a batch of instances and return the results."""
        if isinstance(criteria, list):
            if len(criteria) != len(instances):
                raise ValueError(
                    f"The provided criteria list must be equal in length with the instances. {len(criteria)} != {len(instances)}"
                )
            if len(criteria) == 0:
                raise ValueError("Criteria list is empty")

        elif not isinstance(criteria, str) and not isinstance(criteria, Criteria):
            raise ValueError(
                f"criteria parameter must be of type Criteria or str, you provided a {type(criteria)}"
            )

        parsed_criteria = self._get_parsed_criteria(
            cast(list[Criteria] | list[str], [criteria] * len(instances))
            if isinstance(criteria, str) or isinstance(criteria, Criteria)
            else criteria
        )
        parsed_instances: list[Instance] = self._get_instances_from_str(instances)

        self._verify_fields(parsed_instances, parsed_criteria)

        if self.self_consistency > 0:
            parsed_instances = [
                instance
                for instance in parsed_instances
                for _ in range(self.self_consistency)
            ]
            parsed_criteria = [
                criterion
                for criterion in parsed_criteria
                for _ in range(self.self_consistency)
            ]

        results: list[ReturnVarType] = self._evaluate(
            instances=parsed_instances,
            criteria=parsed_criteria,
        )
        # import json
        # result_str = json.dumps([result.model_dump() for result in results])
        # with open("output2.json", "w") as file:
        #     file.write(result_str)
        return results

    def __call__(
        self,
        instances: list[InstanceTypeVar] | list[str] | list[list[str]],
        criteria: Criteria | list[Criteria] | str,
    ) -> list[ReturnVarType]:
        return self.evaluate(
            instances=instances,
            criteria=criteria,
        )

    @abstractmethod
    def _evaluate(
        self,
        instances: list[Instance],
        criteria: list[Criteria],
    ) -> list[ReturnVarType]: ...

    @abstractmethod
    def _run(
        self,
        instances: list[InstanceTypeVar],
        criteria: list[Criteria],
    ) -> list[ReturnVarType]: ...

    @abstractmethod
    def _get_instances_from_str(
        self, instances: list[InstanceTypeVar] | list[str] | list[list[str]]
    ) -> list[Instance]: ...

    @abstractmethod
    def _get_parsed_criteria(
        self, criteria: list[Criteria] | list[str]
    ) -> list[Criteria]: ...

    # @abstractmethod
    # def get_predictions(self, instances: list[InstanceTypeVar]) -> Any:
    #     """Return the raw predictions (e.g., LLM responses) for the given instances."""
    #     ...

    @abstractmethod
    def get_descriptor(self) -> JudgeDescriptor:
        """Get an object with primary information of the judge"""
        ...

    @abstractmethod
    def get_name(self) -> str:
        """Get the name of the judge"""
        ...

    def __str__(self):
        return str(self.get_descriptor())


# ----------------------------------------------------------------------
# Concrete abstract subclasses for the two main evaluation modes
# ----------------------------------------------------------------------
class BaseDirectJudge(BaseJudge[Instance, DirectInstanceResult], ABC):
    def _evaluate(
        self,
        instances: list[Instance],
        criteria: list[Criteria],
    ) -> list[DirectInstanceResult]:
        if self.check_positional_bias:
            results: list[DirectInstanceResult] = self._run(
                instances=[*instances, *instances],
                criteria=[
                    *criteria,
                    *[
                        Criteria(
                            name=criterion.name,
                            description=criterion.description,
                            to_evaluate_field=criterion.to_evaluate_field,
                            context_fields=criterion.context_fields,
                            examples=criterion.examples,
                            options=list(reversed(criterion.options)),
                        )
                        for criterion in criteria
                    ],
                ],
            )

            results_len: int = len(instances)
            results = [
                DirectInstanceResult(
                    criteria=results[i].criteria,
                    instance=instances[i],
                    selected_option=results[i].selected_option,
                    explanation=results[i].explanation,
                    feedback=results[i].feedback,
                    metadata=results[i].metadata,
                    positional_bias=DirectPositionalBiasResult(
                        detected=results[i].selected_option
                        != results[i + results_len].selected_option,
                        result=results[results_len + i],
                    ),
                )
                for i in range(results_len)
            ]
        else:
            results = self._run(instances=instances, criteria=criteria)

        # add numeric scores from criteria if possible
        for r in results:
            score: float | None = (
                r.criteria.get_score_from_option(r.selected_option)
                if r.criteria is not None
                else None
            )
            if score is None:
                try:
                    # try to use the option name as the numeric score
                    score = float(r.selected_option)
                except (ValueError, TypeError):
                    pass
            r.score = score

        if self.self_consistency > 0:
            # apply majority voting for each of the self consistency evaluations
            parsed_results = []
            for i in range(0, len(results), self.self_consistency):
                selected_options = [
                    results[j].selected_option
                    for j in range(i, i + self.self_consistency)
                ]
                most_common_option = Counter(selected_options).most_common(1)[0][0]
                index_of_most_common = selected_options.index(most_common_option)
                to_update_result_index = i + index_of_most_common
                to_update_result = results[to_update_result_index]
                to_update_result.selected_option = most_common_option
                if all(
                    r.score is not None for r in results[i : i + self.self_consistency]
                ):
                    # set the mean of the scores as the score
                    to_update_result.score = (
                        sum(
                            cast(float, r.score)
                            for r in results[i : i + self.self_consistency]
                        )
                        / self.self_consistency
                    )

                to_update_result.metadata["self_consistency"] = {
                    "selected_options": selected_options,
                }
                parsed_results.append(results[to_update_result_index])
            return parsed_results

        return results

    def evaluate_multi_criteria(
        self,
        instances: list[Instance] | list[str],
        multi_criteria: MultiCriteria | list[str] | list[Criteria],
    ) -> list[MultiCriteriaDirectInstanceResult]:
        if isinstance(multi_criteria, list):
            criteria = self._get_parsed_criteria(multi_criteria)
            parsed_multi_criteria = MultiCriteria.from_criteria(criteria)
        else:
            parsed_multi_criteria = multi_criteria

        multi_criteria_items = parsed_multi_criteria.items
        criteria_count = len(multi_criteria_items)
        replicated_instances, replicated_criteria = zip(
            *[
                [instance, weighted_criterion.criterion]
                for instance in instances
                for weighted_criterion in multi_criteria_items
            ]
        )
        replicated_instances = list(replicated_instances)
        replicated_criteria = list(replicated_criteria)

        results = self.evaluate(replicated_instances, replicated_criteria)  # type: ignore

        final_results: list[MultiCriteriaDirectInstanceResult] = []
        for i in range(0, len(replicated_instances), criteria_count):
            criteria_results: list[DirectInstanceResult] = results[
                i : i + criteria_count
            ]
            final_results.append(parsed_multi_criteria.get_result(criteria_results))
        return final_results

    def _get_instances_from_str(
        self, instances: list[Instance] | list[str] | list[list[str]]
    ) -> list[Instance]:
        parsed_instances: list[Instance]
        if isinstance(instances, list) and all(
            isinstance(instance, str) for instance in instances
        ):
            parsed_instances = cast(
                list[Instance],
                [
                    Instance(
                        fields={TO_EVALUATE_FIELD_DEFAULT: instance},
                    )
                    for instance in cast(list[str], instances)
                ],
            )
        else:
            parsed_instances = cast(list[Instance], instances)
        return parsed_instances

    def _get_parsed_criteria(
        self, criteria: list[Criteria] | list[str]
    ) -> list[Criteria]:
        if isinstance(criteria, list) and all(isinstance(x, str) for x in criteria):
            return [
                Criteria(
                    name="",
                    description=description,
                    options=[
                        CriteriaOption(name="Yes", description="", score=1.0),
                        CriteriaOption(name="No", description="", score=0.0),
                    ],
                    to_evaluate_field="response",
                )
                for description in cast(list[str], criteria)
            ]
        else:
            return [
                Criteria(
                    name=criterion.name,
                    description=criterion.description,
                    options=criterion.options,
                    to_evaluate_field=criterion.to_evaluate_field,
                    context_fields=criterion.context_fields,
                    examples=criterion.examples,
                )
                for criterion in cast(list[Criteria], criteria)
            ]

    @abstractmethod
    def _run(
        self,
        instances: list[Instance],
        criteria: list[Criteria],
    ) -> list[DirectInstanceResult]: ...

    # def get_predictions(self, instances: list[Instance]) -> list[str]:
    #     return [i.response for i in instances]

    def get_descriptor(self) -> JudgeDescriptor:
        return JudgeDescriptor(self.get_name(), "direct", "")


class BasePairwiseJudge(BaseJudge[Instance, PairwiseInstanceResult], ABC):
    def __init__(
        self,
        *args,
        **kwargs,
    ):
        super().__init__(
            *args,
            **kwargs,
        )
        if self.self_consistency is True or self.self_consistency > 1:
            raise ValueError(
                "Self consistency is not supported on pairwise comparison judges yet."
            )

    def _evaluate(
        self,
        instances: list[Instance],
        criteria: list[Criteria],
    ) -> list[PairwiseInstanceResult]:
        if self.check_positional_bias:
            results: list[PairwiseInstanceResult] = self._run(
                instances=[
                    *instances,
                    *[
                        Instance(
                            fields={
                                **{
                                    k: v
                                    for k, v in instance.fields.items()
                                    if k != criterion.to_evaluate_field
                                },
                                **{
                                    k: list(reversed(cast(list[str], v)))
                                    for k, v in instance.fields.items()
                                    if k == criterion.to_evaluate_field
                                },
                            }
                        )
                        for instance, criterion in zip(instances, criteria)
                    ],
                ],
                criteria=[
                    *criteria,
                    *criteria,
                ],
            )

            results_len: int = int(
                len(results) / 2
            )  # half of positional bias result count, e.g. original number of instances to evaluate

            for (
                instance_result,
                positional_bias_instance_result,
                instance,
                criterion,
            ) in zip(results[:results_len], results[results_len:], instances, criteria):
                # update winner index if there is not a tie
                positional_bias_instance_result.selected_option = (
                    (
                        len(cast(list[str], get_to_evaluate_text(instance, criterion)))
                        - cast(int, positional_bias_instance_result.selected_option)
                        - 1
                    )
                    if positional_bias_instance_result.selected_option != "tie"
                    else "tie"
                )

                if (
                    instance_result.per_system_results is not None
                    and positional_bias_instance_result.per_system_results is not None
                ):
                    responses_count = len(instance_result.per_system_results)
                    for i, response_result in enumerate(
                        instance_result.per_system_results
                    ):
                        positional_bias_result_response = list(
                            positional_bias_instance_result.per_system_results
                        )[responses_count - i - 1]
                        response_result.positional_bias = [
                            a != b
                            for a, b in zip(
                                response_result.contest_results,
                                reversed(
                                    positional_bias_result_response.contest_results
                                ),
                            )
                        ]

                instance_result.positional_bias = PairwisePositionalBiasResult(
                    detected=instance_result.selected_option
                    != positional_bias_instance_result.selected_option,
                    result=positional_bias_instance_result,
                )

            return results[:results_len]
        else:
            return self._run(instances=instances, criteria=criteria)

    @abstractmethod
    def _run(
        self,
        instances: list[Instance],
        criteria: list[Criteria],
    ) -> list[PairwiseInstanceResult]: ...

    # def get_predictions(self, instances: list[Instance]) -> list[list[str]]:
    #     return [i.responses for i in instances]

    def get_descriptor(self) -> JudgeDescriptor:
        return JudgeDescriptor(self.get_name(), "pairwise", "")

    def _get_instances_from_str(
        self,
        instances: list[Instance] | list[str] | list[list[str]],
    ) -> list[Instance]:
        parsed_instances: list[Instance]
        if (
            isinstance(instances, list)
            and all(isinstance(x, list) for x in instances)
            and all(isinstance(y, str) for x in instances for y in x)
        ):
            parsed_instances = cast(
                list[Instance],
                [
                    Instance(
                        fields={TO_EVALUATE_FIELD_DEFAULT: instance},
                    )
                    for instance in cast(list[list[str]], instances)
                ],
            )
        elif all(isinstance(x, Instance) for x in instances):
            parsed_instances = cast(list[Instance], instances)
        else:
            raise ValueError(
                f"Invalid instance type. Type must be list[list[str]] or list[PairwiseInstance]. Receive {type(instances[0])}"
            )
        return parsed_instances

    def _get_parsed_criteria(
        self, criteria: list[Criteria] | list[str]
    ) -> list[Criteria]:
        if isinstance(criteria, list) and all(isinstance(x, str) for x in criteria):
            return [
                Criteria(
                    name="",
                    description=description,
                    to_evaluate_field=TO_EVALUATE_FIELD_DEFAULT,
                )
                for description in cast(list[str], criteria)
            ]
        else:
            return cast(list[Criteria], criteria)
