"""按列存放的单元格，读取、拆表和类型推断共用

每列是若干段 Arrow 数组，每段一种类型：字符串（已去掉首尾空白，空串记为空值）、布尔、int64、float64、
date32、timestamp[us]，或全是空值。读取时每 CHUNK_ROWS 行转一次，同一时间只有一个块是 Python 对象；
块里同一列的类型混杂时（表头和数值在同一列、数值列里夹着"待定"），这一段存成 Arrow 的 dense union：
每种类型一个子数组，原值和原类型都保留。Arrow 的 is_valid 等函数不认 union 的子数组里的空值，
所以按段处理的地方都经过 each，对子数组分别计算再按位置拼回
"""

from collections.abc import Iterable, Sequence
from dataclasses import dataclass
from datetime import date, datetime, time
from itertools import zip_longest
from typing import Any

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

from app.datafiles.cells import (
    EXACT_FLOAT_INT,
    INT64_RANGE,
    NULL_TEXT,
    WHITESPACE,
    as_text,
    mask,
    union_layout,
)

CHUNK_ROWS = 65_536

NoneType = type(None)


class Column:
    def __init__(self, pieces: Sequence[pa.Array]):
        self.pieces = [piece for piece in pieces if len(piece)]
        self.offsets = np.cumsum([0, *map(len, self.pieces)])

    @classmethod
    def from_values(cls, values: Sequence[Any]) -> "Column":
        return cls([_convert(values)])

    def __len__(self) -> int:
        return int(self.offsets[-1])

    def present(self) -> np.ndarray:
        """每个单元格是否有值"""
        return _join(
            [
                each(piece, lambda part: mask(part.is_valid()), bool)
                for piece in self.pieces
            ],
            bool,
        )

    def matches(self, test) -> np.ndarray:
        """字符串单元格中 test（Arrow 字符串数组 → 布尔数组）为真的位置；其他类型的单元格都不算"""

        def check(part: pa.Array) -> np.ndarray:
            if pa.types.is_string(part.type):
                return mask(test(part))
            return np.zeros(len(part), bool)

        return _join([each(piece, check, bool) for piece in self.pieces], bool)

    def slice(self, start: int, stop: int) -> "Column":
        pieces = []
        for piece, offset in zip(self.pieces, self.offsets):
            begin, end = max(start - offset, 0), min(stop - offset, len(piece))
            if begin < end:
                pieces.append(piece.slice(begin, end - begin))
        return Column(pieces)

    def take(self, rows: np.ndarray) -> "Column":
        """按行号（升序）取出若干行；行号连续时不复制数据"""
        if len(rows) and rows[-1] - rows[0] + 1 == len(rows):
            return self.slice(int(rows[0]), int(rows[-1]) + 1)
        bounds = np.searchsorted(rows, self.offsets)
        return Column(
            [
                piece.take(pa.array(rows[bounds[i] : bounds[i + 1]] - self.offsets[i]))
                for i, piece in enumerate(self.pieces)
            ]
        )

    def to_pylist(self) -> list[Any]:
        return [value for piece in self.pieces for value in piece.to_pylist()]


@dataclass
class Grid:
    """不带表头的原始网格；第 i 行是文件或 sheet 中的第 i + 1 行，各列一样长"""

    sheet: str | None
    columns: list[Column]
    height: int

    @classmethod
    def from_rows(cls, sheet: str | None, rows: Iterable[Sequence[Any]]) -> "Grid":
        builder = GridBuilder()
        for row in rows:
            builder.add(row)
        return builder.finish(sheet)

    def rows(
        self, start: int, stop: int, columns: Sequence[int] | None = None
    ) -> list[list[Any]]:
        """[start, stop) 行的 Python 值，只取 columns 这些列"""
        return self._rows(
            [column.slice(start, stop) for column in self._pick(columns)], stop - start
        )

    def rows_at(
        self, rows: np.ndarray, columns: Sequence[int] | None = None
    ) -> list[list[Any]]:
        return self._rows(
            [column.take(rows) for column in self._pick(columns)], len(rows)
        )

    def _pick(self, columns: Sequence[int] | None) -> list[Column]:
        return self.columns if columns is None else [self.columns[c] for c in columns]

    @staticmethod
    def _rows(columns: list[Column], count: int) -> list[list[Any]]:
        if not columns:
            return [[] for _ in range(count)]
        return [list(row) for row in zip(*(column.to_pylist() for column in columns))]


