from __future__ import annotations

import json
import textwrap
from collections import Counter
from collections.abc import (
    Awaitable,
    Callable,
    Iterable,
    Iterator,
    Mapping,
)
from copy import copy, deepcopy
from dataclasses import dataclass, field, fields
from datetime import datetime
from enum import Enum
from functools import cached_property
from importlib.metadata import version
from random import getrandbits
from typing import Any, Optional, TypeVar, Union, cast, overload

import pandas as pd
from typing_extensions import TypeAlias
from wrapt import ObjectProxy

from phoenix.datetime_utils import local_now
from phoenix.experiments.utils import get_experiment_url


class AnnotatorKind(Enum):
    CODE = "CODE"
    LLM = "LLM"


JSONSerializable: TypeAlias = Optional[Union[dict[str, Any], list[Any], str, int, float, bool]]
ExperimentId: TypeAlias = str
DatasetId: TypeAlias = str
DatasetVersionId: TypeAlias = str
ExampleId: TypeAlias = str
RepetitionNumber: TypeAlias = int
ExperimentRunId: TypeAlias = str
TraceId: TypeAlias = str

TaskOutput: TypeAlias = JSONSerializable

ExampleOutput: TypeAlias = Mapping[str, JSONSerializable]
ExampleMetadata: TypeAlias = Mapping[str, JSONSerializable]
ExampleInput: TypeAlias = Mapping[str, JSONSerializable]

Score: TypeAlias = Optional[Union[bool, int, float]]
Label: TypeAlias = Optional[str]
Explanation: TypeAlias = Optional[str]

EvaluatorName: TypeAlias = str
EvaluatorKind: TypeAlias = str
EvaluatorOutput: TypeAlias = Union[
    "EvaluationResult", bool, int, float, str, tuple[Score, Label, Explanation]
]

DRY_RUN: ExperimentId = "DRY_RUN"


def _dry_run_id() -> str:
    suffix = getrandbits(24).to_bytes(3, "big").hex()
    return f"{DRY_RUN}_{suffix}"


@dataclass(frozen=True)
class Example:
    id: ExampleId
    updated_at: datetime
    input: Mapping[str, JSONSerializable] = field(default_factory=dict)
    output: Mapping[str, JSONSerializable] = field(default_factory=dict)
    metadata: Mapping[str, JSONSerializable] = field(default_factory=dict)

    def __post_init__(self) -> None:
        object.__setattr__(self, "input", _make_read_only(self.input))
        object.__setattr__(self, "output", _make_read_only(self.output))
        object.__setattr__(self, "metadata", _make_read_only(self.metadata))

    @classmethod
    def from_dict(cls, obj: Mapping[str, Any]) -> Example:
        return cls(
            input=obj["input"],
            output=obj["output"],
            metadata=obj.get("metadata") or {},
            id=obj["id"],
            updated_at=obj["updated_at"],
        )

    def __repr__(self) -> str:
        spaces = " " * 4
        name = self.__class__.__name__
        identifiers = [f'{spaces}id="{self.id}",']
        contents = [
            spaces
            + f"{_blue(key)}="
            + json.dumps(
                _shorten(value),
                ensure_ascii=False,
                sort_keys=True,
                indent=len(spaces),
            )
            .replace("\n", f"\n{spaces}")
            .replace(' "..."\n', " ...\n")
            + ","
            for key in ("input", "output", "metadata")
            if (value := getattr(self, key, None))
        ]
        return "\n".join([f"{name}(", *identifiers, *contents, ")"])


@dataclass(frozen=True)
class Dataset:
    id: DatasetId
    version_id: DatasetVersionId
    examples: Mapping[ExampleId, Example] = field(repr=False, default_factory=dict)

    def __post_init__(self) -> None:
        object.__setattr__(self, "examples", _ReadOnly(self.examples))

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

    def __iter__(self) -> Iterator[Example]:
        return iter(self.examples.values())

    @cached_property
    def _keys(self) -> tuple[str, ...]:
        return tuple(self.examples.keys())

    @overload
    def __getitem__(self, key: int) -> Example: ...
    @overload
    def __getitem__(self, key: slice) -> list[Example]: ...
    def __getitem__(self, key: Union[int, slice]) -> Union[Example, list[Example]]:
        if isinstance(key, int):
            return self.examples[self._keys[key]]
        return [self.examples[k] for k in self._keys[key]]

    def as_dataframe(self, drop_empty_columns: bool = True) -> pd.DataFrame:
        df = pd.DataFrame.from_records(
            [
                {
                    "example_id": example.id,
                    "input": deepcopy(example.input),
                    "output": deepcopy(example.output),
                    "metadata": deepcopy(example.metadata),
                }
                for example in self.examples.values()
            ]
        ).set_index("example_id")
        if drop_empty_columns:
            return df.reindex([k for k, v in df.items() if v.astype(bool).any()], axis=1)
        return df

    @classmethod
    def from_dict(cls, obj: Mapping[str, Any]) -> Dataset:
        examples = tuple(map(Example.from_dict, obj.get("examples") or ()))
        return cls(
            id=obj["dataset_id"],
            version_id=obj["version_id"],
            examples={ex.id: ex for ex in examples},
        )


