"""单元格的类型规则：单个值用 classify 判断，一段 Arrow 数组用 classify_piece 批量判断

classify 是规则的唯一定义：一个值是什么类型、转换成什么值。大表不能逐个调用它，classify_piece 对常见写法
（空值记号、布尔、ASCII 数字和百分数、编号、ISO 日期和常见的日期时间）用正则和 Arrow 转换批量得出结论，
其余的值去重后逐个交给 classify。批量规则只认它有把握的写法，认不出的一律交给 classify，所以结论和逐个
调用 classify 相同；测试里逐值对比两者
"""

import json
import math
import re
from dataclasses import dataclass
from datetime import UTC, date, datetime, time
from typing import Any

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

NULL_TOKENS = {"", "-", "--", "—", "——", "N/A", "n/a", "NULL", "null", "#N/A"}
EXCEL_ERROR = re.compile(r"^#(?:NULL!|DIV/0!|VALUE!|REF!|NAME\?|NUM!|N/A)$")

CURRENCY_SYMBOLS = "¥￥$€"
NUMBER = re.compile(r"^[+-]?(?:\d+|\d{1,3}(?:,\d{3})+)(?:\.\d+)?(?:[eE][+-]?\d+)?$")
INTEGER = re.compile(r"^[+-]?\d+$")
# 16 位及以上的纯数字是证件号、卡号一类的编号，转成数值会丢失精度
MIN_ID_DIGITS = 16

# 按年份后的第一个字符选格式；%m、%d 也接受一位数。没有对应格式的直接判为不是日期：
# strptime 每次失败都要抛异常，一个值试完所有格式约 16 μs，还要抢一把模块级的锁，多线程时更慢
DATE_FORMATS = {
    "-": ("%Y-%m-%d", "%Y-%m"),
    "/": ("%Y/%m/%d", "%Y/%m"),
    ".": ("%Y.%m.%d",),
    "年": ("%Y年%m月%d日", "%Y年%m月"),
}
DATETIME_FORMATS = {
    "-": ("%Y-%m-%d %H:%M:%S", "%Y-%m-%d %H:%M"),
    "/": ("%Y/%m/%d %H:%M:%S", "%Y/%m/%d %H:%M"),
}

INT64_RANGE = range(-(2**63), 2**63)
# 超过这个范围的整数存成 float 会丢精度
EXACT_FLOAT_INT = 2**53

# str.isspace() 为真的全部字符，也就是 str.strip() 去掉的字符；Arrow 自带的去空白函数按另一套定义，
# 读取时显式传入这些字符，结果才和 Python 一致
WHITESPACE = "\t\n\x0b\x0c\r\x1c\x1d\x1e\x1f \x85\xa0            

  　"
# 正则里匹配上面这些字符的写法（RE2 的 \s 只认 ASCII 空白）
WHITESPACE_CLASS = "".join(f"\\x{{{ord(char):x}}}" for char in WHITESPACE)

# 类型代码，classify 返回的类型名按这个顺序编号
KINDS = ("null", "bool", "int", "float", "percent", "date", "datetime", "text")
NULL, BOOL, INT, FLOAT, PERCENT, DATE, DATETIME, TEXT = range(len(KINDS))
KIND_CODES = {name: code for code, name in enumerate(KINDS)}

MICROS_PER_DAY = 86_400_000_000
EPOCH = datetime(1970, 1, 1)


def classify(value: Any) -> tuple[str, Any]:
    """单个值的类型和转换后的值；类型是 KINDS 之一"""
    if value is None:
        return "null", None
    if isinstance(value, bool):
        return "bool", value
    if isinstance(value, int):
        return ("int", value) if value in INT64_RANGE else ("text", str(value))
    if isinstance(value, float):
        return (
            ("null", None)
            if math.isnan(value) or math.isinf(value)
            else ("float", value)
        )
    if isinstance(value, datetime):
        # Excel 只设了日期格式的单元格读出来是零点的 datetime
        if value.time() == time() and value.tzinfo is None:
            return "date", value.date()
        return "datetime", _naive(value)
    if isinstance(value, date):
        return "date", value
    if not isinstance(value, str):
        return "text", as_text(value)

    text = value.strip()
    if text in NULL_TOKENS or EXCEL_ERROR.match(text):
        return "null", None
    if text.lower() in ("true", "false"):
        return "bool", text.lower() == "true"
    number = parse_number(text)
    if number is not None:
        if isinstance(number, int):
            # 带千分位的超长整数和超出范围的原生整数一样保留文本
            return ("int", number) if number in INT64_RANGE else ("text", text)
        return "float", number
    if text.endswith("%"):
        number = parse_number(text[:-1].rstrip())
        if number is not None:
            return "percent", number / 100
    moment = parse_datetime(text)
    if moment is not None:
        return ("datetime" if isinstance(moment, datetime) else "date"), moment
    return "text", text


