#!/usr/bin/env python3
# Copyright (c) 2026 Lark Technologies Pte. Ltd.
# SPDX-License-Identifier: MIT
"""Detect occupied subtable regions in a Lark sheet."""

from __future__ import annotations

import argparse
import csv
import io
import re
from collections import deque
from dataclasses import dataclass
from typing import Any

from lark_sheet_range import RangeBounds, col_to_index, format_range, index_to_col, parse_range, range_union
from lark_sheet_read_cli import (
    LarkCliError,
    add_spreadsheet_args,
    emit_error,
    emit_success,
    envelope_data,
    resolve_target_sheets,
    run_sheets,
    sheet_identifier,
    sheet_locator,
    sheet_title,
)

ACTION = "detect_subtables"
ROW_PREFIX_RE = re.compile(r"^\[row=(\d+)\]\s?(.*)$")
MAX_EXTERNAL_MERGE_ANCHOR_CHECKS = 10


def _inside_quoted_field(lines: list[str]) -> bool:
    """True when the accumulated record has an unterminated quoted field.

    RFC 4180 escapes a literal quote by doubling it, so both halves of a `""`
    pair count and parity still tracks whether a field is left open. A record
    that is still open must swallow the next physical line verbatim — even one
    that looks like a `[row=N]` prefix, because inside quotes that text is
    ordinary cell content, not a new record.
    """
    return sum(line.count('"') for line in lines) % 2 == 1


@dataclass
class CsvGrid:
    row_numbers: list[int]
    col_letters: list[str]
    values: list[list[str]]
    row_numbers_inferred: bool = False


@dataclass
class Component:
    bounds: RangeBounds
    occupied_count: int


def parse_annotated_csv(
    text: str,
    col_indices: list[str] | None = None,
    row_indices: list[int] | None = None,
    source_range: str | None = None,
) -> CsvGrid:
    row_numbers: list[int] = []
    values: list[list[str]] = []
    max_cols = 0
    lines = (text or "").splitlines()
    row_numbers_inferred = False
    has_authoritative_rows = isinstance(row_indices, list) and len(row_indices) > 0

    first_meaningful = next((line for line in lines if line.strip()), "")
    if ROW_PREFIX_RE.match(first_meaningful):
        records = []
        current_lines: list[str] | None = None
        current_row_number: int | None = None
        for line in lines:
            match = ROW_PREFIX_RE.match(line)
            if match and current_lines is not None and _inside_quoted_field(current_lines):
                # Prefix-looking text inside an open quoted field is content.
                current_lines.append(line)
                continue
            if match:
                if current_lines is not None and current_row_number is not None:
                    records.append("\n".join(current_lines))
                    row_numbers.append(current_row_number)
                current_row_number = int(match.group(1))
                current_lines = [match.group(2)]
            elif current_lines is not None:
                current_lines.append(line)
        if current_lines is not None and current_row_number is not None:
            records.append("\n".join(current_lines))
            row_numbers.append(current_row_number)

        # Cross-check against the row numbers the server itself reported.
        # _inside_quoted_field decides whether a "[row=N]" line starts a new
        # record or is content inside an open quoted field; when the payload's
        # quoting is malformed (a lone unescaped quote in a cell), that call
        # goes the wrong way and every following line is swallowed into the
        # previous cell — the rows simply vanish, and everything downstream
        # (data_range, last data row, column profiles) is quietly computed from
        # a short grid. row_indices is authoritative and already in hand, so
        # refuse rather than profile a grid that does not match it.
        if has_authoritative_rows and row_numbers != [int(r) for r in row_indices]:
            raise ValueError(
                "annotated_csv did not parse into the rows the server reported "
                f"(parsed {len(row_numbers)} rows {row_numbers[:5]}…, expected "
                f"{len(row_indices)} rows {list(row_indices)[:5]}…) — most likely "
                "an unbalanced quote in a cell. Re-read a narrower --range, or use "
                "+cells-get for this region instead of the CSV path."
            )

        for record in records:
            parsed = next(csv.reader([record]))
            values.append(parsed)
            max_cols = max(max_cols, len(parsed))
    else:
        reader = csv.reader(io.StringIO(text or ""))
        fallback_start = 1
        if source_range:
            fallback_start = parse_range(
                source_range,
                max_row=1_048_576,
                max_col=18_278,
            ).start_row
        for offset, row in enumerate(reader):
            row_num = None
            if has_authoritative_rows and offset < len(row_indices):
                try:
                    row_num = int(row_indices[offset])
                except (TypeError, ValueError):
                    # Non-numeric row index: leave row_num as None so the
                    # inferred fallback numbering below takes over.
                    pass
            if row_num is None:
                row_numbers_inferred = True
            row_numbers.append(row_num if row_num is not None else fallback_start + offset)
            values.append(row)
            max_cols = max(max_cols, len(row))

    if col_indices:
        col_letters = [str(col) for col in col_indices[:max_cols]]
        while len(col_letters) < max_cols:
            next_col = col_to_index(col_letters[-1]) + 1 if col_letters else len(col_letters) + 1
            col_letters.append(index_to_col(next_col))
    else:
        col_letters = [index_to_col(i) for i in range(1, max_cols + 1)]

    for row in values:
        row.extend([""] * (len(col_letters) - len(row)))
    inferred = bool(values) and not ROW_PREFIX_RE.match(first_meaningful) and (
        row_numbers_inferred or not has_authoritative_rows
    )
    return CsvGrid(
        row_numbers=row_numbers,
        col_letters=col_letters,
        values=values,
        row_numbers_inferred=inferred,
    )


