from __future__ import annotations

from typing import Any, Dict, List, Literal, Optional

from pydantic import (
    BaseModel,
    ConfigDict,
    Field,
    model_validator,
)

LOOKUP_SEP = "__"


class DataFrameParams(BaseModel):
    """
    Schema for parameters passed to the underlying DataFrame constructor/loader.
    """

    model_config = ConfigDict(extra="ignore", arbitrary_types_allowed=True)

    fieldnames: Optional[List[str]] = None
    index_col: Optional[str] = None
    db_index: Optional[str | List[str]] = None
    coerce_float: bool = False
    verbose: bool = True
    datetime_index: bool = False
    column_names: Optional[List[str]] = None
    chunk_size: int = 100_000

    def get(self, key: str, default: Any = None) -> Any:
        return getattr(self, key, default)


class DataFrameOptions(BaseModel):
    """
    Schema for options applied to the DataFrame after loading.
    """

    model_config = ConfigDict(extra="ignore")

    debug: bool = False
    duplicate_expr: Optional[str] = None
    duplicate_keep: Literal["first", "last", False] = "last"
    sort_field: Optional[str] = None
    group_by_expr: Optional[str] = None
    group_expr: Optional[str] = None

    def get(self, key: str, default: Any = None) -> Any:
        return getattr(self, key, default)


class ParamsConfig(BaseModel):
    """
    Immutable configuration object.
    Automatically distributes flat kwargs into nested configs (df_params, df_options, filters).
    """

    model_config = ConfigDict(frozen=True, extra="ignore")

    # Configuration
    field_map: Dict[str, str] = Field(default_factory=dict)
    legacy_filters: bool = False
    sticky_filters: Dict[str, Any] = Field(default_factory=dict)

    # Buckets
    df_params: DataFrameParams = Field(default_factory=DataFrameParams)
    df_options: DataFrameOptions = Field(default_factory=DataFrameOptions)
    filters: Dict[str, Any] = Field(default_factory=dict)

    @model_validator(mode="before")
    @classmethod
    def distribute_params(cls, data: Any) -> Any:
        """
        Root validator (Pre-Init):
        Iterates over input arguments and sorts them into the correct
        submodel (df_params, df_options) or treats them as filters.
        """
        if not isinstance(data, dict):
            return data

        # 1. Identify known fields
        known_root_fields = set(cls.model_fields.keys())
        known_df_params = set(DataFrameParams.model_fields.keys())
        known_df_options = set(DataFrameOptions.model_fields.keys())

        # 2. Extract Existing Sub-Models (Handling Models OR Dicts)
        # [FIX] If user passes DataFrameParams object, converting to dict preserves values.
        def _extract_dict(val: Any) -> Dict[str, Any]:
            if isinstance(val, BaseModel):
                return val.model_dump(exclude_unset=True)
            if isinstance(val, dict):
                return val.copy()
            return {}

        df_p_data = _extract_dict(data.get("df_params"))
        df_o_data = _extract_dict(data.get("df_options"))
        filters_data = _extract_dict(data.get("filters"))

        root_data = {}

        # 3. Sort flat keys into buckets (Merging logic)
        for k, v in data.items():
            # Skip keys we already extracted
            if k in ("df_params", "df_options", "filters"):
                continue

            if k in known_root_fields:
                root_data[k] = v
            elif k in known_df_params:
                df_p_data[k] = v
            elif k in known_df_options:
                df_o_data[k] = v
            else:
                filters_data[k] = v

        # 4. Handle Sticky Filters & Legacy Logic
        sticky = (
            root_data.get("sticky_filters", {}) or data.get("sticky_filters", {}) or {}
        )

        if "params" in data and isinstance(data["params"], dict):
            sticky.update(data["params"])

        final_filters = {**sticky, **filters_data}

        # 5. Legacy Filter Conversion
        use_legacy = root_data.get("legacy_filters", data.get("legacy_filters", False))
        field_map = root_data.get("field_map", data.get("field_map", {}))

        # [AUTO-DETECT] If field_map is provided, assume legacy_filters=True unless explicitly False
        if field_map and "legacy_filters" not in data:
            use_legacy = True

        if use_legacy and field_map:
            final_filters = cls._convert_legacy(final_filters, field_map)

        # 6. Return structured data to Pydantic
        return {
            **root_data,
            "legacy_filters": use_legacy,  # Sync back the detected value
            "df_params": df_p_data,
            "df_options": df_o_data,
            "filters": final_filters,
        }

    @staticmethod
    def _convert_legacy(
        filters: Dict[str, Any], field_map: Dict[str, str]
    ) -> Dict[str, Any]:
        """Static helper to rename filter keys based on field_map."""
        reverse_map = {v: k for k, v in field_map.items()}
        new_filters = {}

        for key, value in filters.items():
            parts = key.split(LOOKUP_SEP)
            new_parts = [reverse_map.get(p, p) for p in parts]
            new_key = LOOKUP_SEP.join(new_parts)
            new_filters[new_key] = value

        return new_filters

    def with_updates(self, **kwargs) -> "ParamsConfig":
        """
        Creates a modified copy by merging current state with new kwargs.
        Re-runs the distribution logic to correctly sort new params.
        """
        # Dump current state
        current_data = self.model_dump()

        # Flatten strictly needed buckets
        # Note: We do NOT flatten 'df_params' into the root dict here,
        # we keep it nested so the validator treats it as the "base" to update.
        flat_update = {
            # Start with existing complex objects as bases
            "df_params": current_data.get("df_params", {}),
            "df_options": current_data.get("df_options", {}),
            "filters": current_data.get("filters", {}),
            # Root config
            "field_map": self.field_map,
            "legacy_filters": self.legacy_filters,
            "sticky_filters": self.sticky_filters,
        }

        # Apply new flat arguments
        # The validator will merge these into the dictionaries above
        flat_update.update(kwargs)

        return ParamsConfig.model_validate(flat_update)
