"""天璇类型系统 — 仅8种固定类型，面向对象接口。

每个类型继承 :class:`ColumnType`，实现:
- ``distance()`` — 为该类型列计算距离矩阵（追加到 LazyFrame）
- ``display()`` — 将值格式化为可读字符串
- ``polars_dtype`` — 映射到 Polars 数据类型
- ``name`` — 中文名称

:data:`TYPE_REGISTRY` 是唯一的类型注册表，按名称索引。
"""
from __future__ import annotations

from abc import ABC, abstractmethod
from typing import Any

import numpy as np
import polars as pl


# ═══════════════════════════════════════════════════════════════════════
#  Base type — all 8 types inherit from this
# ═══════════════════════════════════════════════════════════════════════

class ColumnType(ABC):
    """Base class for all column types.

    Subclasses MUST define:
    - ``name`` (class attribute, Chinese)
    - ``polars_dtype`` (class attribute)
    - ``distance()``
    - ``display()``
    """

    name: str = ""
    polars_dtype: pl.DataType = pl.Utf8

    @abstractmethod
    def distance(self, lf: pl.LazyFrame, columns: list[str]) -> pl.LazyFrame:
        """Compute per-row distances for *columns* of this type.

        Returns a new LazyFrame with distance columns appended
        (prefixed ``_dist_``).
        """
        ...

    @abstractmethod
    def display(self, value: Any) -> str:
        """Format *value* for human-readable display."""
        ...


# ═══════════════════════════════════════════════════════════════════════
#  8 concrete types
# ═══════════════════════════════════════════════════════════════════════

class IPv4Type(ColumnType):
    name = "IPv4"
    polars_dtype = pl.Utf8

    def distance(self, lf: pl.LazyFrame, columns: list[str]) -> pl.LazyFrame:
        if len(columns) < 2:
            return lf
        from analysis.distance import ip_distance
        return ip_distance(lf, columns)

    def display(self, value: Any) -> str:
        return str(value) if value else "-"


class StringType(ColumnType):
    name = "字符串"
    polars_dtype = pl.Utf8

    def distance(self, lf: pl.LazyFrame, columns: list[str]) -> pl.LazyFrame:
        from analysis.distance import string_distance
        return string_distance(lf, columns)

    def display(self, value: Any) -> str:
        return str(value)[:80] if value else "-"


class EnumType(ColumnType):
    name = "枚举"
    polars_dtype = pl.Utf8

    def distance(self, lf: pl.LazyFrame, columns: list[str]) -> pl.LazyFrame:
        # 0-1 match distance: compare to mode
        for col in columns:
            mode_val = lf.select(pl.col(col).mode().first()).collect(streaming=True).item()
            lf = lf.with_columns(
                (pl.col(col) != mode_val).cast(pl.Float64).alias(f'_enum_dist_{col}')
            )
        return lf

    def display(self, value: Any) -> str:
        return str(value) if value else "-"


class FloatType(ColumnType):
    name = "浮点"
    polars_dtype = pl.Float64

    def distance(self, lf: pl.LazyFrame, columns: list[str]) -> pl.LazyFrame:
        # Normalised euclidean distance — handled in preprocessing pipeline
        for col in columns:
            lf = lf.with_columns(
                pl.col(col).cast(pl.Float64).alias(f'_raw_{col}')
            )
        return lf

    def display(self, value: Any) -> str:
        try:
            return f"{float(value):.2f}"
        except (ValueError, TypeError):
            return str(value)


class IntType(ColumnType):
    name = "整数"
    polars_dtype = pl.Int64

    def distance(self, lf: pl.LazyFrame, columns: list[str]) -> pl.LazyFrame:
        # Normalised manhattan distance — handled in preprocessing pipeline
        for col in columns:
            lf = lf.with_columns(
                pl.col(col).cast(pl.Int64).alias(f'_raw_{col}')
            )
        return lf

    def display(self, value: Any) -> str:
        try:
            return str(int(value))
        except (ValueError, TypeError):
            return str(value)


class BoolType(ColumnType):
    name = "布尔"
    polars_dtype = pl.Boolean

    def distance(self, lf: pl.LazyFrame, columns: list[str]) -> pl.LazyFrame:
        # 0-1 match distance
        for col in columns:
            lf = lf.with_columns(
                pl.col(col).cast(pl.Float64).alias(f'_bool_dist_{col}')
            )
        return lf

    def display(self, value: Any) -> str:
        if value is None:
            return "-"
        return "是" if str(value).lower() in ('true', '1', '+', 't') else "否"


class BytesType(ColumnType):
    name = "字节"
    polars_dtype = pl.Utf8

    def distance(self, lf: pl.LazyFrame, columns: list[str]) -> pl.LazyFrame:
        from analysis.distance import bytes_distance
        return bytes_distance(lf, columns)

    def display(self, value: Any) -> str:
        s = str(value).strip() if value else ""
        if len(s) > 20:
            return s[:20] + "..."
        return s or "-"