def parse_number(text: str) -> int | float | None:
    """千分位和货币符号可以有；有前导零或 16 位以上的数字串是编号，不算数值

    判断数字一律用 isdecimal：isdigit 对 ²、① 也返回 True，但 int() 解析不了它们
    """
    if text and text[0] in CURRENCY_SYMBOLS:
        text = text[1:].lstrip()
    # 最常见的纯数字不用正则
    plain = text.isdecimal()
    if not plain and not NUMBER.match(text):
        return None
    digits = text.lstrip("+-")
    if len(digits) > 1 and digits[0] == "0" and digits[1].isdecimal():
        return None
    if digits.isdecimal() and len(digits) >= MIN_ID_DIGITS:
        return None
    if plain:
        return int(text)
    text = text.replace(",", "")
    return int(text) if INTEGER.match(text) else float(text)


def parse_datetime(text: str) -> date | datetime | None:
    if len(text) < 6 or not text[:4].isdecimal():
        return None
    # 标准 ISO 写法先用 fromisoformat，比 strptime 快一个数量级
    if ":" in text:
        try:
            return _naive(datetime.fromisoformat(text))
        except ValueError:
            pass
        for pattern in DATETIME_FORMATS.get(text[4], ()):
            try:
                return datetime.strptime(text, pattern)
            except ValueError:
                continue
        return None
    if len(text) == 10 and text[4] == "-":
        try:
            return date.fromisoformat(text)
        except ValueError:
            pass
    for pattern in DATE_FORMATS.get(text[4], ()):
        try:
            return datetime.strptime(text, pattern).date()
        except ValueError:
            continue
    return None


def as_text(value: Any) -> str:
    """单元格转成文本：列名、保留为文本的值、警告里引用的内容都用它；Excel 里年份等数字读出来是 2025.0"""
    if value is None:
        return ""
    if isinstance(value, str):
        return value.strip()
    if isinstance(value, bool):
        return "true" if value else "false"
    if isinstance(value, float) and value.is_integer():
        return str(int(value))
    if isinstance(value, datetime):
        return value.isoformat(sep=" ")
    if isinstance(value, (date, time)):
        return value.isoformat()
    if isinstance(value, (dict, list)):
        # Parquet 的嵌套值里可能有日期、Decimal、bytes，JSON 表示不了的写成字符串
        return json.dumps(value, ensure_ascii=False, default=str)
    return str(value)


def _naive(value: datetime) -> datetime:
    # 带时区的时间统一换成 UTC 再去掉时区，同一列里的值才能比较
    if value.tzinfo is None:
        return value
    return value.astimezone(UTC).replace(tzinfo=None)


# ---- 批量判断 ----

# classify 只会把这些开头的字符串判成文本以外的类型（空值记号和布尔另外判断）：
# 货币符号，或可选的正负号加数字（数值、百分数、日期都以数字开头）
CANDIDATE = r"^(?:[¥￥$€]|[+-]?\p{Nd})"
# 最多 15 位整数部分的 ASCII 数值，保证整数在 int64 内、百分数除以 100 与 Python 结果相同；
# 有前导零的（编号）、货币符号后带空格的等写法不在这里，交给 classify
_DIGITS = r"[¥￥$€]?[+-]?(?:0|[1-9][0-9]{0,14}|[1-9][0-9]{0,2}(?:,[0-9]{3}){1,4})(?:\.[0-9]+)?(?:[eE][+-]?[0-9]{1,3})?"
FAST_NUMBER = f"^{_DIGITS}$"
FAST_PERCENT = f"^{_DIGITS}%$"
# 有前导零或 16 位以上的纯数字是编号：classify 判成文本
FAST_ID = r"^0[0-9]+$|^[0-9]{16,}$"
FAST_DATE = r"^[1-9][0-9]{3}-[0-9]{2}-[0-9]{2}$"
# fromisoformat 能解析的写法，以及 DATETIME_FORMATS 里分隔符为 - 或 /、各部分一两位数字的写法
FAST_DATETIMES = (
    r"^[1-9][0-9]{3}-[0-9]{2}-[0-9]{2}[T ][0-9]{2}:[0-9]{2}(?::[0-9]{2}(?:\.[0-9]{3}(?:[0-9]{3})?)?)?$",
    r"^[1-9][0-9]{3}-[0-9]{1,2}-[0-9]{1,2} [0-9]{1,2}:[0-9]{1,2}(?::[0-9]{1,2})?$",
    r"^[1-9][0-9]{3}/[0-9]{1,2}/[0-9]{1,2} [0-9]{1,2}:[0-9]{1,2}(?::[0-9]{1,2})?$",
)