def _merged_ranges(layout: dict[str, Any]) -> list[str]:
    merges = layout.get("merged_cells") or layout.get("merges") or []
    result = []
    for item in merges:
        if isinstance(item, str):
            result.append(item)
        elif isinstance(item, dict):
            value = item.get("range") or item.get("a1_range") or item.get("range_ref")
            if isinstance(value, str):
                result.append(value)
    return result


def _scan_bounds(grid: CsvGrid) -> RangeBounds | None:
    col_numbers = [col_to_index(col) for col in grid.col_letters]
    if grid.row_numbers and grid.col_letters:
        return RangeBounds(
            min(grid.row_numbers),
            min(col_numbers),
            max(grid.row_numbers),
            max(col_numbers),
        )
    return None


def _external_merge_anchors(grid: CsvGrid, layout: dict[str, Any]) -> dict[str, str]:
    scan_bounds = _scan_bounds(grid)
    if scan_bounds is None:
        return {}
    result = {}
    for merge_ref in _merged_ranges(layout):
        try:
            bounds = parse_range(merge_ref)
        except ValueError:
            continue
        intersects = not (
            bounds.end_row < scan_bounds.start_row
            or bounds.start_row > scan_bounds.end_row
            or bounds.end_col < scan_bounds.start_col
            or bounds.start_col > scan_bounds.end_col
        )
        anchor_in_scan = (
            scan_bounds.start_row <= bounds.start_row <= scan_bounds.end_row
            and scan_bounds.start_col <= bounds.start_col <= scan_bounds.end_col
        )
        if intersects and not anchor_in_scan:
            result[merge_ref] = format_range(
                bounds.start_row, bounds.start_col, bounds.start_row, bounds.start_col
            )
    return result


def _has_value(grid: CsvGrid) -> bool:
    return any(value.strip() for row in grid.values for value in row)


def build_occupancy(
    grid: CsvGrid,
    layout: dict[str, Any],
    *,
    confirmed_external_merges: set[str] | None = None,
) -> set[tuple[int, int]]:
    col_numbers = [col_to_index(col) for col in grid.col_letters]
    occupied: set[tuple[int, int]] = set()
    for row_idx, row_num in enumerate(grid.row_numbers):
        for col_idx, value in enumerate(grid.values[row_idx]):
            if value.strip():
                occupied.add((row_num, col_numbers[col_idx]))

    scan_bounds = _scan_bounds(grid)

    for merge_ref in _merged_ranges(layout):
        try:
            bounds = parse_range(merge_ref)
        except ValueError:
            continue
        if scan_bounds is None:
            continue
        anchor_in_scan = (
            scan_bounds.start_row <= bounds.start_row <= scan_bounds.end_row
            and scan_bounds.start_col <= bounds.start_col <= scan_bounds.end_col
        )
        if anchor_in_scan and (bounds.start_row, bounds.start_col) not in occupied:
            continue
        if not anchor_in_scan and merge_ref not in (confirmed_external_merges or set()):
            continue
        sr = max(bounds.start_row, scan_bounds.start_row)
        er = min(bounds.end_row, scan_bounds.end_row)
        sc = max(bounds.start_col, scan_bounds.start_col)
        ec = min(bounds.end_col, scan_bounds.end_col)
        if sr > er or sc > ec:
            continue
        for row in range(sr, er + 1):
            for col in range(sc, ec + 1):
                occupied.add((row, col))
    return occupied


