import json
from abc import ABC
from dataclasses import dataclass, field
from itertools import product
from pathlib import Path
from types import MappingProxyType
from typing import Any, Callable, Dict, List, Mapping, Optional, Sequence, Set, Tuple, Type, Union
from uuid import UUID, uuid4

import pandas as pd
from pandas.api.types import is_integer_dtype, is_numeric_dtype, is_string_dtype
from pyarrow import Table, parquet

from phoenix.config import TRACE_DATASET_DIR
from phoenix.exceptions import PhoenixException

EVAL_NAME_COLUMN_PREFIX = "eval."


class InvalidParquetMetadataError(PhoenixException):
    pass


class NeedsNamedIndex(ABC):
    index_names: Mapping[Tuple[str, ...], Callable[[Any], bool]]
    all_valid_index_name_sorted_combos: Set[Tuple[str, ...]]

    @classmethod
    def preferred_names(cls) -> List[str]:
        return [choices[0] for choices in cls.index_names.keys()]

    @classmethod
    def aliases(cls) -> Mapping[str, str]:
        return {alias: choices[0] for choices in cls.index_names.keys() for alias in choices[1:]}

    @classmethod
    def unalias(cls, name: str) -> str:
        return cls.aliases().get(name, name)

    @classmethod
    def is_valid_index_names(cls, names: Sequence[str]) -> bool:
        return (
            len(names) == len(cls.index_names)
            and tuple(sorted(names)) in cls.all_valid_index_name_sorted_combos
        )

    @classmethod
    def find_valid_index_names(cls, dtypes: "pd.Series[Any]") -> Optional[List[str]]:
        valid_names = []
        for names, check_type in cls.index_names.items():
            for name in names:
                if name in dtypes.index and check_type(dtypes[name]):
                    valid_names.append(name)
                    break
            else:
                return None
        return valid_names


class NeedsResultColumns(ABC):
    result_column_names: Mapping[str, Callable[[Any], bool]] = MappingProxyType(
        {
            "score": is_numeric_dtype,
            "label": is_string_dtype,
            "explanation": is_string_dtype,
        }
    )

    @classmethod
    def is_valid_result_columns(cls, dtypes: "pd.Series[Any]") -> bool:
        names = cls.result_column_names.keys()
        intersection = dtypes.index.intersection(names)  # type: ignore
        if not len(intersection):
            return False
        for name in intersection:
            check_type = cls.result_column_names[name]
            if not check_type(dtypes[name]):
                return False
        return True


