"""LLM rater for tasks with a LLM model server."""

import re
from collections import Counter
from typing import Any, Dict, List

from uniflow.op.model.constants import (
    AVERAGE_SCORE,
    MAJORITY_VOTE,
    RESPONSE,
    SCORES,
    VOTES,
)
from uniflow.op.model.llm_processor import JsonFormattedDataProcessor, LLMDataProcessor
from uniflow.op.prompt_schema import GuidedPrompt


class LLMRater(LLMDataProcessor):
    """LLM Rater"""

    def __init__(
        self,
        guided_prompt_template: GuidedPrompt,
        model_config: Dict[str, Any],
        label2score: Dict[str, float],
    ) -> None:
        """LLM Rater Constructor.

        Args:
            guided_prompt_template (GuidedPrompt): Guided prompt template.
            model_config (Dict[str, Any]): Model config.
            label2score (Dict[str, float]): String to score mapping.
        """
        super().__init__(guided_prompt_template, model_config)
        example_keys = list(guided_prompt_template.examples[0].dict().keys())
        pattern = r"^[^A-Za-z]+|[^A-Za-z]+$"
        self._rater_key = example_keys[-1]
        self._label2score = {
            re.sub(pattern, "", k).lower().lower(): float(v)
            for k, v in label2score.items()
        }
        self._score2label = {v: k for k, v in self._label2score.items()}

    def _deserialize(self, data: List[str]) -> List[Dict[str, Any]]:
        """Deserialize data.

        Args:
            data (List[str]): Data to deserialize.

        Returns:
            List[Dict[str, Any]]: Deserialized data.
        """

        def _extract_label(text):
            # Create a regular expression pattern from the dictionary keys
            # The pattern will look like '(yes|no)', case-insensitive
            pattern = (
                r"\b("
                + "|".join(  # noqa: W503
                    re.escape(key) for key in self._label2score.keys()
                )  # noqa: W503
                + r")\b"  # noqa: W503
            )

            matches = list(re.finditer(pattern, text, re.IGNORECASE))
            if matches:
                # Return the last matched group (Yes or No), preserving the original case
                return (
                    matches[-1].group(1).lower()
                )  # Convert to lowercase to match dictionary keys
            else:
                # Return None if no match is found
                return None

        data = super()._deserialize(data)
        response = data[RESPONSE]

        labels = [_extract_label(d) for d in response]
        scores = []
        for label in labels:
            if label is not None:
                scores.append(self._label2score[label])
        majority_vote = Counter(labels).most_common(1)[0][0]
        mean_score = sum(scores) / len(scores) if len(scores) > 0 else None

        data.update(
            {
                MAJORITY_VOTE: majority_vote,
                AVERAGE_SCORE: mean_score,
                VOTES: labels,
                SCORES: scores,
            }
        )

        return data


class JsonFormattedLLMRater(JsonFormattedDataProcessor):
    """Json Formatted LLM Rater"""

    def __init__(
        self,
        guided_prompt_template: GuidedPrompt,
        model_config: Dict[str, Any],
        label2score: Dict[str, float],
    ) -> None:
        """Json Formatted LLM Rater Constructor.

        Args:
            guided_prompt_template (GuidedPrompt): Guided prompt template.
            model_config (Dict[str, Any]): Model config.
            label2score (Dict[str, float]): String to score mapping.
        """
        super().__init__(guided_prompt_template, model_config)
        example_keys = list(guided_prompt_template.examples[0].dict().keys())
        self._pattern = r"^[^A-Za-z]+|[^A-Za-z]+$"
        self._rater_key = example_keys[-1]
        self._label2score = {
            re.sub(self._pattern, "", k).lower(): float(v)
            for k, v in label2score.items()
        }
        self._score2label = {v: k for k, v in self._label2score.items()}

    def _deserialize(self, data: List[str]) -> List[Dict[str, Any]]:
        """Deserialize data.

        Args:
            data (List[str]): Data to deserialize.

        Returns:
            List[Dict[str, Any]]: Deserialized data.
        """
        data = super()._deserialize(data)
        response = data[RESPONSE]

        labels = [
            re.sub(self._pattern, "", r[self._rater_key]).lower()
            if self._rater_key in r
            else None
            for r in response
        ]
        scores = []
        for label in labels:
            if label is not None and label in self._label2score:
                scores.append(self._label2score[label])
        majority_vote = Counter(labels).most_common(1)[0][0]
        mean_score = sum(scores) / len(scores) if len(scores) > 0 else None
        data.update(
            {
                MAJORITY_VOTE: majority_vote,
                AVERAGE_SCORE: mean_score,
                VOTES: labels,
                SCORES: scores,
            }
        )

        return data
