"""把各种格式读成原始网格（CSV、Excel）或按列的记录（JSON、JSONL），Parquet 直接读成 Arrow 表

CSV 和 Excel 可能有标题行、多行表头、合计行，所以先读成不带表头的网格，再由 tables.py 识别表区域
文本格式边读边转成列（grid.py），不把整个文件解码进内存
"""

import codecs
import csv
import io
import json
from dataclasses import dataclass, field
from itertools import islice
from pathlib import Path
from typing import Any

import pyarrow as pa
import pyarrow.csv as pacsv
import pyarrow.parquet as pq
from python_calamine import (
    CalamineWorkbook,
    PasswordError,
    SheetTypeEnum,
    SheetVisibleEnum,
)

from app.datafiles.base import (
    MAX_TABLE_ROWS,
    DataFileError,
    Reporter,
    brief,
    no_report,
    rows_exceeded,
    scaled,
)
from app.datafiles.cells import as_text
from app.datafiles.grid import CHUNK_ROWS, Column, Grid, GridBuilder
from app.schemas.data_file import ParseWarning

# 原始网格比数据多出的标题、表头、空行，读取时先粗略拦住明显超限的文件，精确的行数在拆表后检查
ROW_SLACK = 1000

# JSON 嵌套对象展开的层数，更深的部分存成 JSON 字符串
JSON_FLATTEN_LEVELS = 2

# JSONL 坏行超过这个比例时整个文件判为失败
MAX_BAD_LINE_RATIO = 0.01

# 逐行读取时每隔这么多行报告一次进度
PROGRESS_ROWS = 20_000

# Arrow 读 CSV 每次处理的字节数；单行超过它时 Arrow 报错，改用标准库
ARROW_BLOCK_BYTES = 4 * 1024 * 1024

# 单元格最长的字符数；csv 模块默认 128KB，长文本单元格会报错
csv.field_size_limit(16 * 1024 * 1024)


@dataclass
class Records:
    """已有列结构的数据，按列存放"""

    names: list[str]
    columns: list[Column]
    warnings: list[ParseWarning] = field(default_factory=list)


def read_delimited(
    path: Path, codec: str, delimiter: str, report: Reporter = no_report
) -> Grid:
    """codec 是 detect 确认过的编码，整个文件都能按它解码

    先用 Arrow 读（C++ 实现，快一个数量级）；它要求每行列数相同，遇到标题行、说明行这类长短不齐的行，
    或引号不成对、单行过长，就报错。这时改用标准库 csv 从头读一遍。Arrow 读得出来时两者结果相同，测试里对比
    """
    report = _increasing(report)
    try:
        return _read_delimited_arrow(path, codec, delimiter, report)
    except pa.ArrowInvalid:
        return _read_delimited_python(path, codec, delimiter, report)


def _read_delimited_arrow(
    path: Path, codec: str, delimiter: str, report: Reporter
) -> Grid:
    with path.open("rb") as raw, io.TextIOWrapper(raw, codec, newline="") as text:
        first = next(csv.reader(text, delimiter=delimiter), [])
    if not first:
        raise pa.ArrowInvalid("第一行是空行，列数要从后面的行推断")
    names = [f"c{i}" for i in range(len(first))]
    builder = GridBuilder(strings=True)
    size = max(path.stat().st_size, 1)
    with path.open("rb") as raw:
        if codec == "utf-8-sig" and raw.read(3) != codecs.BOM_UTF8:
            raw.seek(0)
        reader = pacsv.open_csv(
            raw,
            read_options=pacsv.ReadOptions(
                column_names=names,
                block_size=ARROW_BLOCK_BYTES,
                encoding="utf8" if codec.startswith("utf-8") else codec,
            ),
            parse_options=pacsv.ParseOptions(
                delimiter=delimiter, newlines_in_values=True, ignore_empty_lines=False
            ),
            convert_options=pacsv.ConvertOptions(
                column_types={name: pa.string() for name in names},
                strings_can_be_null=False,
                quoted_strings_can_be_null=False,
            ),
        )
        batches: list[pa.RecordBatch] = []
        waiting = 0
        for batch in reader:
            batches.append(batch)
            waiting += len(batch)
            # 攒到一个块再转：每段太短时，类型推断的固定开销会变大
            if waiting >= CHUNK_ROWS:
                _add_batches(builder, batches)
                batches, waiting = [], 0
            if builder.height > MAX_TABLE_ROWS + ROW_SLACK:
                raise rows_exceeded(None)
            report(raw.tell() / size)
        if batches:
            _add_batches(builder, batches)
    if builder.height > MAX_TABLE_ROWS + ROW_SLACK:
        raise rows_exceeded(None)
    return builder.finish(None)


