"""
Label API - Declarative Label Definitions.

This module provides a fluent builder pattern for defining labels,
analogous to the FeatureBuilder for features.
"""

from dataclasses import dataclass
from typing import Any, Dict, List, Optional, Tuple, Union

from ..enums import ScalarType
from ..resources.label import LabelVariant
from ..resources.mappings import EntityMapping, EntityMappings, ResourceColumnMapping


@dataclass
class BuiltLabel:
    """Wrapper for a built label variant."""

    label: LabelVariant
    base_name: str
    base_variant: str

    def get_label(self) -> LabelVariant:
        """Get the underlying LabelVariant."""
        return self.label

    def name_variant(self) -> Tuple[str, str]:
        """Get the label's name and variant as a tuple."""
        return (self.base_name, self.base_variant)

    def get_all_labels(self) -> List[LabelVariant]:
        """Get all labels (for consistency with BuiltFeatures)."""
        return [self.label]


class LabelBuilder:
    """
    Builder class for declaratively defining labels using a fluent API.

    The LabelBuilder provides a chainable interface for defining labels that
    will be registered with Featureform. Labels can have single or multiple
    entities.

    Example - Single Entity Label:
        ```python
        @ff.entity
        class Customer:
            is_fraud = (
                ff.Label()
                .from_dataset(transactions)
                .value("is_fraud")
                .timestamp("ts")
                .entity("customer", "customer_id")
                .description("Binary fraud classification label")
            )
        ```

    Example - Multi-Entity Label:
        ```python
        @ff.entity
        class Transaction:
            is_funded = (
                ff.Label()
                .from_dataset(transactions)
                .value("is_funded")
                .timestamp("ts")
                .entities({
                    "customer": "customer_id",
                    "loan": "loan_id"
                })
            )
        ```

    Note:
        When used inside an @ff.entity decorated class, the label name and
        variant are automatically derived from the attribute name and the
        current run, respectively.
    """

    def __init__(self):
        """Initialize an empty LabelBuilder."""
        # Label identification
        self.base_name: Optional[str] = None
        self.initial_variant: Optional[str] = None

        # Source configuration
        self.source: Optional[Any] = None

        # Column configuration
        self.value_column: Optional[str] = None
        self.timestamp_column: Optional[str] = None

        # Entity configuration (supports single or multiple entities)
        self.entity_mappings: List[EntityMapping] = []

        # Metadata
        self._owner: str = ""
        self._description: str = ""
        self._tags: List[str] = []
        self._properties: Dict[str, str] = {}
        self._value_type: Optional[Union[ScalarType, str]] = None

        # Built label (set after build() is called)
        self._built_label: Optional["BuiltLabel"] = None

        # Tracking for validation
        self._entity_called: bool = False
        self._entities_called: bool = False

    def name(self, name: str) -> "LabelBuilder":
        """Set the label name."""
        self.base_name = name
        return self

    def variant(self, variant: str) -> "LabelBuilder":
        """Set the label variant."""
        self.initial_variant = variant
        return self

    def from_dataset(self, source: Any) -> "LabelBuilder":
        """Set the source dataset for the label."""
        self.source = source
        return self

    def value(self, column: str) -> "LabelBuilder":
        """Set the value column for the label."""
        self.value_column = column
        return self

    def timestamp(self, column: str) -> "LabelBuilder":
        """Set the timestamp column for the label."""
        self.timestamp_column = column
        return self

    def entity(self, name: Any, column: str) -> "LabelBuilder":
        """
        Set a single entity mapping for the label.

        Args:
            name: Entity reference (string, EntityRegistrar, or Entity object)
            column: Column name in the source that contains entity values
        """
        if self._entities_called:
            raise ValueError("Cannot use both .entity() and .entities()")
        self._entity_called = True

        resolved_name = self._resolve_entity_name(name)
        self.entity_mappings.append(
            EntityMapping(name=resolved_name, entity_column=column)
        )
        return self

    def entities(self, mappings: Dict[Any, str]) -> "LabelBuilder":
        """
        Set multiple entity mappings for the label.

        Args:
            mappings: Dictionary mapping entity references to column names.
                      Keys can be strings, EntityRegistrar, or Entity objects.
        """
        if self._entity_called:
            raise ValueError("Cannot use both .entity() and .entities()")
        self._entities_called = True

        for entity_ref, column_name in mappings.items():
            resolved_name = self._resolve_entity_name(entity_ref)
            self.entity_mappings.append(
                EntityMapping(name=resolved_name, entity_column=column_name)
            )
        return self

    def _resolve_entity_name(self, entity: Any) -> str:
        """
        Resolve an entity reference to its string name.

        Handles:
        - str: returned as-is
        - EntityRegistrar: calls .name() method
        - Entity: accesses .name property

        Args:
            entity: Entity reference (string, EntityRegistrar, or Entity)

        Returns:
            The entity name as a string

        Raises:
            ValueError: If entity type is not recognized
        """
        if isinstance(entity, str):
            return entity
        elif hasattr(entity, "name"):
            # Check if it's a method (EntityRegistrar) or property (Entity)
            name_attr = getattr(entity, "name")
            if callable(name_attr):
                return name_attr()
            else:
                return name_attr
        else:
            raise ValueError(
                f"Cannot resolve entity name from type {type(entity).__name__}. "
                f"Expected str, EntityRegistrar, or Entity."
            )

    def description(self, desc: str) -> "LabelBuilder":
        """Set the label description."""
        self._description = desc
        return self

    def owner(self, owner: Any) -> "LabelBuilder":
        """
        Set the label owner.

        Args:
            owner: Owner name (string) or UserRegistrar-like object with .name() method
        """
        if isinstance(owner, str):
            self._owner = owner
        elif hasattr(owner, "name"):
            # UserRegistrar has a .name() method
            name_attr = getattr(owner, "name")
            self._owner = name_attr() if callable(name_attr) else name_attr
        else:
            raise ValueError(f"Cannot resolve owner from type {type(owner)}")
        return self

    def tags(self, tags: List[str]) -> "LabelBuilder":
        """Set the label tags."""
        self._tags = tags
        return self

    def properties(self, props: Dict[str, str]) -> "LabelBuilder":
        """Set the label properties."""
        self._properties = props
        return self

    def type(self, value_type: Union[ScalarType, str]) -> "LabelBuilder":
        """Set the label value type."""
        self._value_type = value_type
        return self

    def _validate(self) -> None:
        """Validate the builder configuration before building."""
        errors = []

        if self.base_name is None:
            errors.append("Label name must be set via .name() or entity decorator")

        if self.source is None:
            errors.append("Label source must be set via .from_dataset()")

        if self.value_column is None:
            errors.append("Label value column must be set via .value()")

        if not self.entity_mappings:
            errors.append(
                "At least one entity mapping must be set via .entity() or .entities()"
            )

        if errors:
            raise ValueError(
                "Invalid LabelBuilder configuration:\n"
                + "\n".join(f"  - {e}" for e in errors)
            )

    def build(self) -> BuiltLabel:
        """Build the LabelVariant from the current configuration."""
        self._validate()

        # Store source reference directly - don't resolve to tuple here.
        # The LabelVariant._get_and_set_equivalent_variant() method will
        # resolve it at registration time, which allows the source variant
        # to be updated if an equivalent source is found during apply().
        source_ref = self.source

        # Create location based on entity count
        if len(self.entity_mappings) == 1:
            # Single entity: use ResourceColumnMapping
            location = ResourceColumnMapping(
                entity=self.entity_mappings[0].entity_column,
                value=self.value_column,
                timestamp=self.timestamp_column or "",
            )
            entity_name = self.entity_mappings[0].name
        else:
            # Multiple entities: use EntityMappings
            location = EntityMappings(
                mappings=self.entity_mappings,
                value_column=self.value_column,
                timestamp_column=self.timestamp_column or "",
            )
            # For multi-entity, use the first entity as the primary
            entity_name = self.entity_mappings[0].name

        # Create LabelVariant with source reference (not tuple)
        # This allows the source variant to be updated at registration time
        label = LabelVariant(
            name=self.base_name,
            variant=self.initial_variant,
            source=source_ref,
            entity=entity_name,
            owner=self._owner,
            description=self._description,
            location=location,
            value_type=self._value_type or ScalarType.STRING,
            tags=self._tags,
            properties=self._properties,
        )

        built = BuiltLabel(
            label=label,
            base_name=self.base_name,
            base_variant=self.initial_variant,
        )
        self._built_label = built
        return built

    def _get_source_name_variant(self) -> Tuple[str, str]:
        """Extract name and variant from the source."""
        if hasattr(self.source, "name_variant"):
            return self.source.name_variant()
        elif hasattr(self.source, "name") and hasattr(self.source, "variant"):
            return (self.source.name, self.source.variant)
        else:
            raise ValueError(
                f"Cannot extract name/variant from source: {type(self.source)}"
            )


def Label(*args, **kwargs):
    """
    Factory function for creating labels.

    This is the main entry point for defining labels. It supports both
    the legacy API and the new builder pattern.

    **Legacy API** (pass transformation_args as first positional argument):
        ```python
        @ff.entity
        class User:
            is_fraud = ff.Label(
                transformation[["CustomerID", "is_fraud"]],
                type=ff.Bool,
                variant="v1",
            )
        ```

    **Builder Pattern** (no arguments, use method chaining):
        ```python
        @ff.entity
        class Customer:
            is_fraud = (
                ff.Label()
                .name("is_fraud")
                .variant("v1")
                .from_dataset(transactions)
                .value("is_fraud")
                .entity("customer", "customer_id")
            )
        ```

    Both patterns can be used together in the same entity class, allowing
    for seamless migration from the legacy API to the builder pattern.

    Returns:
        LabelBuilder (no args) or LabelColumnResource (with args)
    """
    if args or kwargs:
        # Legacy API - delegate to LabelColumnResource
        from .column_resources import LabelColumnResource

        return LabelColumnResource(*args, **kwargs)
    else:
        # Builder Pattern - return new LabelBuilder
        return LabelBuilder()
