from __future__ import annotations

from typing import Any

import pandas as pd
from pandas.api.types import (
    is_bool_dtype,
    is_categorical_dtype,
    is_datetime64_any_dtype,
    is_numeric_dtype,
    is_object_dtype,
    is_string_dtype,
)


ID_NAME_TOKENS = (
    "id",
    "uuid",
    "guid",
    "key",
    "code",
    "token",
    "hash",
    "identifier",
)

TEXT_NAME_HINTS = (
    "text",
    "description",
    "comment",
    "message",
    "body",
    "title",
    "name",
)


def _safe_nunique(series: pd.Series, dropna: bool = True) -> int:
    try:
        return int(series.nunique(dropna=dropna))
    except Exception:
        return 0


def _nonnull_ratio(series: pd.Series) -> float:
    if len(series) == 0:
        return 0.0
    return float(series.notna().mean())


def _avg_string_length(series: pd.Series) -> float:
    nonnull = series.dropna()
    if nonnull.empty:
        return 0.0

    try:
        as_str = nonnull.astype(str)
        return float(as_str.str.len().mean())
    except Exception:
        return 0.0


def _looks_text_like(series: pd.Series) -> bool:
    """
    Heuristic:
    object/string column with fairly long average value length
    and not close to unique IDs.
    """
    if not (is_object_dtype(series) or is_string_dtype(series) or is_categorical_dtype(series)):
        return False

    nonnull_ratio = _nonnull_ratio(series)
    if nonnull_ratio == 0.0:
        return False

    nunique = _safe_nunique(series)
    unique_ratio = nunique / max(len(series), 1)
    avg_len = _avg_string_length(series)

    return avg_len >= 20 and unique_ratio < 0.95


def _looks_identifier_like(column_name: str, series: pd.Series) -> bool:
    name = str(column_name).lower()

    name_hit = any(token in name for token in ID_NAME_TOKENS)
    nunique = _safe_nunique(series)
    unique_ratio = nunique / max(len(series), 1)
    avg_len = _avg_string_length(series)

    if name_hit and unique_ratio >= 0.5:
        return True

    # High-cardinality string-like columns are often identifiers.
    if (is_object_dtype(series) or is_string_dtype(series) or is_categorical_dtype(series)) and unique_ratio >= 0.95:
        return True

    # Long opaque strings often behave like hashes / IDs.
    if (is_object_dtype(series) or is_string_dtype(series)) and unique_ratio >= 0.80 and avg_len >= 12:
        return True

    return False


def infer_column_schema(column_name: str, series: pd.Series) -> dict[str, Any]:
    """
    Infer a compact schema summary for a single column.
    """

    dtype_str = str(series.dtype)
    nunique = _safe_nunique(series)
    unique_ratio = nunique / max(len(series), 1)
    missing_ratio = float(series.isna().mean()) if len(series) else 0.0

    inferred_type = "unknown"
    if is_datetime64_any_dtype(series):
        inferred_type = "datetime"
    elif is_bool_dtype(series):
        inferred_type = "boolean"
    elif is_numeric_dtype(series):
        inferred_type = "numeric"
    elif is_categorical_dtype(series):
        inferred_type = "categorical"
    elif is_string_dtype(series) or is_object_dtype(series):
        inferred_type = "text" if _looks_text_like(series) else "categorical"

    return {
        "name": column_name,
        "dtype": dtype_str,
        "inferred_type": inferred_type,
        "nunique": nunique,
        "unique_ratio": unique_ratio,
        "missing_ratio": missing_ratio,
        "nonnull_ratio": _nonnull_ratio(series),
        "avg_string_length": _avg_string_length(series),
        "identifier_like": _looks_identifier_like(column_name, series),
        "text_like": _looks_text_like(series),
    }


def infer_schema(df: pd.DataFrame) -> dict[str, Any]:
    """
    Infer dataset-level schema information.

    Returns a dict with:
    - per-column details
    - grouped column lists by inferred type
    """

    columns: dict[str, dict[str, Any]] = {}
    numeric_cols: list[str] = []
    categorical_cols: list[str] = []
    datetime_cols: list[str] = []
    boolean_cols: list[str] = []
    text_cols: list[str] = []
    identifier_cols: list[str] = []

    for col in df.columns:
        info = infer_column_schema(col, df[col])
        columns[col] = info

        inferred_type = info["inferred_type"]
        if inferred_type == "numeric":
            numeric_cols.append(col)
        elif inferred_type == "categorical":
            categorical_cols.append(col)
        elif inferred_type == "datetime":
            datetime_cols.append(col)
        elif inferred_type == "boolean":
            boolean_cols.append(col)
        elif inferred_type == "text":
            text_cols.append(col)

        if info["identifier_like"]:
            identifier_cols.append(col)

    return {
        "n_rows": int(len(df)),
        "n_columns": int(len(df.columns)),
        "columns": columns,
        "numeric_columns": numeric_cols,
        "categorical_columns": categorical_cols,
        "datetime_columns": datetime_cols,
        "boolean_columns": boolean_cols,
        "text_columns": text_cols,
        "identifier_columns": identifier_cols,
    }