#!/usr/bin/env python3

from __future__ import annotations

import sys
from typing import Any

from _pdf_common import (
    SkillArgumentParser,
    input_pdf,
    run_cli,
    selected_page_window,
)

DEFAULT_MAX_PAGES = 5
DEFAULT_MAX_TABLES = 20
DEFAULT_MAX_CELLS = 2000


def _parse_args(argv: list[str]):
    parser = SkillArgumentParser(description="按页提取 PDF 表格")
    parser.add_argument("--input", required=True, help="本地 PDF 路径")
    parser.add_argument("--start-page", type=int, default=1, help="起始页，1-based")
    parser.add_argument("--end-page", type=int, help="结束页，1-based，默认到末页")
    parser.add_argument(
        "--start-table",
        type=int,
        default=0,
        help="起始页内从第几个表格继续，0-based",
    )
    parser.add_argument(
        "--max-pages",
        type=int,
        default=DEFAULT_MAX_PAGES,
        help=f"单次最多处理页数，默认 {DEFAULT_MAX_PAGES}",
    )
    parser.add_argument(
        "--max-tables",
        type=int,
        default=DEFAULT_MAX_TABLES,
        help=f"单次最多返回表格数，默认 {DEFAULT_MAX_TABLES}",
    )
    parser.add_argument(
        "--max-cells",
        type=int,
        default=DEFAULT_MAX_CELLS,
        help=f"单次最多返回单元格数，默认 {DEFAULT_MAX_CELLS}",
    )
    args = parser.parse_args(argv)
    if args.start_table < 0:
        raise ValueError("start-table 不能小于 0")
    if args.max_tables < 1:
        raise ValueError("max-tables 必须大于 0")
    if args.max_cells < 1:
        raise ValueError("max-cells 必须大于 0")
    return args


def _clean_table(table) -> list[list[str]]:
    return [
        ["" if cell is None else str(cell).replace("\x00", "") for cell in row]
        for row in table
    ]


def _page_count(path) -> int:
    from pypdf import PdfReader

    with path.open("rb") as stream:
        reader = PdfReader(stream, strict=False)
        if reader.is_encrypted:
            raise ValueError("PDF 已加密，无法提取表格")
        return len(reader.pages)


def _extract(args) -> dict[str, Any]:
    import pdfplumber

    path = input_pdf(args.input)
    page_count = _page_count(path)
    result_pages: list[dict[str, Any]] = []
    table_count = 0
    cell_count = 0
    truncated = False
    next_table = 0

    with pdfplumber.open(path) as pdf:
        observed_page_count = len(pdf.pages)
        if observed_page_count != page_count:
            raise RuntimeError(
                "pdfplumber 页数与 PDF 检查结果不一致："
                f"{observed_page_count} != {page_count}；"
                "无法可靠提取表格"
            )
        start_page, actual_end, next_page = selected_page_window(
            page_count,
            args.start_page,
            args.end_page,
            args.max_pages,
        )

        for page_number in range(start_page, actual_end + 1):
            page_tables = []
            all_tables = pdf.pages[page_number - 1].extract_tables() or []
            table_offset = args.start_table if page_number == start_page else 0
            if table_offset > len(all_tables):
                raise ValueError(
                    f"start-table 超过第 {page_number} 页表格数量 {len(all_tables)}"
                )
            for table_index, table in enumerate(
                all_tables[table_offset:],
                start=table_offset,
            ):
                cleaned = _clean_table(table)
                cells = sum(len(row) for row in cleaned)
                if cells > args.max_cells and table_count == 0:
                    raise ValueError(
                        f"第 {page_number} 页第 {table_index} 个表格有 {cells} "
                        "个单元格，请提高 max-cells"
                    )
                if (
                    table_count >= args.max_tables
                    or cell_count + cells > args.max_cells
                ):
                    truncated = True
                    next_page = page_number
                    next_table = table_index
                    break
                page_tables.append(cleaned)
                table_count += 1
                cell_count += cells

            result_pages.append({"page": page_number, "tables": page_tables})
            if truncated:
                break

    return {
        "path": str(path),
        "page_count": page_count,
        "start_page": start_page,
        "end_page": result_pages[-1]["page"],
        "window_end_page": actual_end,
        "table_count": table_count,
        "cell_count": cell_count,
        "truncated": truncated,
        "pages": result_pages,
        "has_more": next_page is not None,
        "next_page": next_page,
        "next_table": next_table,
    }


def main(argv: list[str] | None = None) -> int:
    arguments = sys.argv[1:] if argv is None else argv
    return run_cli(lambda: _extract(_parse_args(arguments)))


if __name__ == "__main__":
    raise SystemExit(main())