NULL_SET = pa.array(sorted(NULL_TOKENS))
BOOL_SET = pa.array(["true", "false"])
NULL_TEXT = pa.scalar(None, pa.string())


@dataclass
class Classified:
    """一段单元格逐个分类的结果，和对每个值调用 classify 的结论相同

    kinds 是每个值的类型代码；ints、reals、moments、flags 和这段等长，只在对应类型的位置上有意义，
    这段里没有这种类型时为空
    """

    piece: pa.Array
    kinds: np.ndarray
    # int
    ints: np.ndarray | None = None
    # float、percent（已除以 100）
    reals: np.ndarray | None = None
    # date、datetime：自 1970 年起的微秒数，日期按零点
    moments: np.ndarray | None = None
    # bool
    flags: np.ndarray | None = None

    def values(self, name: str) -> np.ndarray:
        """取某类转换后的值，没有时补零"""
        array = getattr(self, name)
        if array is None:
            array = np.zeros(len(self.kinds), _DTYPES[name])
            setattr(self, name, array)
        return array

    def count(self) -> np.ndarray:
        return np.bincount(self.kinds, minlength=len(KINDS))


_DTYPES = {"ints": np.int64, "reals": np.float64, "moments": np.int64, "flags": bool}


def union_layout(piece: pa.UnionArray) -> tuple[np.ndarray, np.ndarray]:
    """dense union 每个位置属于第几个子数组、是子数组里的第几个

    切片后的 union，pyarrow 的 type_codes、offsets 不计切片的偏移，所以直接读缓冲区
    """
    start, stop = piece.offset, piece.offset + len(piece)
    _, codes, offsets = piece.buffers()[:3]
    return (
        np.frombuffer(codes, np.int8)[start:stop],
        np.frombuffer(offsets, np.int32)[start:stop],
    )


def mask(values: pa.Array | pa.ChunkedArray) -> np.ndarray:
    """Arrow 布尔数组转成 numpy，空值当作 False"""
    return pc.fill_null(values, False).to_numpy(zero_copy_only=False)


def classify_piece(piece: pa.Array, order: list | None = None) -> Classified:
    """批量判断一段单元格；段的类型是读取时产生的几种之一，见 grid.py

    order 是字符串段依次尝试的批量规则，同一列的各段传同一个列表：这些规则认的写法互不重叠，先试哪个
    结论都一样，所以每段之后把认出最多的规则挪到前面，后面的段大多一次就认完
    """
    kind = piece.type
    if pa.types.is_union(kind):
        return _classify_union(piece, list(FAST_RULES) if order is None else order)
    valid = mask(piece.is_valid())
    if pa.types.is_string(kind):
        return _classify_strings(
            piece, valid, list(FAST_RULES) if order is None else order
        )
    result = Classified(piece, np.where(valid, TEXT, NULL).astype(np.int8))
    if pa.types.is_null(kind):
        return result
    if pa.types.is_boolean(kind):
        result.kinds[valid] = BOOL
        result.flags = mask(piece)
    elif pa.types.is_integer(kind):
        result.kinds[valid] = INT
        result.ints = pc.fill_null(piece, 0).to_numpy()
    elif pa.types.is_floating(kind):
        reals = pc.fill_null(piece, 0.0).to_numpy()
        result.kinds[valid] = np.where(np.isfinite(reals[valid]), FLOAT, NULL)
        result.reals = reals
    elif pa.types.is_date32(kind):
        result.kinds[valid] = DATE
        result.moments = (
            pc.fill_null(piece, 0).cast(pa.int32()).to_numpy().astype(np.int64)
            * MICROS_PER_DAY
        )
    elif pa.types.is_timestamp(kind):
        moments = pc.fill_null(piece.cast(pa.int64()), 0).to_numpy()
        midnight = moments % MICROS_PER_DAY == 0
        result.kinds[valid] = np.where(midnight[valid], DATE, DATETIME)
        result.moments = moments
    else:
        raise TypeError(f"不支持的单元格类型：{kind}")
    return result


