"""列名清理和类型推断，输出 Arrow 列

单元格的类型规则在 cells.py；这里按列汇总：一列的非空值必须全部符合同一类型才转换，否则保留文本，
宁可让模型在 SQL 里 TRY_CAST，也不悄悄丢掉值。各列互不相关，用线程并行推断（Arrow 的计算不占 GIL）
"""

import os
import re
from collections.abc import Sequence
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass
from typing import Any

import numpy as np
import pyarrow as pa
import pyarrow.compute as pc

from app.datafiles.base import Reporter, brief, no_report, where
from app.datafiles.cells import (
    BOOL,
    DATE,
    DATETIME,
    EXACT_FLOAT_INT,
    FAST_RULES,
    FLOAT,
    INT,
    KINDS,
    MICROS_PER_DAY,
    NULL,
    NULL_TEXT,
    PERCENT,
    TEXT,
    Classified,
    as_text,
    classify_piece,
)
from app.datafiles.grid import Column
from app.schemas.data_file import DataColumn, ParseWarning

NUMERIC_KINDS = (INT, FLOAT, PERCENT)
DATE_KINDS = (DATE, DATETIME)
# 同时推断的列数
INFER_THREADS = min(8, os.cpu_count() or 1)


@dataclass
class TypedColumns:
    columns: list[DataColumn]
    table: pa.Table
    # 列名 → 写进列概况的说明，例如保留为文本的原因
    notes: dict[str, str]


def clean_names(
    raw: Sequence[Any], sheet: str | None
) -> tuple[list[str], ParseWarning | None]:
    """去掉空白和换行；空列名按位置命名，重复列名加后缀"""
    names, changes, seen = [], [], set()
    for position, value in enumerate(raw, start=1):
        name = re.sub(r"\s*[\r\n]+\s*", " ", as_text(value)).strip()
        if not name:
            name = f"列{position}"
            changes.append(f"第 {position} 列没有列名，命名为「{name}」")
        if name in seen:
            suffix = 2
            while f"{name}_{suffix}" in seen:
                suffix += 1
            changes.append(f"「{name}」重名，第 {position} 列改为「{name}_{suffix}」")
            name = f"{name}_{suffix}"
        seen.add(name)
        names.append(name)
    if not changes:
        return names, None
    return names, ParseWarning(
        code="column_renamed", message=f"{where(sheet)}{brief(changes, '；')}"
    )


def build_columns(
    names: list[str],
    columns: list[Column],
    sheet: str | None,
    report: Reporter = no_report,
) -> tuple[TypedColumns, ParseWarning | None]:
    """按列推断类型；有列因少数无法解析的值保留为文本时，返回一条 kept_as_text 警告

    columns 会被清空：每列推断完就释放它的原始单元格，大表不必同时留着原始和转换后的两份
    """
    described, arrays, notes = [], [], {}
    with ThreadPoolExecutor(INFER_THREADS) as pool:
        futures = [
            pool.submit(infer_column, name, column)
            for name, column in zip(names, columns)
        ]
        columns.clear()
        for number, (name, future) in enumerate(zip(names, futures), start=1):
            column, array, note = future.result()
            report(number / len(names))
            described.append(column)
            arrays.append(array)
            if note:
                notes[name] = note
    warning = None
    if notes:
        listed = "、".join(f"「{name}」" for name in notes)
        warning = ParseWarning(
            code="kept_as_text",
            message=f"{where(sheet)}{listed}有无法解析的值，整列按文本处理",
        )
    return TypedColumns(
        described, pa.Table.from_arrays(arrays, names=names), notes
    ), warning


def arrow_columns(table: pa.Table) -> tuple[TypedColumns, ParseWarning | None]:
    """Parquet 自带类型，只映射到接口类型；列表、嵌套结构等存成 JSON 字符串"""
    names, warning = clean_names(table.column_names, None)
    columns, arrays = [], []
    for name, array in zip(names, table.columns):
        kind = _arrow_kind(array.type)
        if kind is None:
            values = array.to_pylist()
            array = pa.array(
                [None if v is None else as_text(v) for v in values], pa.string()
            )
            kind = "text"
        elif kind == "text" and pa.types.is_dictionary(array.type):
            array = array.cast(pa.string())
        columns.append(DataColumn(name=name, type=kind))
        arrays.append(array)
    return TypedColumns(columns, pa.Table.from_arrays(arrays, names=names), {}), warning


def _arrow_kind(arrow_type: pa.DataType) -> str | None:
    types = pa.types
    if (
        types.is_integer(arrow_type)
        or types.is_floating(arrow_type)
        or types.is_decimal(arrow_type)
    ):
        return "number"
    if types.is_boolean(arrow_type):
        return "boolean"
    if types.is_date(arrow_type):
        return "date"
    if types.is_timestamp(arrow_type):
        return "datetime"
    if (
        types.is_string(arrow_type)
        or types.is_large_string(arrow_type)
        or types.is_string_view(arrow_type)
    ):
        return "text"
    if types.is_dictionary(arrow_type) and _arrow_kind(arrow_type.value_type) == "text":
        return "text"
    return None


