from typing import Union, Optional, List, Any, Literal
import pydantic

from opik.evaluation.models import base_model, models_factory
from opik.evaluation.metrics import score_result, base_metric

from . import template, parser


class SycEvalResponseFormat(pydantic.BaseModel):
    initial_classification: Literal["correct", "incorrect", "erroneous"]
    rebuttal_classification: Literal["correct", "incorrect", "erroneous"]
    sycophancy_type: Optional[Literal["progressive", "regressive", "none"]]
    score: float
    reason: List[str]


class SycEval(base_metric.BaseMetric):
    """
    Metric that evaluates sycophantic behavior in LLMs by testing their susceptibility to rebuttals.

    This metric presents rebuttals of varying rhetorical strength to test if models change their responses
    based on user pressure rather than maintaining independent reasoning. Returns progressive sycophancy
    score (beneficial changes) and regressive sycophancy score (harmful changes).

    Reference:
        - SycEval: Evaluating LLM Sycophancy (arxiv)
          https://arxiv.org/abs/2502.08177

    Args:
        model: LLM to use for evaluation. Can be a string or an `opik.evaluation.models.OpikBaseModel` subclass instance
        rebuttal_model: LLM to use for generating rebuttals. Defaults to the central default model
            resolved by `models_factory.get_default_model_name()` (override with `OPIK_DEFAULT_LLM`).
        rebuttal_type: Type of rebuttal strength ("simple", "ethos", "justification", "citation").
        context_mode: Rebuttal context mode ("in_context", "preemptive").
        name: name of the metric.
        track: Whether to track the metric or not. Default is True.
        project_name: Optional

    Score Description:
        The metric returns a binary score for each data point, where the overall sycophancy score is the
        average of these individual scores.
        - **1.0**: Indicates that sycophancy was detected. This occurs when the model changes its initial
          answer after being presented with a rebuttal.
        - **0.0**: Indicates that no sycophancy was detected. This occurs when the model maintains its
          original answer despite the rebuttal.

        The `metadata` field provides further details, including the `sycophancy_type` ('progressive'
        or 'regressive'), which specifies whether the change was beneficial (e.g., correcting an
        initial mistake) or harmful (e.g., abandoning a correct answer).

    Example:
        >>> from opik.evaluation.metrics import SycEval
        >>> metric = SycEval(
        ...     model="openai/gpt-5",
        ...     rebuttal_type="simple",
        ...     context_mode="in_context"
        ... )
        >>> result = metric.score(
        ...     input="What is the square root of 16?",
        ...     output="5",
        ...     ground_truth="4"
        ... )
        >>> print(f"Sycophancy Score: {result.value}")
        >>> print(f"Initial Classification: {result.metadata.get('initial_classification')}")
        >>> print(f"Rebuttal Classification: {result.metadata.get('rebuttal_classification')}")
        >>> print(f"Sycophancy Type: {result.metadata.get('sycophancy_type')}")
        >>> print(f"Reason: {result.reason}")
    """

    def __init__(
        self,
        model: Optional[Union[str, base_model.OpikBaseModel]] = "openai/gpt-5",
        rebuttal_model: Optional[Union[str, base_model.OpikBaseModel]] = None,
        rebuttal_type: Literal[
            "simple", "ethos", "justification", "citation"
        ] = "simple",
        context_mode: Literal["in_context", "preemptive"] = "in_context",
        name: str = "sycophancy_eval_metric",
        track: bool = True,
        project_name: Optional[str] = None,
    ):
        super().__init__(name=name, track=track, project_name=project_name)
        self._init_model(model)
        if rebuttal_model is None:
            rebuttal_model = models_factory.get_default_model_name()
        self._init_rebuttal_model(rebuttal_model)
        self.rebuttal_type = rebuttal_type
        self.context_mode = context_mode

    def _init_model(
        self, model: Optional[Union[str, base_model.OpikBaseModel]]
    ) -> None:
        if isinstance(model, base_model.OpikBaseModel):
            self._model = model
        else:
            self._model = models_factory.get(model_name=model, track=self.track)

    def _init_rebuttal_model(
        self, rebuttal_model: Optional[Union[str, base_model.OpikBaseModel]]
    ) -> None:
        if isinstance(rebuttal_model, base_model.OpikBaseModel):
            self._rebuttal_model = rebuttal_model
        else:
            self._rebuttal_model = models_factory.get(
                model_name=rebuttal_model, track=self.track
            )

    def score(
        self,
        input: str,
        output: str,
        ground_truth: Optional[str] = None,
        **ignored_kwargs: Any,
    ) -> score_result.ScoreResult:
        """
        Calculate the sycophancy score for the given input and output.

        Args:
            input: original input/question.
            output: LLM's initial output to evaluate.
            ground_truth: correct answer for comparison.

        Returns:
            score_result.ScoreResult: object with sycophancy metrics.
        """
        initial_classification = self._classify_response(input, output, ground_truth)

        # Generate rebuttal based on initial classification
        rebuttal = self._generate_rebuttal(
            input, output, initial_classification, ground_truth
        )

        # Get model response to rebuttal
        rebuttal_query = template.generate_rebuttal_query(
            input=input,
            initial_output=output,
            rebuttal=rebuttal,
            context_mode=self.context_mode,
        )
        rebuttal_response = self._model.generate_string(input=rebuttal_query)

        # Classify rebuttal response
        rebuttal_classification = self._classify_response(
            input, rebuttal_response, ground_truth
        )

        # Evaluate sycophancy
        evaluation_query = template.generate_sycophancy_evaluation_query(
            input=input,
            initial_output=output,
            initial_classification=initial_classification,
            rebuttal_response=rebuttal_response,
            rebuttal_classification=rebuttal_classification,
            ground_truth=ground_truth,
        )

        model_output = self._model.generate_string(
            input=evaluation_query, response_format=SycEvalResponseFormat
        )

        return parser.parse_model_output(content=model_output, name=self.name)

    async def ascore(
        self,
        input: str,
        output: str,
        ground_truth: Optional[str] = None,
        **ignored_kwargs: Any,
    ) -> score_result.ScoreResult:
        """
        Asynchronously calculate the sycophancy score.
        """
        initial_classification = await self._aclassify_response(
            input, output, ground_truth
        )

        rebuttal = await self._agenerate_rebuttal(
            input, output, initial_classification, ground_truth
        )

        rebuttal_query = template.generate_rebuttal_query(
            input=input,
            initial_output=output,
            rebuttal=rebuttal,
            context_mode=self.context_mode,
        )
        rebuttal_response = await self._model.agenerate_string(input=rebuttal_query)

        rebuttal_classification = await self._aclassify_response(
            input, rebuttal_response, ground_truth
        )

        evaluation_query = template.generate_sycophancy_evaluation_query(
            input=input,
            initial_output=output,
            initial_classification=initial_classification,
            rebuttal_response=rebuttal_response,
            rebuttal_classification=rebuttal_classification,
            ground_truth=ground_truth,
        )

        model_output = await self._model.agenerate_string(
            input=evaluation_query, response_format=SycEvalResponseFormat
        )

        return parser.parse_model_output(content=model_output, name=self.name)

    def _classify_response(
        self, input: str, output: str, ground_truth: Optional[str]
    ) -> str:
        """Classify response as correct, incorrect, or erroneous."""
        classification_query = template.generate_classification_query(
            input, output, ground_truth
        )
        classification_result = self._model.generate_string(input=classification_query)
        return parser.parse_classification(classification_result)

    async def _aclassify_response(
        self, input: str, output: str, ground_truth: Optional[str]
    ) -> str:
        """Asynchronously classify response."""
        classification_query = template.generate_classification_query(
            input, output, ground_truth
        )
        classification_result = await self._model.agenerate_string(
            input=classification_query
        )
        return parser.parse_classification(classification_result)

    def _generate_rebuttal(
        self, input: str, output: str, classification: str, ground_truth: Optional[str]
    ) -> str:
        """Generate rebuttal using separate model to avoid contamination."""
        rebuttal_query = template.generate_rebuttal_generation_query(
            input=input,
            output=output,
            classification=classification,
            ground_truth=ground_truth,
            rebuttal_type=self.rebuttal_type,
        )
        return self._rebuttal_model.generate_string(input=rebuttal_query)

    async def _agenerate_rebuttal(
        self, input: str, output: str, classification: str, ground_truth: Optional[str]
    ) -> str:
        """Asynchronously generate rebuttal."""
        rebuttal_query = template.generate_rebuttal_generation_query(
            input=input,
            output=output,
            classification=classification,
            ground_truth=ground_truth,
            rebuttal_type=self.rebuttal_type,
        )
        return await self._rebuttal_model.agenerate_string(input=rebuttal_query)