from __future__ import annotations

import inspect
from typing import Any
from typing import get_args
from typing import Protocol
from typing import runtime_checkable

import pandas as pd
import pyarrow as pa
import pydantic.dataclasses as pydc
from typing_extensions import Literal
from typing_extensions import TypedDict

from tempora.utils.dataclass import pydc_config

SchemaMismatchPolicy = Literal['raise', 'coerce', 'skip']
LabelUnknownPolicy = Literal['raise', 'skip', 'use_default']
LabelReduction = Literal['first', 'last', 'mode', 'unique']
EventMissingPolicy = Literal['negative', 'skip', 'raise']


class TransformMetadata(TypedDict):
    columns: list[str]
    source: Literal['context', 'target']
    time_column: str | None
    entity_keys: list[str] | None


@runtime_checkable
class TransformFunction(Protocol):
    def __call__(
        self, data: pd.DataFrame, metadata: TransformMetadata,
    ) -> pd.DataFrame:
        ...


@pydc.dataclass(frozen=True, config=pydc_config)
class TransformSpec:
    """Serializable transform specification.

    Args:
        transform: User-defined transform callable.
        output_schema: Optional output schema for transformed data.
        on_schema_mismatch: Policy to apply if transformed schema differs from the
            expected schema. With `'coerce'`, values are cast to expected types where
            possible; if a required column is missing, behavior falls through to
            `'skip'` for that segment.
    """
    transform: TransformFunction
    output_schema: pa.Schema | None = None
    on_schema_mismatch: SchemaMismatchPolicy = 'raise'

    def __post_init__(self) -> None:
        if not callable(self.transform):
            raise TypeError(
                f'"transform" must be callable, got {type(self.transform).__name__}',
            )

        try:
            sig = inspect.signature(self.transform)
        except (TypeError, ValueError) as e:
            raise TypeError('"transform" must expose an inspectable signature') from e
        params = list(sig.parameters.values())
        has_var_positional = any(
            p.kind is inspect.Parameter.VAR_POSITIONAL for p in params
        )
        if len(params) < 2 and not has_var_positional:
            raise TypeError(
                '"transform" must accept at least two positional args: '
                '(data, metadata)',
            )

        if self.on_schema_mismatch not in get_args(SchemaMismatchPolicy):
            raise ValueError(
                f'Invalid value for "on_schema_mismatch": "{self.on_schema_mismatch}"',
            )


@pydc.dataclass(frozen=True, config=pydc_config)
class LabelEncoder:
    """Encode target-window values into numeric class labels."""
    reduction: LabelReduction = 'last'
    class_map: dict[Any, int] | None = None
    unknown_policy: LabelUnknownPolicy = 'skip'
    default_label: int | None = None

    def __post_init__(self) -> None:
        if self.reduction not in get_args(LabelReduction):
            raise ValueError(f'Invalid value for "reduction": "{self.reduction}"')

        if self.unknown_policy not in get_args(LabelUnknownPolicy):
            raise ValueError(
                f'Invalid value for "unknown_policy": "{self.unknown_policy}"',
            )

        if self.unknown_policy == 'use_default' and self.default_label is None:
            raise ValueError(
                '"default_label" is required when "unknown_policy" is "use_default"',
            )

        if self.class_map is not None:
            for raw_label, encoded_label in self.class_map.items():
                if not isinstance(encoded_label, int):
                    raise TypeError(
                        f'Class map value for label "{raw_label}" must be an int, '
                        f'got {type(encoded_label).__name__}',
                    )

        if self.default_label is not None and not isinstance(self.default_label, int):
            raise TypeError(
                f'"default_label" must be an int, got '
                f'{type(self.default_label).__name__}',
            )

        if self.unknown_policy == 'use_default' and not isinstance(
            self.default_label, int,
        ):
            raise TypeError(
                '"default_label" must be a valid int when "unknown_policy" is '
                '"use_default"',
            )


@pydc.dataclass(frozen=True, config=pydc_config)
class EventEncoder:
    """Encode event date/datetime targets as binary in-window labels."""
    positive_label: int = 1
    negative_label: int = 0
    include_start: bool = True
    include_end: bool = False
    missing_policy: EventMissingPolicy = 'negative'

    def __post_init__(self) -> None:
        if self.positive_label == self.negative_label:
            raise ValueError('"positive_label" and "negative_label" must differ')

        if self.missing_policy not in get_args(EventMissingPolicy):
            raise ValueError(
                f'Invalid value for "missing_policy": "{self.missing_policy}"',
            )