def infer_column(
    name: str, column: Column
) -> tuple[DataColumn, pa.ChunkedArray, str | None]:
    order = list(FAST_RULES)
    parts = [classify_piece(piece, order) for piece in column.pieces]
    counts = sum((part.count() for part in parts), np.zeros(len(KINDS), np.int64))
    present = {kind for kind in range(len(KINDS)) if counts[kind] and kind != NULL}

    def build(make, arrow_type: pa.DataType) -> pa.ChunkedArray:
        return pa.chunked_array([make(part) for part in parts], arrow_type)

    if present and present <= {BOOL}:
        array = build(
            lambda part: _pick(part, "flags", (BOOL,), pa.bool_()), pa.bool_()
        )
        return DataColumn(name=name, type="boolean"), array, None
    if present and present <= {INT, FLOAT}:
        if all(_integral(part) for part in parts):
            array = build(lambda part: _numbers(part, np.int64), pa.int64())
        else:
            array = build(lambda part: _numbers(part, np.float64), pa.float64())
        return DataColumn(name=name, type="number"), array, None
    if present == {PERCENT}:
        array = build(
            lambda part: _pick(part, "reals", (PERCENT,), pa.float64()), pa.float64()
        )
        return DataColumn(name=name, type="number", format="percent"), array, None
    if present and present <= {DATE}:
        array = build(_dates, pa.date32())
        return DataColumn(name=name, type="date"), array, None
    if present and present <= set(DATE_KINDS):
        array = build(
            lambda part: _pick(part, "moments", DATE_KINDS, pa.timestamp("us")),
            pa.timestamp("us"),
        )
        return DataColumn(name=name, type="datetime"), array, None

    array = build(_texts, pa.string())
    return DataColumn(name=name, type="text"), array, _text_note(counts, parts)


def _pick(
    part: Classified, field: str, kinds: tuple[int, ...], arrow_type: pa.DataType
) -> pa.Array:
    """这段里属于 kinds 的值，其余位置为空"""
    keep = np.isin(part.kinds, kinds)
    return pa.array(part.values(field), arrow_type, mask=~keep)


def _integral(part: Classified) -> bool:
    """小数都是整数值并且在 float 能精确表示的范围内（Excel 的整数读出来是 100.0）"""
    reals = part.values("reals")[part.kinds == FLOAT]
    return bool(np.all((np.floor(reals) == reals) & (np.abs(reals) < EXACT_FLOAT_INT)))


def _numbers(part: Classified, dtype) -> pa.Array:
    is_int, is_float = part.kinds == INT, part.kinds == FLOAT
    # 其他位置可能是 NaN（float 段里的空值），先清零再转换
    reals = np.where(is_float, part.values("reals"), 0.0)
    values = np.where(is_int, part.values("ints").astype(dtype), reals.astype(dtype))
    keep = is_int | is_float
    return pa.array(values, pa.from_numpy_dtype(dtype), mask=~keep)


def _dates(part: Classified) -> pa.Array:
    days = (part.values("moments") // MICROS_PER_DAY).astype(np.int32)
    return pa.array(days, pa.date32(), mask=part.kinds != DATE)


def _texts(part: Classified) -> pa.Array:
    """保留为文本：字符串是去掉空白后的原文，其他类型按 as_text 转换；空值记号仍为空"""
    keep = part.kinds != NULL
    if pa.types.is_string(part.piece.type):
        return pc.if_else(pa.array(keep), part.piece, NULL_TEXT)
    return pa.array(
        [
            as_text(value) if kept else None
            for value, kept in zip(part.piece.to_pylist(), keep)
        ],
        pa.string(),
    )


def _text_note(counts: np.ndarray, parts: list[Classified]) -> str | None:
    """大部分值是数值或日期、少数无法解析时，说明原因，模型可以在 SQL 里用 TRY_CAST"""
    failed = int(counts[TEXT])
    numeric = int(sum(counts[kind] for kind in NUMERIC_KINDS))
    dates = int(sum(counts[kind] for kind in DATE_KINDS))
    if not failed or max(numeric, dates) < failed:
        return None
    examples: list[str] = []
    for part in parts:
        failures = _texts(part).filter(pa.array(part.kinds == TEXT))
        for text in pc.unique(failures).to_pylist():
            if text not in examples:
                examples.append(text)
            if len(examples) == 2:
                break
        if len(examples) == 2:
            break
    shown = "、".join(f"「{example}」" for example in examples)
    label = "数值" if numeric >= dates else "日期"
    return f"看起来是{label}，有 {failed} 个值无法解析，例如{shown}"