from __future__ import annotations

import json
from enum import Enum
from typing import TYPE_CHECKING, Any

from mlflow.data import Dataset
from mlflow.data.evaluation_dataset_source import EvaluationDatasetSource
from mlflow.data.pyfunc_dataset_mixin import PyFuncConvertibleDatasetMixin
from mlflow.entities._mlflow_object import _MlflowObject
from mlflow.entities.dataset_record import DatasetRecord
from mlflow.entities.dataset_record_source import DatasetRecordSourceType
from mlflow.exceptions import MlflowException
from mlflow.protos.datasets_pb2 import Dataset as ProtoDataset
from mlflow.telemetry.events import DatasetToDataFrameEvent, MergeRecordsEvent
from mlflow.telemetry.track import record_usage_event
from mlflow.tracing.constant import TraceMetadataKey
from mlflow.tracking.context import registry as context_registry
from mlflow.utils.mlflow_tags import MLFLOW_USER

if TYPE_CHECKING:
    import pandas as pd

    from mlflow.entities.trace import Trace


SESSION_IDENTIFIER_FIELDS = frozenset({"goal"})
SESSION_INPUT_FIELDS = frozenset({"persona", "goal", "context", "simulation_guidelines"})
SESSION_ALLOWED_COLUMNS = SESSION_INPUT_FIELDS | {"expectations", "tags", "source"}


class DatasetGranularity(Enum):
    TRACE = "trace"
    SESSION = "session"
    UNKNOWN = "unknown"


