"""Combine datasets using joins, binds, or set operations."""

import polars as pl

# All supported combine types
JOIN_TYPES = {"inner", "left", "right", "full", "semi", "anti"}
BIND_TYPES = {"bind_rows", "bind_cols"}
SET_TYPES = {"intersect", "union", "setdiff"}
ALL_TYPES = JOIN_TYPES | BIND_TYPES | SET_TYPES


def _align_join_key_dtypes(
    x_lf: pl.LazyFrame,
    y_lf: pl.LazyFrame,
    left_on: list[str],
    right_on: list[str],
) -> tuple:
    """
    Align join key dtypes between two LazyFrames.

    Smart dtype alignment preserving the "better" type:
    - Cat vs Str: Cast Str → Cat (preserves categorical benefits)
    - Int vs Float: Cast Int → Float (lossless upcast)
    - Other mismatches: Cast to String as fallback

    Returns:
        Tuple of (x_lf, y_lf) with aligned key dtypes
    """
    x_schema = x_lf.collect_schema()
    y_schema = y_lf.collect_schema()

    x_casts = []
    y_casts = []

    # Numeric types for Int vs Float detection
    int_types = {
        pl.Int8,
        pl.Int16,
        pl.Int32,
        pl.Int64,
        pl.UInt8,
        pl.UInt16,
        pl.UInt32,
        pl.UInt64,
    }
    float_types = {pl.Float32, pl.Float64}

    for left_col, right_col in zip(left_on, right_on):
        left_dtype = x_schema.get(left_col)
        right_dtype = y_schema.get(right_col)

        if left_dtype is None or right_dtype is None:
            continue  # Column doesn't exist, let Polars raise the error

        if left_dtype == right_dtype:
            continue  # Already aligned

        # Case 1: Enum vs Other - cast Other to Enum (preserves categorical benefits)
        if left_dtype == pl.Enum and right_dtype in (
            pl.String,
            pl.Utf8,
            pl.Categorical,
        ):
            y_casts.append(pl.col(right_col).cast(pl.Enum(left_dtype.categories)))
        elif right_dtype == pl.Enum and left_dtype in (
            pl.String,
            pl.Utf8,
            pl.Categorical,
        ):
            x_casts.append(pl.col(left_col).cast(pl.Enum(right_dtype.categories)))

        # Case 2: Cat vs Str - cast Str to Cat (preserves categorical benefits)
        if left_dtype == pl.Categorical and right_dtype in (pl.String, pl.Utf8):
            y_casts.append(pl.col(right_col).cast(pl.Categorical))
        elif right_dtype == pl.Categorical and left_dtype in (pl.String, pl.Utf8):
            x_casts.append(pl.col(left_col).cast(pl.Categorical))

        # Case 3: Int vs Float - cast Int to Float (lossless upcast)
        elif left_dtype in int_types and right_dtype in float_types:
            x_casts.append(pl.col(left_col).cast(right_dtype))
        elif right_dtype in int_types and left_dtype in float_types:
            y_casts.append(pl.col(right_col).cast(left_dtype))

        # Case 4: Fallback - cast both to String
        else:
            if left_dtype != pl.String:
                x_casts.append(pl.col(left_col).cast(pl.String))
            if right_dtype != pl.String:
                y_casts.append(pl.col(right_col).cast(pl.String))

    # Apply casts
    if x_casts:
        x_lf = x_lf.with_columns(x_casts)
    if y_casts:
        y_lf = y_lf.with_columns(y_casts)

    return x_lf, y_lf


