"""Data type classifier for column-level type inference.

Provides :class:`DataType` enumeration and functions to classify Polars
columns based on name heuristics and value sampling.

Typical usage::

    from analysis.type_classifier import classify_column, classify_schema, DataType

    # Single column
    dtype = classify_column('src_ip', series, config_type='ipv4')

    # Full schema
    type_map = classify_schema(lazyframe, config_overrides={'src_ip': 'ipv4'})
"""
from __future__ import annotations

import re
from enum import Enum

import polars as pl

# ---------------------------------------------------------------------------
#  DataType enumeration
# ---------------------------------------------------------------------------


class DataType(Enum):
    """Enumerated column types recognised by the classifier engine."""

    INT = 1
    FLOAT = 2
    ENUM = 3
    HEX = 4
    URL = 5
    IPv4 = 6
    LAT_LON = 7
    TIMESTAMP = 8
    BOOL_ENUM = 9
    STRING = 10


# ---------------------------------------------------------------------------
#  Detection patterns (raw strings for regex)
# ---------------------------------------------------------------------------

HEX_PATTERN = r'^[0-9a-f]{2}(?:\s[0-9a-f]{2})*$'
"""Matches ``"0e a6 3f"``, ``"ab cd ef"`` etc. (space-separated hex pairs)."""

URL_PATTERN = r'^https?://'
"""Matches URLs starting with ``http://`` or ``https://``."""

IPv4_PATTERN = r'^\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}$'
"""Matches ``"192.168.1.1"`` style IPv4 addresses."""

MAC_PATTERN = r'^([0-9A-Fa-f]{2}[:.-]){5}[0-9A-Fa-f]{2}$'
"""Matches MAC addresses with ``:``, ``-``, or ``.`` separators."""

BOOL_ENUM_VALUES = (
    '', ' ', '+', '-', '0', '1',
    'true', 'false', 'yes', 'no', 't', 'f',
)
"""String values considered boolean / binary indicators."""

LAT_LON_NAMES = ('lat', 'latitude', 'lon', 'lng', 'longitude')

# Name-to-type heuristics (checked after config override, before values)
_NAME_TYPE_MAP: list[tuple[re.Pattern, DataType]] = [
    # User column names - EXACT match only (no generic fuzzy patterns)
    # IPv4 columns
    (re.compile(r'^:ips$|^:ipd$|^server.ip$|^client.ip$', re.I), DataType.IPv4),
    # ENUM columns
    (re.compile(r'^:prs$|^:prd$', re.I), DataType.ENUM),
    (re.compile(r'^(?:scnt|dcnt|cnrs|isrs|1ipp|4dbn)$', re.I), DataType.ENUM),
    (re.compile(r'^0ver$', re.I), DataType.ENUM),
    # FLOAT columns
    (re.compile(r'^4dur$|^8ses$|^2tmo$|^0rnt$', re.I), DataType.FLOAT),
    # INT columns
    (re.compile(r'^4ksz$|^8ack$|^8ppk$|^row$', re.I), DataType.INT),
    # HEX columns
    (re.compile(r'^0cph$|^0crv$', re.I), DataType.HEX),
    # TIMESTAMP columns
    (re.compile(r'^8dbd$|^time$|^timestamp$', re.I), DataType.TIMESTAMP),
    # Dot-notation suffixes for lat/lon (e.g. :ips.latd, :ipd.lond)
    (re.compile(r'\.latd$', re.I), DataType.LAT_LON),
    (re.compile(r'\.lond$', re.I), DataType.LAT_LON),
    # Dot-notation STRING suffixes (e.g. :ips.ispn, :ipd.city)
    (re.compile(r'\.(ispn|orgn|city)$', re.I), DataType.STRING),
    # Standalone string columns - exact match
    (re.compile(r'^(?:cnam|snam|tabl|name|0rnd|source.node|ecdhe.named.curve)$', re.I), DataType.STRING),
]