class EvaluationDataset(_MlflowObject, Dataset, PyFuncConvertibleDatasetMixin):
    """
    Evaluation dataset for storing inputs and expectations for GenAI evaluation.

    This class supports lazy loading of records - when retrieved via get_evaluation_dataset(),
    only metadata is loaded. Records are fetched when to_df() or merge_records() is called.
    """

    def __init__(
        self,
        dataset_id: str,
        name: str,
        digest: str,
        created_time: int,
        last_update_time: int,
        tags: dict[str, Any] | None = None,
        schema: str | None = None,
        profile: str | None = None,
        created_by: str | None = None,
        last_updated_by: str | None = None,
    ):
        """Initialize the EvaluationDataset."""
        self.dataset_id = dataset_id
        self.created_time = created_time
        self.last_update_time = last_update_time
        self.tags = tags
        self._schema = schema
        self._profile = profile
        self.created_by = created_by
        self.last_updated_by = last_updated_by
        self._experiment_ids = None
        self._records = None

        source = EvaluationDatasetSource(dataset_id=self.dataset_id)
        Dataset.__init__(self, source=source, name=name, digest=digest)

    def _compute_digest(self) -> str:
        """
        Compute digest for the dataset. This is called by Dataset.__init__ if no digest is provided.
        Since we always have a digest from the dataclass initialization, this should not be called.
        """
        return self.digest

    @property
    def source(self) -> EvaluationDatasetSource:
        """Override source property to return the correct type."""
        return self._source

    @property
    def schema(self) -> str | None:
        """
        Dataset schema information.
        """
        return self._schema

    @property
    def profile(self) -> str | None:
        """
        Dataset profile information.
        """
        return self._profile

    @property
    def experiment_ids(self) -> list[str]:
        """
        Get associated experiment IDs, loading them if necessary.

        This property implements lazy loading - experiment IDs are only fetched from the backend
        when accessed for the first time.
        """
        if self._experiment_ids is None:
            self._load_experiment_ids()
        return self._experiment_ids or []

    @experiment_ids.setter
    def experiment_ids(self, value: list[str]):
        """Set experiment IDs directly."""
        self._experiment_ids = value or []

    def _load_experiment_ids(self):
        """Load experiment IDs from the backend."""
        from mlflow.tracking._tracking_service.utils import _get_store

        tracking_store = _get_store()
        self._experiment_ids = tracking_store.get_dataset_experiment_ids(self.dataset_id)

    @property
    def records(self) -> list[DatasetRecord]:
        """
        Get dataset records, loading them if necessary.

        This property implements lazy loading - records are only fetched from the backend
        when accessed for the first time.
        """
        if self._records is None:
            from mlflow.tracking._tracking_service.utils import _get_store

            tracking_store = _get_store()
            # For lazy loading, we want all records (no pagination)
            self._records, _ = tracking_store._load_dataset_records(
                self.dataset_id, max_results=None
            )
        return self._records or []

    def has_records(self) -> bool:
        """Check if dataset records are loaded without triggering a load."""
        return self._records is not None

    def _process_trace_records(self, traces: list["Trace"]) -> list[dict[str, Any]]:
        """Convert a list of Trace objects to dataset record dictionaries.

        Args:
            traces: List of Trace objects to convert

        Returns:
            List of dictionaries with 'inputs', 'expectations', and 'source' fields
        """
        from mlflow.entities.trace import Trace

        record_dicts = []
        for i, trace in enumerate(traces):
            if not isinstance(trace, Trace):
                raise MlflowException.invalid_parameter_value(
                    f"Mixed types in trace list. Expected all elements to be Trace objects, "
                    f"but element at index {i} is {type(trace).__name__}"
                )

            root_span = trace.data._get_root_span()
            inputs = root_span.inputs if root_span and root_span.inputs is not None else {}
            outputs = root_span.outputs if root_span and root_span.outputs is not None else None

            expectations = {}
            expectation_assessments = trace.search_assessments(type="expectation")
            for expectation in expectation_assessments:
                expectations[expectation.name] = expectation.value

            # Preserve session metadata from the original trace
            source_data = {"trace_id": trace.info.trace_id}
            if session_id := trace.info.trace_metadata.get(TraceMetadataKey.TRACE_SESSION):
                source_data["session_id"] = session_id

            record_dict = {
                "inputs": inputs,
                "outputs": outputs,
                "expectations": expectations,
                "source": {
                    "source_type": DatasetRecordSourceType.TRACE.value,
                    "source_data": source_data,
                },
            }
            record_dicts.append(record_dict)

        return record_dicts

    def _process_dataframe_records(self, df: "pd.DataFrame") -> list[dict[str, Any]]:
        """Process a DataFrame into dataset record dictionaries.

        Args:
            df: DataFrame to process. Can be either:
                - DataFrame from search_traces with 'trace' column containing Trace objects/JSON
                - Standard DataFrame with 'inputs', 'expectations' columns

        Returns:
            List of dictionaries with 'inputs', 'expectations', and optionally 'source' fields
        """
        if "trace" in df.columns:
            from mlflow.entities.trace import Trace

            traces = [
                Trace.from_json(trace_item) if isinstance(trace_item, str) else trace_item
                for trace_item in df["trace"]
            ]

            return self._process_trace_records(traces)
        else:
            return df.to_dict("records")

    @record_usage_event(MergeRecordsEvent)
    def merge_records(
        self, records: list[dict[str, Any]] | "pd.DataFrame" | list["Trace"]
    ) -> "EvaluationDataset":
        """
        Merge new records with existing ones.

        Args:
            records: Records to merge. Can be:
                - List of dictionaries with 'inputs' and optionally 'expectations' and 'tags'
                - Session format with 'persona', 'goal', 'context' nested inside 'inputs'
                - DataFrame from mlflow.search_traces() - automatically parsed and converted
                - DataFrame with 'inputs' column and optionally 'expectations' and 'tags' columns
                - List of Trace objects

        Returns:
            Self for method chaining

        Example:
            .. code-block:: python

                # Direct usage with search_traces DataFrame output
                traces_df = mlflow.search_traces()  # Returns DataFrame by default
                dataset.merge_records(traces_df)  # No extraction needed

                # Or with standard DataFrame
                df = pd.DataFrame([{"inputs": {"q": "What?"}, "expectations": {"a": "Answer"}}])
                dataset.merge_records(df)

                # Session format in inputs
                test_cases = [
                    {
                        "inputs": {
                            "persona": "Student",
                            "goal": "Find articles",
                            "context": {"student_id": "U1"},
                        }
                    },
                ]
                dataset.merge_records(test_cases)
        """
        import pandas as pd

        from mlflow.entities.trace import Trace
        from mlflow.tracking._tracking_service.utils import _get_store, get_tracking_uri

        if isinstance(records, pd.DataFrame):
            record_dicts = self._process_dataframe_records(records)
        elif isinstance(records, list) and records and isinstance(records[0], Trace):
            record_dicts = self._process_trace_records(records)
        else:
            record_dicts = records

        self._validate_record_dicts(record_dicts)

        self._infer_source_types(record_dicts)

        tracking_store = _get_store()

        try:
            existing_dataset = tracking_store.get_dataset(self.dataset_id)
            self._schema = existing_dataset.schema
        except Exception as e:
            raise MlflowException.invalid_parameter_value(
                f"Cannot add records to dataset {self.dataset_id}: Dataset not found. "
                f"Please verify the dataset exists and check your tracking URI is set correctly "
                f"(currently set to: {get_tracking_uri()})."
            ) from e

        self._validate_schema(record_dicts)

        context_tags = context_registry.resolve_tags()
        if user_tag := context_tags.get(MLFLOW_USER):
            for record in record_dicts:
                if "tags" not in record:
                    record["tags"] = {}
                if MLFLOW_USER not in record["tags"]:
                    record["tags"][MLFLOW_USER] = user_tag

        tracking_store.upsert_dataset_records(dataset_id=self.dataset_id, records=record_dicts)
        self._records = None

        return self

    def _validate_record_dicts(self, record_dicts: list[dict[str, Any]]) -> None:
        """Validate that record dictionaries have the required structure.

        Args:
            record_dicts: List of record dictionaries to validate

        Raises:
            MlflowException: If records don't have the required structure
        """
        for record in record_dicts:
            if not isinstance(record, dict):
                raise MlflowException.invalid_parameter_value("Each record must be a dictionary")
            if "inputs" not in record:
                raise MlflowException.invalid_parameter_value(
                    "Each record must have an 'inputs' field"
                )

    def _infer_source_types(self, record_dicts: list[dict[str, Any]]) -> None:
        """Infer source types for records without explicit source information.

        Simple inference rules:
        - Records with expectations -> HUMAN (manual test cases/ground truth)
        - Records with inputs but no expectations -> CODE (programmatically generated)

        Inference can be overridden by providing explicit source information.

        Note that trace inputs (from List[Trace] or pd.DataFrame of Trace data) will
        always be inferred as a trace source type when processing trace records.

        Args:
            record_dicts: List of record dictionaries to process (modified in place)
        """
        for record in record_dicts:
            if "source" in record:
                continue

            if "expectations" in record and record["expectations"]:
                record["source"] = {
                    "source_type": DatasetRecordSourceType.HUMAN.value,
                    "source_data": {},
                }
            elif "inputs" in record and "expectations" not in record:
                record["source"] = {
                    "source_type": DatasetRecordSourceType.CODE.value,
                    "source_data": {},
                }

    def _validate_schema(self, record_dicts: list[dict[str, Any]]) -> None:
        """
        Validate schema consistency of new records and compatibility with existing dataset.

        Args:
            record_dicts: List of normalized record dictionaries

        Raises:
            MlflowException: If records have invalid schema, inconsistent schemas within batch,
                or are incompatible with existing dataset schema
        """
        granularity_counts: dict[DatasetGranularity, int] = {}
        has_empty_inputs = False

        for record in record_dicts:
            input_keys = set(record.get("inputs", {}).keys())
            if not input_keys:
                has_empty_inputs = True
                continue

            record_type = self._classify_input_fields(input_keys)

            if record_type == DatasetGranularity.UNKNOWN:
                session_fields = input_keys & SESSION_IDENTIFIER_FIELDS
                other_fields = input_keys - SESSION_INPUT_FIELDS
                raise MlflowException.invalid_parameter_value(
                    f"Invalid input schema: cannot mix session fields {list(session_fields)} "
                    f"with other fields {list(other_fields)}. "
                    f"Consider placing {list(other_fields)} fields inside 'context'."
                )

            granularity_counts[record_type] = granularity_counts.get(record_type, 0) + 1

        if len(granularity_counts) > 1:
            counts_str = ", ".join(
                f"{count} records with {granularity.value} granularity"
                for granularity, count in granularity_counts.items()
            )
            raise MlflowException.invalid_parameter_value(
                f"All records must use the same granularity. Found {counts_str}."
            )

        batch_granularity = next(iter(granularity_counts), DatasetGranularity.UNKNOWN)
        existing_granularity = self._get_existing_granularity()

        if has_empty_inputs and DatasetGranularity.SESSION in {
            batch_granularity,
            existing_granularity,
        }:
            raise MlflowException.invalid_parameter_value(
                "Empty inputs are not allowed for session records. The 'goal' field is required."
            )

        if DatasetGranularity.UNKNOWN in {batch_granularity, existing_granularity}:
            return

        if batch_granularity != existing_granularity:
            raise MlflowException.invalid_parameter_value(
                f"New records use {batch_granularity.value} granularity, but existing "
                f"dataset uses {existing_granularity.value}. Cannot mix granularities."
            )

    def _get_existing_granularity(self) -> DatasetGranularity:
        """
        Get granularity from the dataset's stored schema.

        Returns:
            DatasetGranularity based on existing records, or UNKNOWN if empty/unparseable
        """
        if self._schema is None:
            if self.has_records():
                return self._classify_input_fields(set(self.records[0].inputs.keys()))
            return DatasetGranularity.UNKNOWN
        try:
            schema = json.loads(self._schema)
            input_keys = set(schema.get("inputs", {}).keys())
            return self._classify_input_fields(input_keys)
        except (json.JSONDecodeError, TypeError):
            return DatasetGranularity.UNKNOWN

    @staticmethod
    def _classify_input_fields(input_keys: set[str]) -> DatasetGranularity:
        """
        Classify a set of input field names into a granularity type:
        - SESSION: Has 'goal' field, and only session fields (persona, goal, context)
        - TRACE: No 'goal' field present
        - UNKNOWN: Empty or has 'goal' mixed with non-session fields

        Args:
            input_keys: Set of field names from a record's inputs

        Returns:
            DatasetGranularity classification for the input fields
        """
        if not input_keys:
            return DatasetGranularity.UNKNOWN

        has_session_identifier = bool(input_keys & SESSION_IDENTIFIER_FIELDS)

        if not has_session_identifier:
            return DatasetGranularity.TRACE

        if input_keys <= SESSION_INPUT_FIELDS:
            return DatasetGranularity.SESSION

        return DatasetGranularity.UNKNOWN

    def delete_records(self, record_ids: list[str]) -> int:
        """
        Delete specific records from the dataset.

        Args:
            record_ids: List of record IDs to delete.

        Returns:
            The number of records deleted.

        Example:
            .. code-block:: python

                # Get record IDs to delete
                df = dataset.to_df()
                record_ids_to_delete = df["dataset_record_id"].tolist()[:2]

                # Delete the records
                deleted_count = dataset.delete_records(record_ids_to_delete)
                print(f"Deleted {deleted_count} records")
        """
        from mlflow.tracking._tracking_service.utils import _get_store

        tracking_store = _get_store()
        deleted_count = tracking_store.delete_dataset_records(
            dataset_id=self.dataset_id,
            dataset_record_ids=record_ids,
        )
        self._records = None  # Clear cached records
        return deleted_count

    @record_usage_event(DatasetToDataFrameEvent)
    def to_df(self) -> "pd.DataFrame":
        """
        Convert dataset records to a pandas DataFrame.

        This method triggers lazy loading of records if they haven't been loaded yet.

        Returns:
            DataFrame with columns for inputs, outputs, expectations, tags, and metadata
        """
        import pandas as pd

        records = self.records

        if not records:
            return pd.DataFrame(
                columns=[
                    "inputs",
                    "outputs",
                    "expectations",
                    "tags",
                    "source_type",
                    "source_id",
                    "source",
                    "created_time",
                    "dataset_record_id",
                ]
            )

        data = [
            {
                "inputs": record.inputs,
                "outputs": record.outputs,
                "expectations": record.expectations,
                "tags": record.tags,
                "source_type": record.source_type,
                "source_id": record.source_id,
                "source": record.source,
                "created_time": record.created_time,
                "dataset_record_id": record.dataset_record_id,
            }
            for record in records
        ]

        return pd.DataFrame(data)

    def to_proto(self) -> ProtoDataset:
        """Convert to protobuf representation."""
        proto = ProtoDataset()

        proto.dataset_id = self.dataset_id
        proto.name = self.name
        if self.tags is not None:
            proto.tags = json.dumps(self.tags)
        if self.schema is not None:
            proto.schema = self.schema
        if self.profile is not None:
            proto.profile = self.profile
        proto.digest = self.digest
        proto.created_time = self.created_time
        proto.last_update_time = self.last_update_time
        if self.created_by is not None:
            proto.created_by = self.created_by
        if self.last_updated_by is not None:
            proto.last_updated_by = self.last_updated_by
        if self._experiment_ids is not None:
            proto.experiment_ids.extend(self._experiment_ids)

        return proto

    @classmethod
    def from_proto(cls, proto: ProtoDataset) -> "EvaluationDataset":
        """Create instance from protobuf representation."""
        tags = None
        if proto.HasField("tags"):
            tags = json.loads(proto.tags)

        dataset = cls(
            dataset_id=proto.dataset_id,
            name=proto.name,
            digest=proto.digest,
            created_time=proto.created_time,
            last_update_time=proto.last_update_time,
            tags=tags,
            schema=proto.schema if proto.HasField("schema") else None,
            profile=proto.profile if proto.HasField("profile") else None,
            created_by=proto.created_by if proto.HasField("created_by") else None,
            last_updated_by=proto.last_updated_by if proto.HasField("last_updated_by") else None,
        )
        if proto.experiment_ids:
            dataset._experiment_ids = list(proto.experiment_ids)
        return dataset

    def to_dict(self) -> dict[str, Any]:
        """Convert to dictionary representation."""
        result = super().to_dict()

        result.update({
            "dataset_id": self.dataset_id,
            "tags": self.tags,
            "schema": self.schema,
            "profile": self.profile,
            "created_time": self.created_time,
            "last_update_time": self.last_update_time,
            "created_by": self.created_by,
            "last_updated_by": self.last_updated_by,
            "experiment_ids": self.experiment_ids,
        })

        result["records"] = [record.to_dict() for record in self.records]

        return result

    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> "EvaluationDataset":
        """Create instance from dictionary representation."""
        if "dataset_id" not in data:
            raise ValueError("dataset_id is required")
        if "name" not in data:
            raise ValueError("name is required")
        if "digest" not in data:
            raise ValueError("digest is required")
        if "created_time" not in data:
            raise ValueError("created_time is required")
        if "last_update_time" not in data:
            raise ValueError("last_update_time is required")

        dataset = cls(
            dataset_id=data["dataset_id"],
            name=data["name"],
            digest=data["digest"],
            created_time=data["created_time"],
            last_update_time=data["last_update_time"],
            tags=data.get("tags"),
            schema=data.get("schema"),
            profile=data.get("profile"),
            created_by=data.get("created_by"),
            last_updated_by=data.get("last_updated_by"),
        )
        if "experiment_ids" in data:
            dataset._experiment_ids = data["experiment_ids"]

        if "records" in data:
            dataset._records = [
                DatasetRecord.from_dict(record_data) for record_data in data["records"]
            ]

        return dataset