def _add_batches(builder: GridBuilder, batches: list[pa.RecordBatch]) -> None:
    builder.add_strings(
        [
            pa.concat_arrays([batch.column(i) for batch in batches])
            for i in range(batches[0].num_columns)
        ]
    )


def _read_delimited_python(
    path: Path, codec: str, delimiter: str, report: Reporter
) -> Grid:
    builder = GridBuilder(strings=True)
    size = max(path.stat().st_size, 1)
    with path.open("rb") as raw, io.TextIOWrapper(raw, codec, newline="") as text:
        reader = csv.reader(text, delimiter=delimiter)
        while rows := list(islice(reader, PROGRESS_ROWS)):
            builder.add_rows(rows)
            if builder.height > MAX_TABLE_ROWS + ROW_SLACK:
                raise rows_exceeded(None)
            report(raw.tell() / size)
    return builder.finish(None)


def _increasing(report: Reporter) -> Reporter:
    """Arrow 读到一半失败、换标准库从头读时，进度不往回走"""
    reached = 0.0

    def inner(fraction: float) -> None:
        nonlocal reached
        if fraction > reached:
            reached = fraction
            report(fraction)

    return inner


def read_excel(
    path: Path, report: Reporter = no_report
) -> tuple[list[Grid], list[ParseWarning]]:
    try:
        book = CalamineWorkbook.from_path(str(path))
    except PasswordError:
        raise DataFileError(422, "Excel 文件设置了密码，请去掉密码后重新上传") from None
    except Exception:
        raise DataFileError(422, "无法读取 Excel 文件，文件可能已损坏") from None
    grids, warnings = [], []
    try:
        sheets = book.sheets_metadata
        for number, meta in enumerate(sheets):
            report(number / len(sheets))
            # 图表页、宏表等没有单元格数据
            if meta.typ != SheetTypeEnum.WorkSheet:
                continue
            if meta.visible != SheetVisibleEnum.Visible:
                warnings.append(
                    ParseWarning(
                        code="hidden_sheet_skipped",
                        message=f"跳过了隐藏的 sheet「{meta.name}」",
                    )
                )
                continue
            sheet = book.get_sheet_by_name(meta.name)
            if sheet.total_height > MAX_TABLE_ROWS + ROW_SLACK:
                raise rows_exceeded(meta.name)
            # 解析 sheet 的 XML 占读取的大部分时间，逐行转成列占剩下的
            rows = scaled(
                report, (number + 0.8) / len(sheets), (number + 1) / len(sheets)
            )
            # iter_rows 保留顶部的空行，行号和 Excel 里看到的一致；左侧的空列不影响识别表区域；
            # 没有单元格的 sheet 上 iter_rows 会崩溃，不调用
            builder = GridBuilder()
            cells = sheet.iter_rows() if sheet.height else iter(())
            while batch := list(islice(cells, PROGRESS_ROWS)):
                builder.add_rows(batch)
                rows(builder.height / sheet.height)
            grid = builder.finish(meta.name)
            if not any(column.present().any() for column in grid.columns):
                warnings.append(
                    ParseWarning(
                        code="empty_sheet_skipped",
                        message=f"sheet「{meta.name}」是空的，已跳过",
                    )
                )
                continue
            grids.append(grid)
    finally:
        book.close()
    return grids, warnings


def read_json(value: Any, report: Reporter = no_report) -> Grid | Records:
    """对象数组按记录读取；数组的数组当作网格，第一行通常是表头"""
    warnings = []
    if isinstance(value, dict):
        records, path = _find_records(value)
        if path:
            warnings.append(
                ParseWarning(
                    code="json_path_selected",
                    message=f"使用了 JSON 中「{path}」下的 {len(records)} 条记录",
                )
            )
    else:
        records = value
    if not records:
        raise DataFileError(422, "JSON 里没有找到数据")
    if all(isinstance(item, list) for item in records):
        if len(records) > MAX_TABLE_ROWS + ROW_SLACK:
            raise rows_exceeded(None)
        return Grid.from_rows(
            None, ([_scalar(cell) for cell in row] for row in records)
        )
    if all(isinstance(item, dict) for item in records):
        flattener = _Flattener()
        for number, record in enumerate(records):
            if number % PROGRESS_ROWS == 0:
                report(number / len(records))
            flattener.add(record)
        result = flattener.finish()
        result.warnings = warnings
        return result
    if not any(isinstance(item, (dict, list)) for item in records):
        if len(records) > MAX_TABLE_ROWS:
            raise rows_exceeded(None)
        return Records(["值"], [Column.from_values(records)], warnings)
    raise DataFileError(422, "JSON 数组里的元素类型不一致，需要都是对象或都是数组")