def _classify_union(piece: pa.UnionArray, order: list) -> Classified:
    """类型混杂的段（见 grid.py）：每个子数组各判断一次，再按位置拼回"""
    result = Classified(piece, np.full(len(piece), NULL, np.int8))
    codes, offsets = union_layout(piece)
    for code in range(piece.type.num_fields):
        at = np.flatnonzero(codes == code)
        if not len(at):
            continue
        child = classify_piece(piece.field(code), order)
        rows = offsets[at]
        result.kinds[at] = child.kinds[rows]
        for name in _DTYPES:
            values = getattr(child, name)
            if values is not None:
                result.values(name)[at] = values[rows]
    return result


def _classify_strings(piece: pa.Array, valid: np.ndarray, order: list) -> Classified:
    """字符串段：按 classify 的顺序（空值记号 → 布尔 → 数值 → 百分数 → 日期时间 → 文本）批量判断"""
    result = Classified(piece, np.where(valid, TEXT, NULL).astype(np.int8))
    # 空值记号、Excel 错误值、布尔都不以货币符号或数字开头，只在其余的值里找；
    # 这些值剩下的就是文本，类型已经记好。数值、日期列整段都是候选，省掉这几步
    pending = valid & mask(pc.match_substring_regex(piece, CANDIDATE))
    others = valid & ~pending
    if others.any():
        hit = others & mask(
            pc.or_(
                pc.is_in(piece, value_set=NULL_SET),
                pc.match_substring_regex(piece, EXCEL_ERROR.pattern),
            )
        )
        result.kinds[hit] = NULL
        others &= ~hit
    if others.any():
        lower = pc.utf8_lower(piece)
        hit = others & mask(pc.is_in(lower, value_set=BOOL_SET))
        if hit.any():
            result.kinds[hit] = BOOL
            result.values("flags")[hit] = mask(pc.equal(lower, "true"))[hit]
    claimed = {}
    for rule in order:
        left = int(pending.sum())
        if not left:
            break
        rule(piece, result, pending)
        claimed[rule] = left - int(pending.sum())
    order.sort(key=lambda rule: -claimed.get(rule, 0))
    if pending.any():
        _delegate(piece, result, pending)
    return result


def _matched(piece: pa.Array, pattern: str, pending: np.ndarray) -> np.ndarray:
    """还没有结论、并且符合 pattern 的值的位置"""
    return np.flatnonzero(pending & mask(pc.match_substring_regex(piece, pattern)))


def _plain_number(strings: pa.Array) -> pa.Array:
    """去掉货币符号、千分位和正号，剩下 Arrow 能直接转换的数字

    只用于符合 FAST_NUMBER 的值：货币符号只在开头，正号在开头或指数里（1e+5 去掉后是 1e5，值不变）
    """
    return pc.replace_substring_regex(strings, r"[¥￥$€,+]", "")


def _fast_numbers(piece: pa.Array, result: Classified, pending: np.ndarray) -> None:
    rows = _matched(piece, FAST_NUMBER, pending)
    if not len(rows):
        return
    plain = _plain_number(piece.take(rows))
    integer = mask(pc.match_substring_regex(plain, r"^-?[0-9]+$"))
    if integer.any():
        at = rows[integer]
        result.kinds[at] = INT
        result.values("ints")[at] = pc.cast(
            plain.filter(integer), pa.int64()
        ).to_numpy()
    if not integer.all():
        at = rows[~integer]
        result.kinds[at] = FLOAT
        result.values("reals")[at] = pc.cast(
            plain.filter(~integer), pa.float64()
        ).to_numpy()
    pending[rows] = False


def _fast_percents(piece: pa.Array, result: Classified, pending: np.ndarray) -> None:
    rows = _matched(piece, FAST_PERCENT, pending)
    if not len(rows):
        return
    plain = _plain_number(pc.utf8_slice_codeunits(piece.take(rows), 0, -1))
    result.kinds[rows] = PERCENT
    result.values("reals")[rows] = pc.cast(plain, pa.float64()).to_numpy() / 100
    pending[rows] = False


def _fast_ids(piece: pa.Array, result: Classified, pending: np.ndarray) -> None:
    # 类型已经记为文本，只需标记为已判断
    pending[_matched(piece, FAST_ID, pending)] = False