class TimestampType(ColumnType):
    name = "时间戳"
    polars_dtype = pl.Float64

    def distance(self, lf: pl.LazyFrame, columns: list[str]) -> pl.LazyFrame:
        from analysis.distance import timestamp_fft_distance
        return timestamp_fft_distance(lf, columns)

    def display(self, value: Any) -> str:
        try:
            t = float(value)
            if t < 1e8 or t > 2e12:
                return str(value)
            from datetime import datetime
            return datetime.fromtimestamp(t).strftime("%Y-%m-%d %H:%M")
        except (ValueError, TypeError, OSError):
            return str(value) if value else "-"


# ═══════════════════════════════════════════════════════════════════════
#  Registry — single source of truth
# ═══════════════════════════════════════════════════════════════════════

TYPE_REGISTRY: dict[str, ColumnType] = {
    "枚举": EnumType(),
    "字节": BytesType(),
    "浮点": FloatType(),
    "整数": IntType(),
    "布尔": BoolType(),
    "字符串": StringType(),
    "IPv4": IPv4Type(),
    "时间戳": TimestampType(),
}

# Alias map for config.yaml type strings
TYPE_ALIASES: dict[str, str] = {
    'enum': '枚举', '枚举': '枚举',
    'bytes': '字节', '字节': '字节', 'hex': '字节',
    'float': '浮点', '浮点': '浮点',
    'int': '整数', '整数': '整数',
    'bool': '布尔', '布尔': '布尔', 'boolean': '布尔',
    'string': '字符串', '字符串': '字符串', 'str': '字符串',
    'ipv4': 'IPv4', 'IPv4': 'IPv4',
    'timestamp': '时间戳', '时间戳': '时间戳',
}


def get_type(name: str) -> ColumnType:
    """Resolve a type name (Chinese or English alias) to a ColumnType."""
    resolved = TYPE_ALIASES.get(name.strip().lower() if name.isascii() else name.strip())
    if resolved is None:
        return StringType()  # fallback
    return TYPE_REGISTRY[resolved]


def classify_column(name: str, series: pl.Series, config_type: str | None = None) -> ColumnType:
    """Classify a single column to a ColumnType.

    Priority: config override → TlsDB lookup → name heuristics → value sampling.
    """
    from analysis.tls_ref import get_tls_ref_dict
    import re

    # 1. Config override
    if config_type is not None:
        t = get_type(config_type)
        if t.name != "字符串" or config_type.lower().strip() == 'string':
            return t

    # 2. TlsDB lookup
    tls_ref = get_tls_ref_dict()
    if name in tls_ref:
        t = get_type(tls_ref[name]['data_type'])
        return t

    # 3. Name heuristics
    _patterns: list[tuple[re.Pattern, str]] = [
        (re.compile(r'^(?:src_ip|dst_ip|:ips|:ipd|server.ip|client.ip)$', re.I), 'IPv4'),
        (re.compile(r'^(?::prs|:prd|scnt|dcnt|1ipp|4dbn|4srs|eiph|name|tabl|0ver|cnrs|isrs)$', re.I), '枚举'),
        (re.compile(r'^(?:0cph|0crv|cipher|ecdhe)', re.I), '字节'),
        (re.compile(r'^(?:4dur|8ses|2tmo|0rnt)$', re.I), '浮点'),
        (re.compile(r'^(?:4ksz|8ack|8ppk|8pak|8seq|8did|8byt|row)$', re.I), '整数'),
        (re.compile(r'^(?:timestamp|time|8dbd|created|updated)', re.I), '时间戳'),
        (re.compile(r'\.(ispn|orgn|city|anon|doma|latd|lond)$', re.I), '字符串'),
        (re.compile(r'^(?:cnam|snam|0rnd|source.node|crcc|orga|orgu|@)', re.I), '字符串'),
    ]
    for pat, type_name in _patterns:
        if pat.search(name):
            return TYPE_REGISTRY[type_name]

    # 4. Value-based sampling
    non_null = series.drop_nulls()
    if len(non_null) == 0:
        return TYPE_REGISTRY["字符串"]

    strs = [str(v).strip().lower() for v in non_null.head(100)]
    ip4 = re.compile(r'^\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}$')
    hex_p = re.compile(r'^[0-9a-f]{2}(?:\s[0-9a-f]{2})*$')

    if all(ip4.match(v) for v in strs):
        return TYPE_REGISTRY["IPv4"]
    if all(hex_p.match(v) for v in strs):
        return TYPE_REGISTRY["字节"]
    if series.dtype.is_integer():
        return TYPE_REGISTRY["整数"]
    if series.dtype == pl.Float64:
        return TYPE_REGISTRY["浮点"]
    if series.dtype == pl.Boolean:
        return TYPE_REGISTRY["布尔"]
    if non_null.n_unique() < 50:
        return TYPE_REGISTRY["枚举"]

    return TYPE_REGISTRY["字符串"]


def classify_schema(
    lf: pl.LazyFrame,
    config_overrides: dict[str, str] | None = None,
) -> dict[str, ColumnType]:
    """Classify all columns in a LazyFrame."""
    config_overrides = config_overrides or {}
    df_sample = lf.collect(streaming=True)
    result = {}
    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