@dataclass(frozen=True)
class Evaluations(NeedsNamedIndex, NeedsResultColumns, ABC):
    eval_name: str  # The name for the evaluation, e.g. 'toxicity'
    dataframe: pd.DataFrame = field(repr=False)
    id: UUID = field(init=False, default_factory=uuid4)

    def __len__(self) -> int:
        return len(self.dataframe)

    def __repr__(self) -> str:
        return (
            f"{self.__class__.__name__}(eval_name={self.eval_name!r}, "
            f"dataframe=<rows: {len(self.dataframe)!r}>)"
        )

    def __dir__(self) -> List[str]:
        return ["get_dataframe"]

    def get_dataframe(self, prefix_columns_with_name: bool = True) -> pd.DataFrame:
        """
        Returns a copy of the dataframe with the evaluation annotations

        Parameters
        __________
        prefix_columns_with_name: bool
            if True, the columns will be prefixed with the eval_name, e.g. 'eval.toxicity.value'
        """
        if prefix_columns_with_name:
            prefix = f"{EVAL_NAME_COLUMN_PREFIX}{self.eval_name}."
            return self.dataframe.add_prefix(prefix)
        return self.dataframe.copy(deep=False)

    def __bool__(self) -> bool:
        return not self.dataframe.empty

    def __post_init__(self) -> None:
        dataframe = (
            pd.DataFrame() if self.dataframe.empty else self._clean_dataframe(self.dataframe)
        )
        object.__setattr__(self, "dataframe", dataframe)

    def _clean_dataframe(self, dataframe: pd.DataFrame) -> pd.DataFrame:
        # Ensure column names are strings.
        column_names = dataframe.columns.astype(str)
        dataframe = dataframe.set_axis(column_names, axis=1)

        # If the dataframe contains the index columns, set the index to those columns
        if not self.is_valid_index_names(dataframe.index.names) and (
            index_names := self.find_valid_index_names(dataframe.dtypes)
        ):
            dataframe = dataframe.set_index(index_names)

        # Validate that the dataframe is indexed correctly.
        if not self.is_valid_index_names(dataframe.index.names):
            raise ValueError(
                f"The dataframe index must be {self.preferred_names()} but was "
                f"'{dataframe.index.name or dataframe.index.names}'"
            )

        # Validate that the dataframe contains result columns of appropriate types.
        if not self.is_valid_result_columns(dataframe.dtypes):
            raise ValueError(
                f"The dataframe must contain one of these columns with appropriate "
                f"value types: {self.result_column_names.keys()} "
            )

        # Un-alias to the preferred names.
        preferred_names = [self.unalias(name) for name in dataframe.index.names]
        dataframe = dataframe.rename_axis(preferred_names, axis=0)

        # Drop the unnecessary columns.
        result_column_names = dataframe.columns.intersection(self.result_column_names.keys())  # type: ignore
        return dataframe.loc[:, result_column_names]  # type: ignore

    def __init_subclass__(
        cls,
        index_names: Mapping[Tuple[str, ...], Callable[[Any], bool]],
        **kwargs: Any,
    ) -> None:
        super().__init_subclass__(**kwargs)
        cls.index_names = index_names
        cls.all_valid_index_name_sorted_combos = set(
            tuple(sorted(prod)) for prod in product(*cls.index_names.keys())
        )

    def to_parquet(self, directory: Optional[Union[str, Path]] = None) -> Path:
        """Persists the evaluations to a parquet file.

        Args:
            directory (Optional[Union[str, Path]], optional): An optional path
            to a directory where the parquet file will be saved. If not
            provided, the parquet file will be saved to a default location.

        Returns:
            Path: The path to the parquet file, including a randomly generated
            filename.
        """
        directory = Path(directory) if directory else TRACE_DATASET_DIR
        path = directory / f"evaluations-{self.id}.parquet"
        table = Table.from_pandas(self.dataframe)
        table = table.replace_schema_metadata(
            {
                **(table.schema.metadata or {}),
                # explicitly encode keys and values, which are automatically encoded regardless
                b"arize": json.dumps(
                    {
                        "eval_id": str(self.id),
                        "eval_name": self.eval_name,
                        "eval_type": self.__class__.__name__,
                    }
                ).encode("utf-8"),
            }
        )
        parquet.write_table(table, path)
        return path

    @classmethod
    def from_parquet(cls, path: Union[str, Path]) -> "Evaluations":
        """Loads the evaluations from a parquet file.

        Args:
            path (Union[str, Path]): Path to a persisted evaluations parquet
            file.

        Returns:
            Evaluations: The loaded evaluations. The type of the returned
            evaluations will be the same as the type of the evaluations that
            were originally persisted.
        """
        schema = parquet.read_schema(path)
        eval_id, eval_name, evaluations_cls = _parse_schema_metadata(schema.metadata)
        table = parquet.read_table(path)
        dataframe = table.to_pandas()
        evaluations = evaluations_cls(eval_name=eval_name, dataframe=dataframe)
        object.__setattr__(evaluations, "id", eval_id)
        return evaluations