@dataclass(frozen=True)
class TestCase:
    example: Example
    repetition_number: RepetitionNumber


@dataclass(frozen=True)
class Experiment:
    id: ExperimentId
    dataset_id: DatasetId
    dataset_version_id: DatasetVersionId
    repetitions: int
    project_name: str = field(repr=False)

    @classmethod
    def from_dict(cls, obj: Mapping[str, Any]) -> Experiment:
        return cls(
            id=obj["id"],
            dataset_id=obj["dataset_id"],
            dataset_version_id=obj["dataset_version_id"],
            repetitions=obj.get("repetitions") or 1,
            project_name=obj.get("project_name") or "",
        )


@dataclass(frozen=True)
class ExperimentRun:
    start_time: datetime
    end_time: datetime
    experiment_id: ExperimentId
    dataset_example_id: ExampleId
    repetition_number: RepetitionNumber
    output: JSONSerializable
    error: Optional[str] = None
    id: ExperimentRunId = field(default_factory=_dry_run_id)
    trace_id: Optional[TraceId] = None

    @classmethod
    def from_dict(cls, obj: Mapping[str, Any]) -> ExperimentRun:
        return cls(
            start_time=obj["start_time"],
            end_time=obj["end_time"],
            experiment_id=obj["experiment_id"],
            dataset_example_id=obj["dataset_example_id"],
            repetition_number=obj.get("repetition_number") or 1,
            output=_make_read_only(obj.get("output")),
            error=obj.get("error"),
            id=obj["id"],
            trace_id=obj.get("trace_id"),
        )

    def __post_init__(self) -> None:
        if self.output is None and self.error is None:
            raise ValueError("Must specify exactly one of experiment_run_output or error")


@dataclass(frozen=True)
class EvaluationResult:
    score: Optional[float] = None
    label: Optional[str] = None
    explanation: Optional[str] = None
    metadata: Mapping[str, JSONSerializable] = field(default_factory=dict)

    @classmethod
    def from_dict(cls, obj: Optional[Mapping[str, Any]]) -> Optional[EvaluationResult]:
        if not obj:
            return None
        return cls(
            score=obj.get("score"),
            label=obj.get("label"),
            explanation=obj.get("explanation"),
            metadata=obj.get("metadata") or {},
        )

    def __post_init__(self) -> None:
        if self.score is None and not self.label:
            raise ValueError("Must specify score or label, or both")
        if self.score is None and not self.label:
            object.__setattr__(self, "score", 0)
        for k in ("label", "explanation"):
            if (v := getattr(self, k, None)) is not None:
                object.__setattr__(self, k, str(v) or None)


@dataclass(frozen=True)
class ExperimentEvaluationRun:
    experiment_run_id: ExperimentRunId
    start_time: datetime
    end_time: datetime
    name: str
    annotator_kind: str
    error: Optional[str] = None
    result: Optional[EvaluationResult] = None
    id: str = field(default_factory=_dry_run_id)
    trace_id: Optional[TraceId] = None

    @classmethod
    def from_dict(cls, obj: Mapping[str, Any]) -> ExperimentEvaluationRun:
        return cls(
            experiment_run_id=obj["experiment_run_id"],
            start_time=obj["start_time"],
            end_time=obj["end_time"],
            name=obj["name"],
            annotator_kind=obj["annotator_kind"],
            error=obj.get("error"),
            result=EvaluationResult.from_dict(obj.get("result")),
            id=obj["id"],
            trace_id=obj.get("trace_id"),
        )

    def __post_init__(self) -> None:
        if self.result is None and self.error is None:
            raise ValueError("Must specify either result or error")


ExperimentTask: TypeAlias = Union[
    Callable[[Example], TaskOutput],
    Callable[[Example], Awaitable[TaskOutput]],
]


@dataclass(frozen=True)
class ExperimentParameters:
    n_examples: int
    n_repetitions: int = 1

    @property
    def count(self) -> int:
        return self.n_examples * self.n_repetitions