def _raw_components(
    occupied: set[tuple[int, int]],
    *,
    adjacent_rows: dict[int, set[int]] | None = None,
    adjacent_cols: dict[int, set[int]] | None = None,
) -> list[Component]:
    remaining = set(occupied)
    components = []
    while remaining:
        start = remaining.pop()
        queue = deque([start])
        cells = [start]
        while queue:
            row, col = queue.popleft()
            vertical_rows = adjacent_rows.get(row, set()) if adjacent_rows else {row - 1, row + 1}
            neighbors = [(neighbor_row, col) for neighbor_row in vertical_rows]
            # Columns get the same treatment as rows: with --skip-hidden the
            # returned columns can be non-consecutive (A, C when B is hidden),
            # so col±1 would split visually adjacent data into two components.
            horizontal_cols = adjacent_cols.get(col, set()) if adjacent_cols else {col - 1, col + 1}
            neighbors.extend((row, neighbor_col) for neighbor_col in horizontal_cols)
            for neighbor in neighbors:
                if neighbor in remaining:
                    remaining.remove(neighbor)
                    queue.append(neighbor)
                    cells.append(neighbor)
        rows = [cell[0] for cell in cells]
        cols = [cell[1] for cell in cells]
        components.append(
            Component(
                RangeBounds(min(rows), min(cols), max(rows), max(cols)),
                len(cells),
            )
        )
    return components


def _box_gap_mergeable(a: RangeBounds, b: RangeBounds, gap_rows: int, gap_cols: int) -> bool:
    cols_overlap = not (a.end_col < b.start_col or b.end_col < a.start_col)
    rows_overlap = not (a.end_row < b.start_row or b.end_row < a.start_row)
    vertical_gap = max(b.start_row - a.end_row - 1, a.start_row - b.end_row - 1, 0)
    horizontal_gap = max(b.start_col - a.end_col - 1, a.start_col - b.end_col - 1, 0)
    return (cols_overlap and vertical_gap <= gap_rows) or (
        rows_overlap and horizontal_gap <= gap_cols
    )


def merge_components(
    components: list[Component], *, gap_rows: int = 1, gap_cols: int = 0
) -> list[Component]:
    merged = components[:]
    changed = True
    while changed:
        changed = False
        next_components: list[Component] = []
        used = [False] * len(merged)
        for i, comp in enumerate(merged):
            if used[i]:
                continue
            current = Component(comp.bounds, comp.occupied_count)
            used[i] = True
            for j in range(i + 1, len(merged)):
                if used[j]:
                    continue
                other = merged[j]
                if _box_gap_mergeable(current.bounds, other.bounds, gap_rows, gap_cols):
                    current = Component(
                        range_union(current.bounds, other.bounds),
                        current.occupied_count + other.occupied_count,
                    )
                    used[j] = True
                    changed = True
            next_components.append(current)
        merged = next_components
    return merged


def _row_values(grid: CsvGrid, bounds: RangeBounds, row_num: int) -> list[str]:
    if row_num not in grid.row_numbers:
        return []
    row = grid.values[grid.row_numbers.index(row_num)]
    values = []
    for col_idx, col_letter in enumerate(grid.col_letters):
        col_num = col_to_index(col_letter)
        if bounds.start_col <= col_num <= bounds.end_col:
            values.append(row[col_idx] if col_idx < len(row) else "")
    return values


def header_candidates(grid: CsvGrid, bounds: RangeBounds) -> list[int]:
    candidates = []
    for row in range(bounds.start_row, min(bounds.end_row, bounds.start_row + 4) + 1):
        values = _row_values(grid, bounds, row)
        non_empty = [value for value in values if value.strip()]
        if len(non_empty) >= max(1, min(2, bounds.col_count)):
            candidates.append(row)
    return candidates


def kind_guess(bounds: RangeBounds, density: float) -> str:
    if bounds.row_count >= 3 and bounds.col_count >= 2 and density >= 0.25:
        return "data_table"
    if bounds.row_count <= 2 or bounds.col_count <= 1:
        return "note_or_label"
    if density < 0.25:
        return "sparse_block"
    return "summary_block"