# Mapping from config.yaml type strings to DataType
_CONFIG_TYPE_MAP: dict[str, DataType] = {
    'int': DataType.INT,
    'float': DataType.FLOAT,
    'enum': DataType.ENUM,
    'hex': DataType.HEX,
    'url': DataType.URL,
    'ipv4': DataType.IPv4,
    'lat_lon': DataType.LAT_LON,
    'timestamp': DataType.TIMESTAMP,
    'bool_enum': DataType.BOOL_ENUM,
    'string': DataType.STRING,
}


# ---------------------------------------------------------------------------
#  Column classification (single column)
# ---------------------------------------------------------------------------


def _name_to_type(name: str) -> DataType | None:
    """Match column name against known heuristics.

    Returns a :class:`DataType` if the name matches strongly, otherwise
    ``None`` (fall through to value-based detection).
    """
    for pattern, dtype in _NAME_TYPE_MAP:
        if pattern.search(name):
            return dtype
    return None


def _is_port_value(v: str) -> bool:
    """Check if *v* is an integer string in the registered port range (1-65535)."""
    try:
        n = int(v)
        return 1 <= n <= 65535
    except ValueError:
        return False


def _can_float(v: str) -> bool:
    """Return ``True`` when *v* can be parsed as a Python float."""
    try:
        float(v)
        return True
    except ValueError:
        return False


def _values_to_type(series: pl.Series) -> DataType:
    """Infer column type by sampling non-null values.

    Detection order: IPv4 → MAC guard → HEX → URL → PORT → BOOL_ENUM
                     → LAT_LON → ENUM → FLOAT → STRING
    """
    dtype = series.dtype

    # Boolean dtype ──────────────────────────────────────────────────────
    if dtype == pl.Boolean:
        return DataType.BOOL_ENUM

    # Numeric dtypes ─────────────────────────────────────────────────────
    if dtype in (pl.Float32, pl.Float64):
        # Check LAT_LON possibility before returning FLOAT
        non_null = series.drop_nulls()
        if len(non_null) > 0:
            sample = non_null.head(100)
            sample_strs = [str(v).strip().lower() for v in sample]
            non_blank = [v for v in sample_strs if v not in ('', ' ', '+', '-')]
            if non_blank:
                try:
                    floats = [float(v) for v in non_blank]
                    max_abs = max(abs(f) for f in floats)
                    if (
                        all(-180 <= f <= 180 for f in floats)
                        and any('.' in v for v in non_blank)
                        and max_abs > 10
                    ):
                        return DataType.LAT_LON
                except ValueError:
                    pass
        return DataType.FLOAT

    if dtype in (pl.Int8, pl.Int16, pl.Int32, pl.Int64,
                 pl.UInt8, pl.UInt16, pl.UInt32, pl.UInt64):
        return DataType.INT

    # String / categorical types ─────────────────────────────────────────
    non_null = series.drop_nulls()
    if len(non_null) == 0:
        return DataType.STRING

    sample = non_null.head(100)
    # Convert all to string representations for pattern matching
    sample_strs = [str(v).strip().lower() for v in sample]

    # ── 1. IPv4 detection ──────────────────────────────────────────────
    if all(re.match(IPv4_PATTERN, v) for v in sample_strs):
        if all(_valid_ip_octets(v) for v in sample_strs):
            return DataType.IPv4

    # ── 2. MAC guard (skip HEX if values look like MAC addresses) ──────
    is_mac = all(re.match(MAC_PATTERN, v) for v in sample_strs)

    # ── 3. HEX detection ───────────────────────────────────────────────
    if not is_mac and all(re.match(HEX_PATTERN, v) for v in sample_strs):
        return DataType.HEX

    # ── 4. URL detection ───────────────────────────────────────────────
    if all(re.match(URL_PATTERN, v) for v in sample_strs):
        return DataType.URL

    # ── 5. PORT detection → ENUM ───────────────────────────────────────
    if all(_is_port_value(v) for v in sample_strs):
        return DataType.ENUM

    # ── 6. BOOL_ENUM detection ─────────────────────────────────────────
    if all(v in BOOL_ENUM_VALUES for v in sample_strs):
        return DataType.BOOL_ENUM

    # ── 7. LAT_LON detection (value-based) ──────────────────────────────
    non_blank = [v for v in sample_strs if v not in ('', ' ', '+', '-')]
    if non_blank:
        try:
            floats = [float(v) for v in non_blank]
            max_abs = max(abs(f) for f in floats)
            if (
                all(-180 <= f <= 180 for f in floats)
                and any('.' in v for v in non_blank)
                and max_abs > 10
            ):
                return DataType.LAT_LON
        except ValueError:
            pass

    # ── 8. ENUM detection (few unique values) ──────────────────────────
    unique_count = non_null.n_unique()
    if unique_count < 20:
        return DataType.ENUM

    # ── 9. FLOAT detection (>80 % can float) ───────────────────────────
    float_count = sum(1 for v in sample_strs if _can_float(v))
    if float_count / len(sample_strs) > 0.8:
        # Guard: values with leading "+"/"-" (e.g., "+5", "-abc") that
        # happen to parse as float should not force a FLOAT classification
        # when non-float values exist in the column.  This catches string
        # columns where a "+" prefix is used as a text marker (phone codes,
        # IDs, error codes) rather than a numeric sign.
        if any(v.startswith('+') or v.startswith('-') for v in sample_strs):
            if float_count < len(sample_strs):
                return DataType.STRING
        return DataType.FLOAT

    return DataType.STRING