@dataclass(frozen=True)
class EvaluationParameters:
    eval_names: frozenset[str]
    exp_params: ExperimentParameters


@dataclass(frozen=True)
class _HasStats:
    _title: str = field(repr=False, default="")
    _timestamp: datetime = field(repr=False, default_factory=local_now)
    stats: pd.DataFrame = field(repr=False, default_factory=pd.DataFrame)

    @property
    def title(self) -> str:
        return f"{self._title} ({self._timestamp:%x %I:%M %p %z})"

    def __str__(self) -> str:
        try:
            assert int(version("pandas").split(".")[0]) >= 1
            # `tabulate` is used by pandas >= 1.0 in DataFrame.to_markdown()
            import tabulate  # noqa: F401
        except (AssertionError, ImportError):
            text = self.stats.__str__()
        else:
            text = self.stats.to_markdown(index=False)
        return f"{self.title}\n{'-'*len(self.title)}\n" + text


@dataclass(frozen=True)
class EvaluationSummary(_HasStats):
    """
    Summary statistics of experiment evaluations.

    Users should not instantiate this directly.
    """

    _title: str = "Experiment Summary"

    @classmethod
    def from_eval_runs(
        cls,
        params: EvaluationParameters,
        *eval_runs: Optional[ExperimentEvaluationRun],
    ) -> EvaluationSummary:
        df = pd.DataFrame.from_records(
            [
                {
                    "evaluator": run.name,
                    "error": run.error,
                    "score": run.result.score if run.result else None,
                    "label": run.result.label if run.result else None,
                }
                for run in eval_runs
                if run is not None
            ]
        )
        if df.empty:
            df = pd.DataFrame.from_records(
                [
                    {"evaluator": name, "error": None, "score": None, "label": None}
                    for name in params.eval_names
                ]
            )
        has_error = bool(df.loc[:, "error"].astype(bool).sum())
        has_score = bool(df.loc[:, "score"].dropna().count())
        has_label = bool(df.loc[:, "label"].astype(bool).sum())
        agg = {
            **(
                dict(n_errors=("error", "count"), top_error=("error", _top_string))
                if has_error
                else {}
            ),
            **(dict(n_scores=("score", "count"), avg_score=("score", "mean")) if has_score else {}),
            **(
                dict(
                    n_labels=("label", "count"),
                    top_2_labels=(
                        "label",
                        lambda s: (dict(Counter(s.dropna()).most_common(2)) or None),
                    ),
                )
                if has_label
                else {}
            ),
        }
        stats = (
            df.groupby("evaluator").agg(**agg)  # type: ignore[call-overload]
            if agg
            else pd.DataFrame()
        )
        sorted_eval_names = sorted(params.eval_names)
        eval_names = pd.DataFrame(
            {
                "evaluator": sorted_eval_names,
                "n": [params.exp_params.count] * len(sorted_eval_names),
            }
        ).set_index("evaluator")
        stats = pd.concat([eval_names, stats], axis=1).reset_index()
        summary: EvaluationSummary = object.__new__(cls)
        summary.__init__(stats=stats)  # type: ignore[misc]
        return summary

    @classmethod
    def __new__(cls, *args: Any, **kwargs: Any) -> Any:
        # Direct instantiation by users is discouraged.
        raise NotImplementedError

    @classmethod
    def __init_subclass__(cls, **kwargs: Any) -> None:
        # Direct sub-classing by users is discouraged.
        raise NotImplementedError


@dataclass(frozen=True)
class TaskSummary(_HasStats):
    """
    Summary statistics of experiment task executions.

    **Users should not instantiate this object directly.**
    """

    _title: str = "Tasks Summary"

    @classmethod
    def from_task_runs(
        cls,
        params: ExperimentParameters,
        task_runs: Iterable[Optional[ExperimentRun]],
    ) -> "TaskSummary":
        df = pd.DataFrame.from_records(
            [
                {
                    "example_id": run.dataset_example_id,
                    "error": run.error,
                }
                for run in task_runs
                if run is not None
            ]
        )
        n_runs = len(df)
        n_errors = 0 if df.empty else df.loc[:, "error"].astype(bool).sum()
        record = {
            "n_examples": params.count,
            "n_runs": n_runs,
            "n_errors": n_errors,
            **(dict(top_error=_top_string(df.loc[:, "error"])) if n_errors else {}),
        }
        stats = pd.DataFrame.from_records([record])
        summary: TaskSummary = object.__new__(cls)
        summary.__init__(stats=stats)  # type: ignore[misc]
        return summary

    @classmethod
    def __new__(cls, *args: Any, **kwargs: Any) -> Any:
        # Direct instantiation by users is discouraged.
        raise NotImplementedError

    @classmethod
    def __init_subclass__(cls, **kwargs: Any) -> None:
        # Direct sub-classing by users is discouraged.
        raise NotImplementedError


