"""在原始网格里识别表：按空行切块，找表头，合并两级表头，去掉标题、说明和合计行

只用规则；判断不了的地方照常产出表，同时记一条警告，调用方据此提示用户
要看每一行的规则（空行切块、合计行）对整列批量计算；表头、标题、说明只看每块开头或末尾的几行，取成 Python 值判断
"""

import math
from dataclasses import dataclass, field
from typing import Any

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

from app.datafiles.base import (
    MAX_TABLE_ROWS,
    Reporter,
    brief,
    no_report,
    rows_exceeded,
    where,
)
from app.datafiles.cells import WHITESPACE_CLASS, as_text, classify, mask
from app.datafiles.grid import Column, Grid
from app.schemas.data_file import ParseWarning

# 只在每块的前几行里找表头
HEADER_SEARCH_ROWS = 10

# 首个或第二个非空单元格是这些词（去掉空白和冒号后）的行是合计行
TOTAL_LABELS = {
    "合计",
    "总计",
    "小计",
    "共计",
    "总和",
    "平均",
    "平均值",
    "total",
    "subtotal",
    "grandtotal",
    "sum",
    "average",
}
# 例如"华东小计""一季度合计"
TOTAL_SUFFIXES = ("小计", "合计", "总计")
_TOTAL_SET = pa.array(sorted(TOTAL_LABELS))
_TOTAL_STRIP = f"[{WHITESPACE_CLASS}:：]"
_TOTAL_SUFFIX = f"(?:{'|'.join(TOTAL_SUFFIXES)})$"
# 每个合计字样（含后缀）至少有其中一个字符；英文按转小写之前的大小写都列出
_TOTAL_HINT = "[计和均tTmMeE]"

# 警告里引用的单元格文本最长显示这么多字
QUOTE_CHARS = 30


@dataclass
class RawTable:
    sheet: str | None
    # 同一 sheet 切出多张表时的序号，从 1 开始；只有一张表时为空
    part: int | None
    header: list[Any]
    # 数据列，和表头一一对应
    columns: list[Column]
    warnings: list[ParseWarning] = field(default_factory=list)


@dataclass
class _Piece:
    columns: list[int]
    # 表头所在行在网格中的下标
    header_row: int
    header: list[Any]
    # 数据行在网格中的范围 [start, stop)
    data: list[tuple[int, int]]
    warnings: list[ParseWarning]


def split_tables(
    grid: Grid, report: Reporter = no_report
) -> tuple[list[RawTable], list[ParseWarning]]:
    """返回识别出的表和不属于任何表的警告（被丢弃的块）

    几步都要扫一遍所有行，每步结束报告一次进度，比例来自大表上的实测耗时
    """
    sheet = grid.sheet
    present = [column.present() for column in grid.columns]
    filled = np.logical_or.reduce(present) if present else np.zeros(grid.height, bool)
    blocks = _blocks(filled)
    used = _used_columns(present, blocks)
    firsts = grid.rows_at(np.array([start for start, _ in blocks], np.int64))
    report(0.3)
    pieces: list[_Piece] = []
    dropped: list[ParseWarning] = []
    for number, ((start, stop), columns, first) in enumerate(
        zip(blocks, used, firsts), start=1
    ):
        last = pieces[-1] if pieces else None
        # 表中间夹着空行时，后面的块以数据行开头、列也落在上一张表内，接在上一张表后面
        if (
            last
            and not _is_label_row([first[c] for c in columns])
            and set(columns) <= set(last.columns)
        ):
            last.data.append((start, stop))
            continue
        if stop - start < 2:
            dropped.append(
                ParseWarning(
                    code="block_dropped",
                    message=f"{where(sheet)}第 {start + 1} 行的{_quote(first)}不是表格，已跳过",
                )
            )
            continue
        pieces.append(_read_block(grid, start, stop, columns, sheet))
        report(0.3 + 0.25 * number / len(blocks))

    tables = []
    for number, piece in enumerate(pieces, start=1):
        table = _finish(grid, piece, sheet)
        report(0.55 + 0.45 * number / len(pieces))
        if table is None:
            dropped.append(
                ParseWarning(
                    code="block_dropped",
                    message=f"{where(sheet)}第 {piece.header_row + 1} 行开始的表只有表头没有数据，已跳过",
                )
            )
        else:
            tables.append(table)
    report(1.0)
    if len(tables) > 1:
        for part, table in enumerate(tables, start=1):
            table.part = part
    return tables, dropped