def read_jsonl(path: Path, codec: str, report: Reporter = no_report) -> Records:
    """逐行读取；codec 是 detect 确认过的编码"""
    flattener, bad, total = _Flattener(), [], 0
    size = max(path.stat().st_size, 1)
    with path.open("rb") as raw, io.TextIOWrapper(raw, codec) as text:
        for number, line in enumerate(text, start=1):
            if number % PROGRESS_ROWS == 0:
                report(raw.tell() / size)
            if not line.strip():
                continue
            total += 1
            try:
                record = json.loads(line)
            except ValueError:
                record = None
            if isinstance(record, dict):
                flattener.add(record)
            else:
                bad.append(number)
    if not flattener.count:
        raise DataFileError(422, "JSONL 里没有找到数据")
    if len(bad) > total * MAX_BAD_LINE_RATIO:
        raise DataFileError(
            422, f"JSONL 有 {len(bad)} 行无法解析（例如第 {bad[0]} 行），超过 1%"
        )
    result = flattener.finish()
    if bad:
        result.warnings.append(
            ParseWarning(
                code="bad_lines_skipped",
                message=f"跳过了 {len(bad)} 行无法解析的内容：{brief([f'第 {n} 行' for n in bad])}",
            )
        )
    return result


def read_parquet(path: Path) -> pa.Table:
    try:
        table = pq.read_table(path)
    except Exception:
        raise DataFileError(422, "无法读取 Parquet 文件，文件可能已损坏") from None
    if table.num_rows > MAX_TABLE_ROWS:
        raise rows_exceeded(None)
    # struct 展开一层，列名写成 a.b；更深的嵌套由 columns.py 存成 JSON 字符串
    return table.flatten()


def _find_records(value: dict) -> tuple[list, str | None]:
    """在对象的前两层里找最长的对象数组，例如 {"data": [...]}；找不到时按列式或单条记录处理"""
    candidates = []
    for key, item in value.items():
        if _is_record_list(item):
            candidates.append((len(item), str(key), item))
        elif isinstance(item, dict):
            for sub_key, sub in item.items():
                if _is_record_list(sub):
                    candidates.append((len(sub), f"{key}.{sub_key}", sub))
    if candidates:
        _, path, records = max(candidates, key=lambda candidate: candidate[0])
        return records, path
    lists = list(value.values())
    # 列式：{"地区": [...], "销售额": [...]}
    if (
        lists
        and all(isinstance(item, list) for item in lists)
        and len({len(item) for item in lists}) == 1
    ):
        return [dict(zip(value, row)) for row in zip(*lists)], None
    return [value], None


def _is_record_list(value: Any) -> bool:
    return (
        isinstance(value, list)
        and bool(value)
        and all(isinstance(item, dict) for item in value)
    )


class _Flattener:
    """逐条展开记录，按列存放：列按首次出现的顺序排列，记录里没有的字段为空"""

    def __init__(self) -> None:
        self._index: dict[str, int] = {}
        self._builder = GridBuilder()

    @property
    def count(self) -> int:
        return self._builder.height

    def add(self, record: dict) -> None:
        if self._builder.height >= MAX_TABLE_ROWS:
            raise rows_exceeded(None)
        flat: dict[str, Any] = {}
        _flatten(record, "", 0, flat)
        row: list[Any] = [None] * len(self._index)
        for name, value in flat.items():
            position = self._index.setdefault(name, len(self._index))
            if position < len(row):
                row[position] = value
            else:
                row.append(value)
        self._builder.add(row)

    def finish(self) -> Records:
        grid = self._builder.finish(None)
        return Records(list(self._index), grid.columns)


def _flatten(record: dict, prefix: str, level: int, out: dict[str, Any]) -> None:
    for key, value in record.items():
        name = f"{prefix}.{key}" if prefix else str(key)
        if isinstance(value, dict) and value and level < JSON_FLATTEN_LEVELS:
            _flatten(value, name, level + 1, out)
        else:
            out[name] = _scalar(value)


def _scalar(value: Any) -> Any:
    # 嵌套的对象、数组存成 JSON 字符串，其余原样交给类型推断
    return as_text(value) if isinstance(value, (dict, list)) else value