def _valid_ip_octets(ip_str: str) -> bool:
    """Check that each octet in an IPv4 string is 0-255."""
    parts = ip_str.split('.')
    if len(parts) != 4:
        return False
    for p in parts:
        try:
            n = int(p)
            if n < 0 or n > 255:
                return False
        except ValueError:
            return False
    return True


def classify_column(
    name: str,
    series: pl.Series,
    config_type: str | None = None,
) -> DataType:
    """Classify a single column based on its *name* and *series* values.

    Priority
    --------
    1. **config_type** — user override from ``config.yaml``.
       If provided (e.g. ``'ipv4'``), returns the corresponding
       :class:`DataType` immediately.
    2. **Value-based** pattern matching — samples up to 100 non-null
       values and tests against IPv4 → HEX → URL → PORT → BOOL_ENUM
       → LAT_LON → ENUM → FLOAT in order.
    3. **Name-based** heuristics — only consulted when value-based
       detection returned ``STRING`` (i.e. no value pattern matched).

    Parameters
    ----------
    name:
        Column name (used for name-based heuristics).
    series:
        Polars :class:`Series` containing the column's data.
    config_type:
        Optional user override (case-insensitive, matched to
        :data:`_CONFIG_TYPE_MAP`).

    Returns
    -------
    DataType
    """
    # 1. Config override (highest priority)
    if config_type is not None:
        mapped = _CONFIG_TYPE_MAP.get(config_type.lower().strip())
        if mapped is not None:
            return mapped

    # 2. Value-based sampling
    result = _values_to_type(series)

    # 3. Name-based heuristics (only when values returned STRING)
    if result == DataType.STRING:
        name_match = _name_to_type(name)
        if name_match is not None:
            return name_match

    return result


# ---------------------------------------------------------------------------
#  Schema classification (all columns)
# ---------------------------------------------------------------------------


def classify_schema(
    lf: pl.LazyFrame,
    config_overrides: dict[str, str] | None = None,
) -> dict[str, DataType]:
    """Classify all columns in a :class:`LazyFrame`.

    Parameters
    ----------
    lf:
        A Polars :class:`LazyFrame` whose columns to classify.
    config_overrides:
        Optional mapping of ``{col_name: type_string}`` from a
        ``config.yaml`` ``columns:`` section.  These take highest
        priority (see :func:`classify_column`).

    Returns
    -------
    dict[str, DataType]
        Mapping from column name to inferred :class:`DataType`.
    """
    config_overrides = config_overrides or {}
    df_sample = lf.collect(streaming=True)
    result: dict[str, DataType] = {}
    for col_name in df_sample.columns:
        config_type = config_overrides.get(col_name)
        result[col_name] = classify_column(
            name=col_name,
            series=df_sample[col_name],
            config_type=config_type,
        )
    return result