def combine(
    x: pl.DataFrame | pl.LazyFrame,
    y: pl.DataFrame | pl.LazyFrame,
    on: str | list[str] | None = None,
    how: str = "inner",
    left_on: str | list[str] | None = None,
    right_on: str | list[str] | None = None,
    add: list[str] | None = None,
    suffix: str = "_right",
) -> pl.DataFrame:
    """
    Combine two datasets using joins, binds, or set operations.

    Parameters
    ----------
    x : pl.DataFrame | pl.LazyFrame
        Left/first dataset.
    y : pl.DataFrame | pl.LazyFrame
        Right/second dataset to combine with ``x``.
    on : str | list[str] | None
        Join key(s) when same column name on both sides.
    how : str
        Combine type. Default: ``"inner"``. Joins: inner, left, right, full,
        semi, anti. Binds: bind_rows, bind_cols. Sets: intersect, union, setdiff.
    left_on : str | list[str] | None
        Join key(s) from ``x`` when names differ from ``y``.
    right_on : str | list[str] | None
        Join key(s) from ``y`` when names differ from ``x``.
    add : list[str] | None
        Columns to select from ``y`` (optional, for joins). If None, uses all.
    suffix : str
        Suffix for overlapping column names. Default: ``"_right"``.

    Returns
    -------
    pl.DataFrame
        Combined dataset.

    Raises
    ------
    ValueError
        If ``how`` is unknown or join keys are missing for join operations.

    Notes
    -----
    Join key dtypes are automatically aligned (e.g., categorical cast to string).

    Examples
    --------
    >>> import polars as pl
    >>> import pyrsm as rsm
    >>> x = pl.DataFrame({"id": [1, 2], "a": ["x", "y"]})
    >>> y = pl.DataFrame({"id": [1, 2], "b": [10, 20]})
    >>> print(rsm.eda.combine(x, y, on="id"))
    shape: (2, 3)
    ┌─────┬─────┬─────┐
    │ id  ┆ a   ┆ b   │
    │ --- ┆ --- ┆ --- │
    │ i64 ┆ str ┆ i64 │
    ╞═════╪═════╪═════╡
    │ 1   ┆ x   ┆ 10  │
    │ 2   ┆ y   ┆ 20  │
    └─────┴─────┴─────┘
    >>> print(rsm.eda.combine(x, x, how="bind_rows"))
    shape: (4, 2)
    ┌─────┬─────┐
    │ id  ┆ a   │
    │ --- ┆ --- │
    │ i64 ┆ str │
    ╞═════╪═════╡
    │ 1   ┆ x   │
    │ 2   ┆ y   │
    │ 1   ┆ x   │
    │ 2   ┆ y   │
    └─────┴─────┘
    """
    # Validate how
    if how not in ALL_TYPES:
        raise ValueError(
            f"Unknown combine type: {how}\n"
            f"Supported: {', '.join(sorted(ALL_TYPES))}"
        )

    is_join = how in JOIN_TYPES

    # Validate join keys
    if is_join:
        if on is None and (left_on is None or right_on is None):
            raise ValueError(
                f"Join type '{how}' requires on= or (left_on= and right_on=).\n"
                "Example: rsm.eda.combine(x, y, on='customer_id', how='left')"
            )

    # Convert to LazyFrame for consistency
    x_lf = x.lazy() if isinstance(x, pl.DataFrame) else x
    y_lf = y.lazy() if isinstance(y, pl.DataFrame) else y

    # Normalize join keys to lists
    if on is not None:
        left_keys = [on] if isinstance(on, str) else list(on)
        right_keys = left_keys
    elif left_on is not None and right_on is not None:
        left_keys = [left_on] if isinstance(left_on, str) else list(left_on)
        right_keys = [right_on] if isinstance(right_on, str) else list(right_on)
    else:
        left_keys = []
        right_keys = []

    # Optionally limit columns from y
    if add and is_join:
        keep_cols = list(set(right_keys + add))
        y_lf = y_lf.select(keep_cols)

    # Execute based on type
    if is_join:
        # Auto-align join key dtypes
        x_lf, y_lf = _align_join_key_dtypes(x_lf, y_lf, left_keys, right_keys)

        if left_keys == right_keys:
            result = x_lf.join(y_lf, on=left_keys, how=how, suffix=suffix)
        else:
            result = x_lf.join(
                y_lf, left_on=left_keys, right_on=right_keys, how=how, suffix=suffix
            )

    elif how == "bind_rows":
        result = pl.concat([x_lf, y_lf], how="diagonal")

    elif how == "bind_cols":
        # Horizontal concat - collect for frame alignment
        x_df = x_lf.collect()
        y_df = y_lf.collect()
        result = pl.concat([x_df, y_df], how="horizontal").lazy()

    elif how == "intersect":
        # Rows present in both datasets
        x_df = x_lf.collect()
        y_df = y_lf.collect()
        result = x_df.join(y_df, on=x_df.columns, how="semi").lazy()

    elif how == "union":
        # All rows from both, deduplicated
        result = pl.concat([x_lf, y_lf]).unique()

    elif how == "setdiff":
        # Rows in x but not in y
        x_df = x_lf.collect()
        y_df = y_lf.collect()
        result = x_df.join(y_df, on=x_df.columns, how="anti").lazy()

    return result.collect()