def summarize_components(grid: CsvGrid, components: list[Component], min_cells: int) -> list[dict[str, Any]]:
    result = []
    for idx, comp in enumerate(
        sorted(components, key=lambda item: (item.bounds.start_row, item.bounds.start_col)),
        start=1,
    ):
        if comp.occupied_count < min_cells:
            continue
        area = comp.bounds.row_count * comp.bounds.col_count
        density = comp.occupied_count / area if area else 0
        samples = []
        for row in range(comp.bounds.start_row, min(comp.bounds.end_row, comp.bounds.start_row + 2) + 1):
            samples.append(_row_values(grid, comp.bounds, row))
        result.append(
            {
                "id": f"T{idx}",
                "range": format_range(
                    comp.bounds.start_row,
                    comp.bounds.start_col,
                    comp.bounds.end_row,
                    comp.bounds.end_col,
                ),
                "rows": comp.bounds.row_count,
                "cols": comp.bounds.col_count,
                "occupied_cells": comp.occupied_count,
                "density": round(density, 4),
                "header_candidates": header_candidates(grid, comp.bounds),
                "kind_guess": kind_guess(comp.bounds, density),
                "sample": samples,
            }
        )
    return result


def detect_subtables(args) -> tuple[dict[str, Any], list[str]]:
    warnings: list[str] = []
    workbook = envelope_data(
        run_sheets(
            "+workbook-info",
            url=args.url,
            spreadsheet_token=args.spreadsheet_token,
            timeout=args.timeout,
        )
    )
    sheet = resolve_target_sheets(
        workbook,
        sheet_id=args.sheet_id,
        sheet_name=args.sheet_name,
        require_one=True,
    )[0]
    sid = sheet_identifier(sheet)
    title = sheet_title(sheet)
    locator = sheet_locator(sheet)

    col_count = min(int(sheet.get("column_count") or args.max_scan_cols), args.max_scan_cols)
    row_count = min(int(sheet.get("row_count") or args.max_scan_rows), args.max_scan_rows)
    scan_range = args.range or f"A1:{index_to_col(max(1, col_count))}{max(1, row_count)}"
    if not args.range:
        if int(sheet.get("column_count") or 0) > args.max_scan_cols:
            warnings.append(f"scan clipped to first {args.max_scan_cols} columns")
        if int(sheet.get("row_count") or 0) > args.max_scan_rows:
            warnings.append(f"scan clipped to first {args.max_scan_rows} rows")

    layout = envelope_data(
        run_sheets(
            "+sheet-info",
            url=args.url,
            spreadsheet_token=args.spreadsheet_token,
            **locator,
            flags={"include": "merges,hidden_rows,hidden_cols"},
            timeout=args.timeout,
        )
    )
    csv_data = envelope_data(
        run_sheets(
            "+csv-get",
            url=args.url,
            spreadsheet_token=args.spreadsheet_token,
            **locator,
            flags={
                "range": scan_range,
                "max_chars": args.max_chars,
                "skip_hidden": True if args.skip_hidden else None,
            },
            timeout=args.timeout,
        )
    )
    actual_range = str(csv_data.get("actual_range") or scan_range)
    if csv_data.get("has_more"):
        raise LarkCliError(
            f"+csv-get truncated the scan range at {actual_range}; narrow --range before detecting subtables"
        )

    grid = parse_annotated_csv(
        csv_data.get("annotated_csv", ""),
        csv_data.get("col_indices"),
        csv_data.get("row_indices"),
        actual_range,
    )
    if grid.row_numbers_inferred:
        warnings.append("CSV row numbers were inferred from the requested range")
    hidden_rows_raw = layout.get("hidden_rows") or []
    hidden_row_indexes = {
        int(value) + 1
        for value in hidden_rows_raw
        if isinstance(value, (int, str)) and str(value).isdigit()
    }
    hidden_columns_raw = layout.get("hidden_cols") or layout.get("hidden_columns") or []
    hidden_col_letters = set()
    for value in hidden_columns_raw if isinstance(hidden_columns_raw, list) else []:
        if isinstance(value, str) and value.isalpha():
            hidden_col_letters.add(value.upper())
        elif isinstance(value, (int, str)) and str(value).isdigit():
            hidden_col_letters.add(index_to_col(int(value) + 1))
    hidden_rows = sorted(row for row in grid.row_numbers if row in hidden_row_indexes)
    hidden_columns = [col for col in grid.col_letters if col.upper() in hidden_col_letters]
    if hidden_rows or hidden_columns:
        warnings.append("scan includes hidden rows or columns; pass --skip-hidden to exclude them")
    confirmed_external_merges = set()
    external_merge_anchors = list(_external_merge_anchors(grid, layout).items())
    if len(external_merge_anchors) > MAX_EXTERNAL_MERGE_ANCHOR_CHECKS:
        warnings.append(
            f"skipped confirmation for {len(external_merge_anchors) - MAX_EXTERNAL_MERGE_ANCHOR_CHECKS} "
            f"external merge anchors (limit: {MAX_EXTERNAL_MERGE_ANCHOR_CHECKS})"
        )
    for merge_ref, anchor in external_merge_anchors[:MAX_EXTERNAL_MERGE_ANCHOR_CHECKS]:
        try:
            anchor_data = envelope_data(
                run_sheets(
                    "+csv-get",
                    url=args.url,
                    spreadsheet_token=args.spreadsheet_token,
                    **locator,
                    flags={
                        "range": anchor,
                        "max_chars": 1024,
                        "skip_hidden": True if args.skip_hidden else None,
                    },
                    timeout=args.timeout,
                )
            )
            anchor_grid = parse_annotated_csv(
                anchor_data.get("annotated_csv", ""),
                anchor_data.get("col_indices"),
                anchor_data.get("row_indices"),
                anchor,
            )
            if _has_value(anchor_grid):
                confirmed_external_merges.add(merge_ref)
        except (LarkCliError, ValueError) as exc:
            warnings.append(f"could not confirm merge anchor {anchor}: {exc}")
    occupied = build_occupancy(
        grid,
        layout,
        confirmed_external_merges=confirmed_external_merges,
    )
    adjacent_rows = None
    adjacent_cols = None
    if args.skip_hidden:
        adjacent_rows = {}
        for previous, current in zip(grid.row_numbers, grid.row_numbers[1:]):
            adjacent_rows.setdefault(previous, set()).add(current)
            adjacent_rows.setdefault(current, set()).add(previous)
        col_numbers = [col_to_index(col) for col in grid.col_letters]
        adjacent_cols = {}
        for previous, current in zip(col_numbers, col_numbers[1:]):
            adjacent_cols.setdefault(previous, set()).add(current)
            adjacent_cols.setdefault(current, set()).add(previous)
    raw_components = _raw_components(
        occupied,
        adjacent_rows=adjacent_rows,
        adjacent_cols=adjacent_cols,
    )
    if len(raw_components) > args.max_merge_components:
        warnings.append(
            f"skipped gap-based component merging for {len(raw_components)} components "
            f"(limit: {args.max_merge_components})"
        )
        components = raw_components
    else:
        components = merge_components(
            raw_components,
            gap_rows=args.gap_rows,
            gap_cols=args.gap_cols,
        )
    subtables = summarize_components(grid, components, args.min_cells)
    return {
        "sheet_id": sid,
        "sheet": title,
        "scan_range": scan_range,
        "actual_range": actual_range,
        "visibility": {
            "skip_hidden": args.skip_hidden,
            "hidden_rows_in_range": hidden_rows,
            "hidden_columns_in_range": hidden_columns,
        },
        "subtables": subtables,
    }, warnings


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    add_spreadsheet_args(parser, require_sheet=True, allow_sheet=True)
    parser.add_argument("--range")
    parser.add_argument("--max-scan-rows", type=int, default=5000)
    parser.add_argument("--max-scan-cols", type=int, default=200)
    parser.add_argument("--gap-rows", type=int, default=1)
    parser.add_argument("--gap-cols", type=int, default=0)
    parser.add_argument("--min-cells", type=int, default=2)
    parser.add_argument("--max-merge-components", type=int, default=2000)
    parser.add_argument("--max-chars", type=int, default=25000)
    parser.add_argument("--skip-hidden", action="store_true")
    parser.add_argument("--timeout", type=int, default=60)
    args = parser.parse_args()
    try:
        data, warnings = detect_subtables(args)
    except (LarkCliError, ValueError, TypeError) as exc:
        emit_error(ACTION, str(exc))
    emit_success(ACTION, data, warnings)


if __name__ == "__main__":
    main()