"""Core enums for SLayer."""

import datetime  # noqa: F401  (kept for downstream imports of TimeGranularity)
from enum import Enum
from typing import Any, Optional


class StrEnum(str, Enum):
    def __str__(self) -> str:
        return self.value


class DataType(StrEnum):
    """SLayer data types — values match sqlglot's ``exp.DataType.Type``
    byte-for-byte so SQL generation can ``CAST`` to the declared type without
    a translation map. (DEV-1361.)"""

    TEXT = "TEXT"
    INT = "INT"
    DOUBLE = "DOUBLE"
    BOOLEAN = "BOOLEAN"
    DATE = "DATE"
    TIMESTAMP = "TIMESTAMP"


# DEV-1361: lenient before-validator absorbs legacy lowercase type spellings
# from older agent input (MCP/REST/CLI), pseudo-types (count/sum/...) drop to
# None so the field falls through to its default. Used by both Column and
# ModelMeasure validators in slayer/core/models.py.
_LEGACY_DATATYPE_ALIASES: dict[str, Optional[str]] = {
    # Pre-rename canonical values.
    "string": "TEXT",
    "number": "DOUBLE",
    "integer": "INT",
    "time": "TIMESTAMP",
    "date": "DATE",
    "boolean": "BOOLEAN",
    # Aggregation pseudo-types — dropped in v5 because they were unused.
    "count": None,
    "count_distinct": None,
    "sum": None,
    "avg": None,
    "min": None,
    "max": None,
    "last": None,
}


def _coerce_legacy_datatype(v: Any) -> Any:
    """Map legacy lowercase ``DataType`` strings to current canonical values.

    Pseudo-types resolve to ``None`` so the calling validator can drop them
    and let the field default fire. Already-canonical values, enum instances,
    and unknown strings pass through untouched (Pydantic's enum coercion will
    raise on unknown).
    """
    if isinstance(v, str):
        mapped = _LEGACY_DATATYPE_ALIASES.get(v)
        if v in _LEGACY_DATATYPE_ALIASES:
            return mapped
    return v