def _fast_dates(piece: pa.Array, result: Classified, pending: np.ndarray) -> None:
    rows = _matched(piece, FAST_DATE, pending)
    if len(rows):
        _put_moments(result, pending, rows, piece.take(rows), pa.date32(), DATE)


def _fast_datetimes(piece: pa.Array, result: Classified, pending: np.ndarray) -> None:
    for pattern in FAST_DATETIMES:
        rows = _matched(piece, pattern, pending)
        if not len(rows):
            continue
        # 统一成 Arrow 认识的 ISO 写法：/ 换成 -，一位数补零（相邻的两位数各占一次，所以替换两遍）
        strings = pc.replace_substring(piece.take(rows), "/", "-")
        for _ in range(2):
            strings = pc.replace_substring_regex(
                strings, r"([-: ])([0-9])([-: ]|$)", r"\10\2\3"
            )
        _put_moments(result, pending, rows, strings, pa.timestamp("us"), DATETIME)


def _put_moments(
    result: Classified,
    pending: np.ndarray,
    rows: np.ndarray,
    strings: pa.Array,
    arrow_type: pa.DataType,
    kind: int,
) -> None:
    """转换成日期或时间并填回；转换不了的（2 月 30 日、25 点）留给 classify"""
    ok = np.zeros(len(rows), bool)
    moments = np.zeros(len(rows), np.int64)
    _cast_parts(strings, arrow_type, 0, ok, moments)
    at = rows[ok]
    result.kinds[at] = kind
    result.values("moments")[at] = moments[ok]
    pending[at] = False


# 二分到这么短还转换失败，就整段交给 classify
MIN_CAST_PART = 64


def _cast_parts(
    strings: pa.Array,
    arrow_type: pa.DataType,
    offset: int,
    ok: np.ndarray,
    moments: np.ndarray,
) -> None:
    """Arrow 的转换遇到一个坏值就整体失败；失败时对半分开重试，只把含坏值的小段留下"""
    try:
        values = pc.cast(strings, arrow_type)
    except pa.ArrowInvalid:
        if len(strings) <= MIN_CAST_PART:
            return
        half = len(strings) // 2
        _cast_parts(strings.slice(0, half), arrow_type, offset, ok, moments)
        _cast_parts(strings.slice(half), arrow_type, offset + half, ok, moments)
        return
    end = offset + len(strings)
    ok[offset:end] = True
    if pa.types.is_date32(arrow_type):
        moments[offset:end] = (
            values.cast(pa.int32()).to_numpy().astype(np.int64) * MICROS_PER_DAY
        )
    else:
        moments[offset:end] = values.cast(pa.int64()).to_numpy()


# 批量规则；认的写法互不重叠，顺序不影响结论
FAST_RULES = (_fast_numbers, _fast_percents, _fast_ids, _fast_dates, _fast_datetimes)


def _delegate(piece: pa.Array, result: Classified, pending: np.ndarray) -> None:
    """批量规则认不出的值去重后逐个交给 classify，再按位置填回"""
    rows = np.flatnonzero(pending)
    values = piece.take(rows)
    unique = pc.unique(values)
    codes = pc.index_in(values, value_set=unique).to_numpy()
    outcomes = [classify(value) for value in unique.to_pylist()]
    kinds = np.array([KIND_CODES[kind] for kind, _ in outcomes], np.int8)
    result.kinds[rows] = kinds[codes]
    for name, wanted, convert in _FIELDS:
        chosen = np.isin(kinds, wanted)
        if not chosen.any():
            continue
        converted = np.array(
            [
                convert(value) if pick else 0
                for (_, value), pick in zip(outcomes, chosen)
            ],
            _DTYPES[name],
        )
        at = chosen[codes]
        result.values(name)[rows[at]] = converted[codes[at]]
    pending[rows] = False


def _micros(value: date) -> int:
    if isinstance(value, datetime):
        return (value - EPOCH) // _MICROSECOND
    return (value - EPOCH.date()).days * MICROS_PER_DAY


_MICROSECOND = datetime(1970, 1, 1, microsecond=1) - EPOCH
# classify 的结果填进 Classified 的哪个字段：字段名、对应的类型代码、怎样转换
_FIELDS = (
    ("ints", (INT,), int),
    ("reals", (FLOAT, PERCENT), float),
    ("moments", (DATE, DATETIME), _micros),
    ("flags", (BOOL,), bool),
)