class GridBuilder:
    """逐行接收单元格，每 CHUNK_ROWS 行转成列；行长短不一时补空值

    strings 为真表示单元格全是字符串（CSV），跳过类型检查；已经按列读好的字符串（Arrow 读 CSV）用 add_strings
    """

    def __init__(self, strings: bool = False):
        self._strings = strings
        self._columns: list[list[pa.Array]] = []
        self._pending: list[Sequence[Any]] = []
        self._done = 0

    @property
    def height(self) -> int:
        return self._done + len(self._pending)

    def add(self, row: Sequence[Any]) -> None:
        self._pending.append(row)
        if len(self._pending) >= CHUNK_ROWS:
            self._flush()

    def add_rows(self, rows: Sequence[Sequence[Any]]) -> None:
        self._pending.extend(rows)
        if len(self._pending) >= CHUNK_ROWS:
            self._flush()

    def add_strings(self, columns: Sequence[pa.Array]) -> None:
        """一批按列存放的字符串单元格，各列等长"""
        self._flush()
        self._append([[_normalize(column)] for column in columns], len(columns[0]))

    def finish(self, sheet: str | None) -> Grid:
        self._flush()
        return Grid(sheet, [Column(pieces) for pieces in self._columns], self._done)

    def _flush(self) -> None:
        rows, count = self._pending, len(self._pending)
        if not count:
            return
        width = max(map(len, rows))
        self._append(
            [
                [
                    _normalize(pa.array(values, pa.string()))
                    if self._strings
                    else _convert(values)
                ]
                for values in (zip_longest(*rows) if width else ())
            ],
            count,
        )
        self._pending = []

    def _append(self, columns: list[list[pa.Array]], count: int) -> None:
        for position, pieces in enumerate(columns):
            if position == len(self._columns):
                # 这一列在前面的块里没有出现过
                self._columns.append([pa.nulls(self._done)])
            self._columns[position].extend(pieces)
        for pieces in self._columns[len(columns) :]:
            pieces.append(pa.nulls(count))
        self._done += count


def each(piece: pa.Array, compute, dtype) -> np.ndarray:
    """对一段计算 numpy 结果；union 段对每个子数组分别计算，再按位置拼回"""
    if not pa.types.is_union(piece.type):
        return compute(piece)
    codes, offsets = union_layout(piece)
    result = np.zeros(len(piece), dtype)
    for code in range(piece.type.num_fields):
        at = np.flatnonzero(codes == code)
        if len(at):
            result[at] = compute(piece.field(code))[offsets[at]]
    return result


def _join(parts: list[np.ndarray], dtype) -> np.ndarray:
    return np.concatenate(parts) if parts else np.zeros(0, dtype)


def _normalize(piece: pa.Array) -> pa.Array:
    """字符串去掉首尾空白（和 str.strip 相同），空串记为空值"""
    if not pa.types.is_string(piece.type):
        return piece
    piece = pc.utf8_trim(piece, WHITESPACE)
    return pc.if_else(pc.equal(piece, ""), NULL_TEXT, piece)


def _convert(values: Sequence[Any]) -> pa.Array:
    """一段单元格转成 Arrow：类型一致时转成一个数组，混杂时转成 dense union"""
    types = set(map(type, values))
    types.discard(NoneType)
    if str in types and len(types) > 1:
        # 空白单元格（Excel 读出来是 ""）不算一种类型
        values = [
            None if type(v) is str and (not v or v.isspace()) else v for v in values
        ]
        types = set(map(type, values))
        types.discard(NoneType)
    piece = _uniform(values, types)
    if piece is not None:
        return _normalize(piece)
    families: dict[type, int] = {}
    children: list[list[Any]] = []
    codes = np.empty(len(values), np.int8)
    offsets = np.empty(len(values), np.int32)
    for index, value in enumerate(values):
        family = _family(value)
        code = families.get(family)
        if code is None:
            code = families[family] = len(children)
            children.append([])
        codes[index] = code
        offsets[index] = len(children[code])
        children[code].append(value)
    return pa.UnionArray.from_dense(
        pa.array(codes),
        pa.array(offsets),
        [_normalize(_child(child)) for child in children],
    )


def _uniform(values: Sequence[Any], types: set[type]) -> pa.Array | None:
    """类型一致（或能无损地合成一种）时转成一个 Arrow 数组，否则返回空"""
    if not types:
        return pa.nulls(len(values))
    if types == {str}:
        return pa.array(values, pa.string())
    if types == {bool}:
        return pa.array(values, pa.bool_())
    if types <= {int, float}:
        if float not in types:
            try:
                return pa.array(values, pa.int64())
            except (OverflowError, pa.ArrowException):
                return None
        # 整数和小数混在一起时存成 float64；整数超出 float 能精确表示的范围就不合并
        if all(
            -EXACT_FLOAT_INT < v < EXACT_FLOAT_INT for v in values if type(v) is int
        ):
            return pa.array(values, pa.float64())
        return None
    if types <= {date, datetime}:
        if any(type(v) is datetime and v.tzinfo is not None for v in values):
            return None
        if types == {date}:
            return pa.array(values, pa.date32())
        # 日期按零点存，类型推断时零点的时间仍算日期，和 classify 一致
        return pa.array(
            [datetime.combine(v, time()) if type(v) is date else v for v in values],
            pa.timestamp("us"),
        )
    return None


def _family(value: Any) -> type:
    """union 里归为同一个子数组的类型：日期和不带时区的时间合在一起，超出 int64 的整数等存成文本，其余按 Python 类型分开"""
    kind = type(value)
    if kind is date or (kind is datetime and value.tzinfo is None):
        return date
    if kind is int and value not in INT64_RANGE:
        return object
    if kind in (NoneType, str, bool, int, float):
        return kind
    return object


def _child(values: list[Any]) -> pa.Array:
    types = set(map(type, values))
    types.discard(NoneType)
    piece = _uniform(values, types)
    if piece is not None:
        return piece
    # 超出 int64 的整数、带时区的时间、时长等：存成文本，classify 对它们的结论与文本相同
    return pa.array([None if v is None else as_text(v) for v in values], pa.string())