def _blocks(filled: np.ndarray) -> list[tuple[int, int]]:
    """连续的非空行组成的块 [start, stop)"""
    edges = np.flatnonzero(np.diff(np.concatenate(([0], filled, [0])).astype(np.int8)))
    return list(zip(edges[0::2].tolist(), edges[1::2].tolist()))


def _used_columns(
    present: list[np.ndarray], blocks: list[tuple[int, int]]
) -> list[list[int]]:
    """每块里有值的列；块之间只隔着空行，按块的起点分段求或即可"""
    if not blocks:
        return []
    starts = np.array([start for start, _ in blocks])
    used = np.array([np.logical_or.reduceat(column, starts) for column in present])
    return [np.flatnonzero(used[:, index]).tolist() for index in range(len(blocks))]


def _read_block(
    grid: Grid, start: int, stop: int, columns: list[int], sheet: str | None
) -> _Piece:
    # 表头最多在第 HEADER_SEARCH_ROWS 行，两级表头还要看它下面两行
    cells = grid.rows(start, min(stop, start + HEADER_SEARCH_ROWS + 2), columns)
    need = max(1, math.ceil(len(columns) / 2))
    dense = [sum(value is not None for value in row) >= need for row in cells]
    search = range(min(HEADER_SEARCH_ROWS, len(cells)))
    warnings = []

    # 表头：前几行里第一个够满、并且全是文字的行；标题行通常只有一个单元格，不够满
    header_at = next((i for i in search if dense[i] and _is_label_row(cells[i])), None)
    confident = header_at is not None
    if header_at is None:
        header_at = next((i for i in search if dense[i]), 0)
        warnings.append(
            ParseWarning(
                code="header_uncertain",
                message=f"{where(sheet)}没有找到明确的表头，用第 {start + header_at + 1} 行作表头，可能识别错误",
            )
        )
    if header_at > 0:
        skipped = "".join(
            _quote(cells[i]) for i in range(header_at) if _quote(cells[i])
        )
        warnings.append(
            ParseWarning(
                code="title_skipped",
                message=f"{where(sheet)}{_row_range(start, start + header_at - 1)}是标题，已跳过：{skipped}",
            )
        )

    header = cells[header_at]
    data_at = header_at + 1
    if (
        confident
        and header_at + 2 < stop - start
        and _is_subheader(cells[header_at], cells[header_at + 1], cells[header_at + 2])
    ):
        header = _merge_header(cells[header_at], cells[header_at + 1])
        data_at = header_at + 2
        warnings.append(
            ParseWarning(
                code="header_merged",
                message=f"{where(sheet)}{_row_range(start + header_at, start + header_at + 1)}合并为表头",
            )
        )
    return _Piece(
        columns, start + header_at, header, [(start + data_at, stop)], warnings
    )


