import re
from dataclasses import dataclass
from string import Formatter
from typing import Callable, List, Mapping, Optional, Tuple, Union

import pandas as pd

from phoenix.evals.exceptions import PhoenixException
from phoenix.evals.utils import NOT_PARSABLE

DEFAULT_START_DELIM = "{"
DEFAULT_END_DELIM = "}"


@dataclass
class PromptOptions:
    provide_explanation: bool = False


class InvalidClassificationTemplateError(PhoenixException):
    pass


class PromptTemplate:
    template: str
    variables: List[str]

    def __init__(
        self,
        template: str,
        delimiters: Tuple[str, str] = (DEFAULT_START_DELIM, DEFAULT_END_DELIM),
    ):
        self.template = template
        self._start_delim, self._end_delim = delimiters
        self.variables = self._parse_variables(self.template)

    def prompt(self, options: Optional[PromptOptions] = None) -> str:
        return self.template

    def format(
        self,
        variable_values: Mapping[str, Union[bool, int, float, str]],
        options: Optional[PromptOptions] = None,
    ) -> str:
        prompt = self.prompt(options)
        if self._start_delim == "{" and self._end_delim == "}":
            prompt = prompt.format(**variable_values)
        else:
            for variable_name in self.variables:
                prompt = prompt.replace(
                    self._start_delim + variable_name + self._end_delim,
                    str(variable_values[variable_name]),
                )
        return prompt

    def _parse_variables(self, text: str) -> List[str]:
        if self._start_delim == "{" and self._end_delim == "}":
            formatter = Formatter()
            variables = [field_name for _, field_name, _, _ in formatter.parse(text) if field_name]
            return variables
        else:
            start = re.escape(self._start_delim)
            end = re.escape(self._end_delim)
            pattern = rf"{start}(.*?){end}"
            variables = re.findall(pattern, text)
            return variables


class ClassificationTemplate(PromptTemplate):
    def __init__(
        self,
        rails: List[str],
        template: str,
        explanation_template: Optional[str] = None,
        explanation_label_parser: Optional[Callable[[str], str]] = None,
        delimiters: Tuple[str, str] = (DEFAULT_START_DELIM, DEFAULT_END_DELIM),
        scores: Optional[List[float]] = None,
    ):
        if scores is not None and len(rails) != len(scores):
            raise InvalidClassificationTemplateError(
                "If scores are provided, each rail must have one and only one score "
                "(i.e., the length of both lists must be the same)."
            )
        self.rails = rails
        self.template = template
        self.explanation_template = explanation_template
        self.explanation_label_parser = explanation_label_parser
        self._start_delim, self._end_delim = delimiters
        self.variables: List[str] = []
        for text in [template, explanation_template]:
            if text is not None:
                self.variables += self._parse_variables(text)
        self._scores = scores

    def __repr__(self) -> str:
        return self.template

    def prompt(self, options: Optional[PromptOptions] = None) -> str:
        if options is None:
            return self.template

        if options.provide_explanation and self.explanation_template:
            return self.explanation_template
        else:
            return self.template

    def extract_label_from_explanation(self, raw_string: str) -> str:
        if parser := self.explanation_label_parser:
            return parser(raw_string)
        return parse_label_from_chain_of_thought_response(raw_string)

    def score(self, rail: str) -> float:
        if self._scores is None:
            return 0.0
        try:
            return self._scores[self.rails.index(rail)]
        except (IndexError, ValueError):
            return 0.0


def parse_label_from_chain_of_thought_response(raw_string: str) -> str:
    label_delimiter = r"\W*label\W*"
    parts = re.split(label_delimiter, raw_string, maxsplit=1, flags=re.IGNORECASE)
    if len(parts) == 2:
        return parts[1]
    return NOT_PARSABLE


def normalize_classification_template(
    rails: List[str], template: Union[PromptTemplate, ClassificationTemplate, str]
) -> ClassificationTemplate:
    """
    Normalizes a template to a ClassificationTemplate object.
    Args:
        template (Union[ClassificationTemplate, str]): The template to be normalized.
    Returns:
        ClassificationTemplate: The normalized template.
    """
    if isinstance(template, ClassificationTemplate):
        return template

    if isinstance(template, PromptTemplate):
        return ClassificationTemplate(rails=rails, template=template.template)

    if isinstance(template, str):
        return ClassificationTemplate(rails=rails, template=template)

    raise TypeError(
        "Invalid type for argument `template`. Expected a string or ClassificationTemplate "
        f"but found {type(template)}."
    )


def normalize_prompt_template(template: Union[PromptTemplate, str]) -> PromptTemplate:
    """
    Normalizes a template to a PromptTemplate object.
    Args:
        template (Union[PromptTemplate, str]): The template to be normalized.
    Returns:
        PromptTemplate: The normalized template.
    """
    if isinstance(template, PromptTemplate):
        return template

    if isinstance(template, str):
        return PromptTemplate(template=template)

    raise TypeError(
        "Invalid type for argument `template`. Expected a string or PromptTemplate "
        f"but found {type(template)}."
    )


def map_template(
    dataframe: pd.DataFrame,
    template: PromptTemplate,
    options: Optional[PromptOptions] = None,
) -> "pd.Series[str]":
    """
    Maps over a dataframe to construct a list of prompts from a template and a dataframe.
    """
    # Was considering to construct the prompts and generate answers concurrently. However,
    # if there's errors in the prompt construction it could interrupt the process and we
    # would've used API credits for nothing. We could solve this problem by streaming the
    # answers so that, if there is an error, we keep the answers obtained up to that point.
    # These are out of scope for M0, but good to keep in mind and consider for the future.
    prompt_options: PromptOptions = PromptOptions() if options is None else options

    try:
        prompts = dataframe.apply(
            lambda row: template.format(
                variable_values={var_name: row[var_name] for var_name in template.variables},
                options=prompt_options,
            ),
            axis=1,
        )
        return prompts
    except KeyError as e:
        raise RuntimeError(
            f"Error while constructing the prompts from the template and dataframe. "
            f"The template variable {e} is not found as a column in the dataframe."
        )
    except Exception as e:
        raise RuntimeError(
            f"Error while constructing the prompts from the template and dataframe variables: {e}."
        )