def _top_string(s: "pd.Series[Any]", length: int = 100) -> Optional[str]:
    if (cnt := s.dropna().str.slice(0, length).value_counts()).empty:
        return None
    return cast(str, cnt.sort_values(ascending=False).index[0])


@dataclass(frozen=True)
class RanExperiment(Experiment):
    """
    An experiment that has been run.

    **Users should not instantiate this object directly.**
    """

    params: ExperimentParameters = field(repr=False)
    dataset: Dataset = field(repr=False)
    runs: Mapping[ExperimentRunId, ExperimentRun] = field(repr=False)
    task_summary: TaskSummary = field(repr=False)
    eval_runs: tuple[ExperimentEvaluationRun, ...] = field(repr=False, default=())
    eval_summaries: tuple[EvaluationSummary, ...] = field(repr=False, default=())

    @property
    def url(self) -> str:
        return get_experiment_url(dataset_id=self.dataset.id, experiment_id=self.id)

    @property
    def info(self) -> str:
        return f"🔗 View this experiment: {self.url}"

    def __post_init__(self) -> None:
        runs = {
            id_: (
                _ExperimentRunWithExample(run, example)
                if (example := self.dataset.examples.get(run.dataset_example_id))
                else run
            )
            for id_, run in self.runs.items()
        }
        object.__setattr__(self, "runs", runs)

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

    def __iter__(self) -> Iterator[ExperimentRun]:
        return iter(self.runs.values())

    @cached_property
    def _keys(self) -> tuple[str, ...]:
        return tuple(self.runs.keys())

    @overload
    def __getitem__(self, key: int) -> ExperimentRun: ...
    @overload
    def __getitem__(self, key: slice) -> list[ExperimentRun]: ...
    def __getitem__(self, key: Union[int, slice]) -> Union[ExperimentRun, list[ExperimentRun]]:
        if isinstance(key, int):
            return self.runs[self._keys[key]]
        return [self.runs[k] for k in self._keys[key]]

    def get_evaluations(
        self,
        drop_empty_columns: bool = True,
    ) -> pd.DataFrame:
        df = pd.DataFrame.from_records(
            [
                {
                    "run_id": run.experiment_run_id,
                    "name": run.name,
                    "error": run.error,
                    "score": run.result.score if run.result else None,
                    "label": run.result.label if run.result else None,
                    "explanation": run.result.explanation if run.result else None,
                }
                for run in self.eval_runs
            ]
        ).set_index("run_id")
        if drop_empty_columns:
            df = df.reindex([k for k, v in df.items() if v.astype(bool).any()], axis=1)
        return df.join(self.as_dataframe())

    def as_dataframe(self, drop_empty_columns: bool = True) -> pd.DataFrame:
        df = pd.DataFrame.from_records(
            [
                {
                    "run_id": run.id,
                    "error": run.error,
                    "output": deepcopy(run.output),
                    "input": deepcopy((ex := self.dataset.examples[run.dataset_example_id]).input),
                    "expected": deepcopy(ex.output),
                    "metadata": deepcopy(ex.metadata),
                    "example_id": run.dataset_example_id,
                }
                for run in self.runs.values()
            ]
        ).set_index("run_id")
        if drop_empty_columns:
            return df.reindex([k for k, v in df.items() if v.astype(bool).any()], axis=1)
        return df

    def add(
        self,
        eval_summary: EvaluationSummary,
        *eval_runs: Optional[ExperimentEvaluationRun],
    ) -> "RanExperiment":
        return _replace(
            self,
            eval_runs=(*self.eval_runs, *filter(bool, eval_runs)),
            eval_summaries=(*self.eval_summaries, eval_summary),
        )

    def __str__(self) -> str:
        summaries = (*reversed(self.eval_summaries), self.task_summary)
        return (
            "\n"
            + ("" if self.id.startswith(DRY_RUN) else f"{self.info}\n\n")
            + "\n\n".join(map(str, summaries))
        )

    @classmethod
    def __new__(cls, *args: Any, **kwargs: Any) -> Any:
        # Direct instantiation by users is discouraged.
        raise NotImplementedError

    @classmethod
    def __init_subclass__(cls, **kwargs: Any) -> None:
        # Direct sub-classing by users is discouraged.
        raise NotImplementedError