def _finish(grid: Grid, piece: _Piece, sheet: str | None) -> RawTable | None:
    ranges = [np.arange(start, stop) for start, stop in piece.data if start < stop]
    rows = np.concatenate(ranges) if ranges else np.zeros(0, np.int64)
    warnings = piece.warnings

    # 表格下方只有一个文字单元格的行是说明，例如"注：数据来源于……"
    notes = []
    while len(rows) and len(piece.columns) >= 3:
        last = int(rows[-1])
        row = grid.rows(last, last + 1, piece.columns)[0]
        if not _is_note(row):
            break
        notes.append((last, row))
        rows = rows[:-1]
    if notes:
        notes.reverse()
        shown = "".join(_quote(row) for _, row in notes)
        warnings.append(
            ParseWarning(
                code="note_skipped",
                message=f"{where(sheet)}{_row_range(notes[0][0], notes[-1][0])}是说明，已跳过：{shown}",
            )
        )

    # 合计、小计行不是明细数据，留着会让求和重复计算
    totals = _totals(grid, piece.columns, rows)
    if totals.any():
        listed = brief([str(index + 1) for index in rows[totals].tolist()])
        warnings.append(
            ParseWarning(
                code="total_rows_removed",
                message=f"{where(sheet)}已排除 {int(totals.sum())} 行合计、小计（第 {listed} 行）",
            )
        )
        rows = rows[~totals]

    if not len(rows):
        return None
    if len(rows) > MAX_TABLE_ROWS:
        raise rows_exceeded(sheet)
    columns = [grid.columns[c].take(rows) for c in piece.columns]
    return RawTable(sheet, None, piece.header, columns, warnings)


def _totals(grid: Grid, columns: list[int], rows: np.ndarray) -> np.ndarray:
    """rows 里哪些是合计行：每行只看前两个非空单元格，其中有合计字样的文字就是"""
    if not len(rows):
        return np.zeros(0, bool)
    low, high = int(rows[0]), int(rows[-1]) + 1
    seen = np.zeros(high - low, np.int64)
    total = np.zeros(high - low, bool)
    for c in columns:
        column = grid.columns[c].slice(low, high)
        present = column.present()
        open_ = present & (seen < 2)
        if open_.any():
            total |= open_ & column.matches(_is_total_label)
        seen += present
        if (seen >= 2).all():
            break
    return total[rows - low]


def _is_total_label(strings: pa.Array) -> pa.Array:
    """去掉空白和冒号、转小写后是合计字样；合计字样都含 _TOTAL_HINT 里的字符，先用它筛掉绝大多数值"""
    found = np.zeros(len(strings), bool)
    rows = np.flatnonzero(mask(pc.match_substring_regex(strings, _TOTAL_HINT)))
    if len(rows):
        key = pc.utf8_lower(
            pc.replace_substring_regex(strings.take(rows), _TOTAL_STRIP, "")
        )
        found[rows] = mask(
            pc.or_(
                pc.is_in(key, value_set=_TOTAL_SET),
                pc.match_substring_regex(key, _TOTAL_SUFFIX),
            )
        )
    return pa.array(found)


def _is_label(value: Any) -> bool:
    return isinstance(value, str) and classify(value)[0] == "text"


def _is_label_row(row: list[Any]) -> bool:
    values = [value for value in row if value is not None]
    return bool(values) and all(_is_label(value) for value in values)


def _is_subheader(top: list[Any], sub: list[Any], first_data: list[Any]) -> bool:
    """第二行也全是文字、紧接着是带数值的数据行，并且两行里有合并单元格留下的空格"""
    return (
        _is_label_row(sub)
        and not _is_label_row(first_data)
        and (None in top or None in sub)
    )


def _merge_header(top: list[Any], sub: list[Any]) -> list[str]:
    """合并单元格只有左上角有值，第一行横向向前填充后与第二行用 - 连接，例如 销售额-一月"""
    merged, carry = [], None
    for upper, lower in zip(top, sub):
        if upper is not None:
            carry = upper
        parts = [as_text(part) for part in (carry, lower) if part is not None]
        merged.append("-".join(parts))
    return merged


def _is_note(row: list[Any]) -> bool:
    values = [value for value in row if value is not None]
    return len(values) == 1 and _is_label(values[0])


def _quote(row: list[Any]) -> str:
    """一行里第一个非空单元格，截断后加引号"""
    value = next((value for value in row if value is not None), None)
    if value is None:
        return ""
    text = as_text(value)
    if len(text) > QUOTE_CHARS:
        text = text[:QUOTE_CHARS] + "…"
    return f"「{text}」"


def _row_range(first: int, last: int) -> str:
    return f"第 {first + 1} 行" if first == last else f"第 {first + 1}–{last + 1} 行"