class TimeGranularity(StrEnum):
    SECOND = "second"
    MINUTE = "minute"
    HOUR = "hour"
    DAY = "day"
    WEEK = "week"
    MONTH = "month"
    QUARTER = "quarter"
    YEAR = "year"

    def period_start(self, date: datetime.date) -> datetime.date:
        if self in (TimeGranularity.SECOND, TimeGranularity.MINUTE, TimeGranularity.HOUR):
            return date
        if self == TimeGranularity.DAY:
            return date
        elif self == TimeGranularity.WEEK:
            return date - datetime.timedelta(days=date.weekday())
        elif self == TimeGranularity.MONTH:
            return date.replace(day=1)
        elif self == TimeGranularity.QUARTER:
            quarter_month = ((date.month - 1) // 3) * 3 + 1
            return date.replace(month=quarter_month, day=1)
        elif self == TimeGranularity.YEAR:
            return date.replace(month=1, day=1)
        raise ValueError(f"Unexpected granularity: {self}")

    def period_end(self, date: datetime.date) -> datetime.date:
        if self in (TimeGranularity.SECOND, TimeGranularity.MINUTE, TimeGranularity.HOUR):
            return date
        if self == TimeGranularity.DAY:
            return date
        elif self == TimeGranularity.WEEK:
            return date + datetime.timedelta(days=6 - date.weekday())
        elif self == TimeGranularity.MONTH:
            if date.month == 12:
                return date.replace(year=date.year + 1, month=1, day=1) - datetime.timedelta(days=1)
            else:
                return date.replace(month=date.month + 1, day=1) - datetime.timedelta(days=1)
        elif self == TimeGranularity.QUARTER:
            quarter_end_month = ((date.month - 1) // 3) * 3 + 3
            if quarter_end_month == 12:
                return datetime.date(date.year, 12, 31)
            else:
                return datetime.date(date.year, quarter_end_month + 1, 1) - datetime.timedelta(days=1)
        elif self == TimeGranularity.YEAR:
            return date.replace(month=12, day=31)
        raise ValueError(f"Unexpected granularity: {self}")


class OrderDirection(StrEnum):
    ASC = "asc"
    DESC = "desc"


class JoinType(StrEnum):
    LEFT = "left"
    INNER = "inner"


# ---------------------------------------------------------------------------
# Aggregation constants
# ---------------------------------------------------------------------------

# Built-in aggregation names (always available without model-level definition).
BUILTIN_AGGREGATIONS: frozenset[str] = frozenset({
    "sum", "avg", "min", "max",
    "count", "count_distinct",
    "first", "last",
    "weighted_avg",
    "median", "percentile",
    "stddev_samp", "stddev_pop",
    "var_samp", "var_pop",
    "corr", "covar_samp", "covar_pop",
})

# Built-in aggregation SQL formulas (for aggregations that use a template).
# {value} = measure's SQL expression; {param_name} = parameter values.
# Note: percentile is dialect-dependent (no single template works on
# SQLite/ClickHouse/MySQL) and lives in generator._build_percentile instead.
BUILTIN_AGGREGATION_FORMULAS: dict[str, str] = {
    "weighted_avg": "SUM({value} * {weight}) / NULLIF(SUM({weight}), 0)",
}

# Built-in aggregations that require specific parameters.
# Percentile's required-param check lives in generator._build_percentile.
BUILTIN_AGGREGATION_REQUIRED_PARAMS: dict[str, list[str]] = {
    "weighted_avg": ["weight"],
    "corr": ["other"],
    "covar_samp": ["other"],
    "covar_pop": ["other"],
}

# Aggregations that only make sense on numeric-valued measures. Applying them
# to a non-numeric measure (e.g. AVG on a VARCHAR column) is always invalid
# and is rejected during query enrichment rather than at SQL execution time.
# min, max, count, count_distinct, first, last work on any type and are NOT
# in this set.
NUMERIC_ONLY_AGGREGATIONS: frozenset[str] = frozenset({
    "sum", "avg", "median", "weighted_avg", "percentile",
    "stddev_samp", "stddev_pop", "var_samp", "var_pop",
    "corr", "covar_samp", "covar_pop",
})


# Default aggregations applicable to a column based on its data type, when the
# column has no explicit ``allowed_aggregations`` whitelist. Used by the engine
# to gate ``column:agg`` expressions (e.g., ``revenue:sum`` requires ``sum`` to
# be eligible for the ``revenue`` column's data type).
_NUMERIC_AGGREGATIONS: frozenset[str] = frozenset({
    "sum", "avg", "min", "max", "count", "count_distinct",
    "median", "weighted_avg", "percentile", "first", "last",
    "stddev_samp", "stddev_pop", "var_samp", "var_pop",
    "corr", "covar_samp", "covar_pop",
})

DEFAULT_AGGREGATIONS_BY_TYPE: dict[DataType, frozenset[str]] = {
    # INT and DOUBLE share the same numeric aggregation set — the type
    # narrowing is for CAST emission, not for what's aggregable. (DEV-1361.)
    DataType.INT: _NUMERIC_AGGREGATIONS,
    DataType.DOUBLE: _NUMERIC_AGGREGATIONS,
    DataType.TEXT: frozenset({
        "count", "count_distinct", "first", "last", "min", "max",
    }),
    DataType.BOOLEAN: frozenset({
        "count", "count_distinct", "sum", "min", "max", "first", "last",
    }),
    DataType.DATE: frozenset({
        "count", "count_distinct", "first", "last", "min", "max",
    }),
    DataType.TIMESTAMP: frozenset({
        "count", "count_distinct", "first", "last", "min", "max",
    }),
}

# Primary-key columns are always restricted to row-counting aggregations,
# regardless of data type. (You can ``count`` customer_ids, but not ``sum`` them.)
PRIMARY_KEY_AGGREGATIONS: frozenset[str] = frozenset({
    "count", "count_distinct",
})