def _asdict(dc: Any) -> dict[str, Any]:
    # non-recursive version of `dataclasses.asdict()`
    return {field.name: getattr(dc, field.name) for field in fields(dc)}


T = TypeVar("T")


def _replace(obj: T, **kwargs: Any) -> T:
    new_obj = object.__new__(obj.__class__)
    new_obj.__init__(**{**_asdict(obj), **kwargs})  # type: ignore[misc]
    return new_obj


def _shorten(obj: Any, width: int = 50) -> Any:
    if isinstance(obj, str):
        return textwrap.shorten(obj, width=width, placeholder="...")
    if isinstance(obj, dict):
        return {k: _shorten(v) for k, v in obj.items()}
    if isinstance(obj, list):
        if len(obj) > 2:
            return [_shorten(v) for v in obj[:2]] + ["..."]
        return [_shorten(v) for v in obj]
    return obj


def _make_read_only(obj: Any) -> Any:
    if isinstance(obj, dict):
        return _ReadOnly({k: _make_read_only(v) for k, v in obj.items()})
    if isinstance(obj, str):
        return obj
    if isinstance(obj, list):
        return _ReadOnly(list(map(_make_read_only, obj)))
    return obj


class _ReadOnly(ObjectProxy):  # type: ignore[misc]
    def __setitem__(self, *args: Any, **kwargs: Any) -> Any:
        raise NotImplementedError

    def __delitem__(self, *args: Any, **kwargs: Any) -> Any:
        raise NotImplementedError

    def __iadd__(self, *args: Any, **kwargs: Any) -> Any:
        raise NotImplementedError

    def pop(self, *args: Any, **kwargs: Any) -> Any:
        raise NotImplementedError

    def append(self, *args: Any, **kwargs: Any) -> Any:
        raise NotImplementedError

    def __copy__(self, *args: Any, **kwargs: Any) -> Any:
        return copy(self.__wrapped__)

    def __deepcopy__(self, *args: Any, **kwargs: Any) -> Any:
        return deepcopy(self.__wrapped__)

    def __repr__(self) -> str:
        return repr(self.__wrapped__)

    def __str__(self) -> str:
        return str(self.__wrapped__)


class _ExperimentRunWithExample(ObjectProxy):  # type: ignore[misc]
    def __init__(self, wrapped: ExperimentRun, example: Example) -> None:
        super().__init__(wrapped)
        self._self_example = example

    @property
    def expected(self) -> ExampleOutput:
        return deepcopy(self._self_example.output)

    @property
    def reference(self) -> ExampleOutput:
        return deepcopy(self._self_example.output)

    @property
    def input(self) -> ExampleInput:
        return deepcopy(self._self_example.input)

    @property
    def metadata(self) -> ExampleMetadata:
        return deepcopy(self._self_example.metadata)

    def __repr__(self) -> str:
        spaces = " " * 4
        name = self.__class__.__name__
        identifiers = [
            f'{spaces}id="{self.id}",',
            f'{spaces}example_id="{self.dataset_example_id}",',
        ]
        outputs = [
            *([f'{spaces}error="{self.error}",'] if self.error else []),
            *(
                [
                    f"{spaces}{_blue('output')}="
                    + json.dumps(
                        _shorten(self.output),
                        ensure_ascii=False,
                        sort_keys=True,
                        indent=len(spaces),
                    )
                    .replace("\n", f"\n{spaces}")
                    .replace(' "..."\n', " ...\n")
                ]
                if not self.error
                else []
            ),
        ]
        dicts = [
            spaces
            + f"{_blue(alias)}={{"
            + (f" # {comment}" if comment else "")
            + json.dumps(
                _shorten(value),
                ensure_ascii=False,
                sort_keys=True,
                indent=len(spaces),
            )[1:]
            .replace("\n", f"\n{spaces}")
            .replace(' "..."\n', " ...\n")
            + ","
            for alias, value, comment in (
                ("expected", self.expected, f"alias for the example.{_blue('output')} dict"),
                ("reference", self.reference, f"alias for the example.{_blue('output')} dict"),
                ("input", self.input, f"alias for the example.{_blue('input')} dict"),
                ("metadata", self.metadata, f"alias for the example.{_blue('metadata')} dict"),
            )
            if value
        ]
        return "\n".join([f"{name}(", *identifiers, *outputs, *dicts, ")"])


def _blue(text: str) -> str:
    return f"\033[1m\033[94m{text}\033[0m"