@dataclass(frozen=True)
class SpanEvaluations(
    Evaluations,
    index_names=MappingProxyType({("context.span_id", "span_id"): is_string_dtype}),
):
    """
    SpanEvaluations is a set of evaluation annotations for a set of spans.
    SpanEvaluations encompasses the evaluation annotations for a single evaluation task
    such as toxicity or hallucinations.

    SpanEvaluations can be appended to TraceDatasets so that the spans and
    evaluations can be joined and analyzed together.

    Parameters
    __________
    eval_name: str
        the name of the evaluation, e.g. 'toxicity'
    dataframe: pandas.DataFrame
        the pandas dataframe containing the evaluation annotations Each row
        represents the evaluations on a span.

    Example
    _______

    DataFrame of evaluations for toxicity may look like:

    | span_id | score              | label              | explanation        |
    |---------|--------------------|--------------------|--------------------|
    | span_1  | 1                  | toxic              | bad language       |
    | span_2  | 0                  | non-toxic          | violence           |
    | span_3  | 1                  | toxic              | discrimination     |
    """


@dataclass(frozen=True)
class DocumentEvaluations(
    Evaluations,
    index_names=MappingProxyType(
        {
            ("context.span_id", "span_id"): is_string_dtype,
            ("document_position", "position"): is_integer_dtype,
        }
    ),
):
    """
    DocumentEvaluations is a set of evaluation annotations for a set of documents.
    DocumentEvaluations encompasses the evaluation annotations for a single evaluation task
    such as relevance.

    Parameters
    __________
    eval_name: str
        the name of the evaluation, e.g. 'relevance'
    dataframe: pandas.DataFrame
        the pandas dataframe containing the evaluation annotations. Each row
        represents the evaluations on a document.

    Example
    _______

    DataFrame of document evaluations for relevance may look like:

    | span_id | position | score | label      | explanation  |
    |---------|----------|-------|------------|--------------|
    | span_1  | 0        | 1     | relevant   | it's apropos |
    | span_1  | 1        | 1     | relevant   | it's germane |
    | span_2  | 0        | 0     | irrelevant | it's rubbish |
    """

    def _clean_dataframe(self, dataframe: pd.DataFrame) -> pd.DataFrame:
        dataframe = super()._clean_dataframe(dataframe)
        if dataframe.index.names != self.preferred_names():
            return dataframe.swaplevel()
        return dataframe


@dataclass(frozen=True)
class TraceEvaluations(
    Evaluations,
    index_names=MappingProxyType({("context.trace_id", "trace_id"): is_string_dtype}),
):
    ...


def _parse_schema_metadata(metadata: Dict[bytes, Any]) -> Tuple[UUID, str, Type[Evaluations]]:
    """Validates and parses the schema metadata. Raises an exception if the
    metadata is invalid.

    Args:
        metadata (Dict[bytes, Any]): A dictionary of schema metadata from a
        parquet file.

    Returns:
        Tuple[str, ModuleType]: The evaluation name and the evaluations class.
    """
    if not (arize_metadata_json := metadata.get(b"arize")):
        raise InvalidParquetMetadataError('Schema metadata is missing "arize" key')
    try:
        arize_metadata = json.loads(arize_metadata_json)
    except json.JSONDecodeError as err:
        raise InvalidParquetMetadataError(
            'Encountered invalid JSON string under "arize" key'
        ) from err
    evaluations_classes = {subclass.__name__: subclass for subclass in Evaluations.__subclasses__()}
    if not (
        isinstance(arize_metadata, dict)
        and (eval_id := _to_uuid(arize_metadata.get("eval_id")))
        and isinstance(eval_name := arize_metadata.get("eval_name"), str)
        and (eval_type := arize_metadata.get("eval_type"))
        and (evaluations_cls := evaluations_classes.get(eval_type))
    ):
        raise InvalidParquetMetadataError(f"Invalid Arize metadata: {arize_metadata}")
    return eval_id, eval_name, evaluations_cls


def _to_uuid(value: Any) -> Optional[UUID]:
    """
    Converts an input to a UUID if possible, otherwise returns None.

    Args:
        value (Any): The value to convert to a UUID.

    Returns:
        Optional[UUID]: A UUID if the value could be converted, otherwise None.
    """
    try:
        return UUID(value)
    except Exception:
        return None
