chatkit/backend/app/parsers/_table_common.py

"""Shared helpers for native-table parsers.

Answer-sheet PDFs (e.g. the `考生作答一覽表` format) contain a grid where each
logical student row is split by the extractor into:
  - a "header" row carrying class / seat / name (and per-row score cells)
  - one or more follow-up "data" rows carrying the letter strings

Example extracted rows:
    ['902', '13', '陳筱琳', '', '', '', '', ..., '51', None, None, None, '51', ...]
    [None,  None, None,     '==CBBBA==C', '===B======', 'B=B==BB==B', ...]

We walk the rows statefully, grouping follow-up data rows into the most recent
header, then extract letters from the combined cell set.
"""

from __future__ import annotations

import re
from dataclasses import dataclass, field
from typing import Optional

from ..answer_grid import normalize_letter

# Markers that identify the answer-key header row.
ANSWER_KEY_MARKERS = ("預設標準答案", "標準答案", "answer key")

# Chars kept when extracting letters from a cell.
_ANSWER_CELL_ALLOWED = re.compile(r"[A-Za-z=\-–—]")

# A cell counts as a "letter cell" if most of its chars survive filtering and
# it has at least 3 chars. This prevents single-char artifacts like a class
# number from being mis-read as one letter.
_LETTER_CELL_MIN_LEN = 3


@dataclass
class _Group:
    """A header row plus any trailing data rows that belong to it."""

    header: list
    data_rows: list[list] = field(default_factory=list)


def _cell_text(cell) -> str:
    if cell is None:
        return ""
    return str(cell).strip()


def _is_classlike(value: str) -> bool:
    return bool(re.fullmatch(r"\d{2,4}", value))


def _is_empty_cell(cell) -> bool:
    return _cell_text(cell) == ""


def _is_answer_key_header(row: list) -> bool:
    text = " ".join(_cell_text(c) for c in row)
    return any(marker in text for marker in ANSWER_KEY_MARKERS)


def _is_student_header(row: list) -> bool:
    """Header row has a class-like number + a non-empty name cell."""
    first_vals = [c for c in row if not _is_empty_cell(c)]
    if len(first_vals) < 2:
        return False
    # Look for class pattern in the leading cells
    head = [_cell_text(c) for c in row[:3]]
    head_nonempty = [h for h in head if h]
    if not head_nonempty:
        return False
    if not _is_classlike(head_nonempty[0]):
        return False
    # At least one later value should look like a name (non-letter-cell)
    return True


def _is_data_row(row: list) -> bool:
    """A data row has empty name/class columns and letter cells elsewhere."""
    head = row[:3]
    if any(not _is_empty_cell(c) for c in head):
        return False
    return any(_looks_like_letter_cell(c) for c in row)


def _looks_like_letter_cell(cell) -> bool:
    text = _cell_text(cell)
    if len(text) < _LETTER_CELL_MIN_LEN:
        return False
    kept = "".join(_ANSWER_CELL_ALLOWED.findall(text))
    return len(kept) >= _LETTER_CELL_MIN_LEN and len(kept) >= len(text) * 0.7


def _extract_letters_from_cell(cell) -> list[Optional[str]]:
    clean = "".join(_ANSWER_CELL_ALLOWED.findall(_cell_text(cell)))
    return [normalize_letter(ch) for ch in clean]


def _collect_letters(group: _Group) -> list[Optional[str]]:
    """Concatenate letter cells across header + data rows, in order."""
    letters: list[Optional[str]] = []
    for row in (group.header, *group.data_rows):
        for cell in row:
            if _looks_like_letter_cell(cell):
                letters.extend(_extract_letters_from_cell(cell))
    return letters


def _extract_name(row: list) -> Optional[str]:
    """Return the first non-empty, non-numeric cell in positions 0-4."""
    for cell in row[:5]:
        text = _cell_text(cell)
        if not text:
            continue
        if re.fullmatch(r"-?\d+(\.\d+)?", text):
            continue
        if _looks_like_letter_cell(text):
            continue
        return text
    return None


def _extract_student_id(row: list) -> Optional[str]:
    """Return 'class-seat' from the header row when present."""
    head = [_cell_text(c) for c in row[:3]]
    class_val = next((h for h in head if _is_classlike(h)), None)
    if not class_val:
        return None
    seat_val = next(
        (h for h in head if h and h != class_val and re.fullmatch(r"\d{1,3}", h)),
        None,
    )
    return f"{class_val}-{seat_val}" if seat_val else class_val


def _group_rows(rows: list[list]) -> tuple[Optional[_Group], list[_Group]]:
    """Walk rows and group each header with its trailing data rows.

    Returns (answer_key_group, student_groups).
    """
    answer_key: Optional[_Group] = None
    students: list[_Group] = []
    current: Optional[_Group] = None
    bucket: Optional[str] = None  # "key" or "student"

    for row in rows:
        if _is_answer_key_header(row):
            current = _Group(header=row)
            answer_key = current
            bucket = "key"
            continue
        if _is_student_header(row):
            current = _Group(header=row)
            students.append(current)
            bucket = "student"
            continue
        if current is not None and _is_data_row(row):
            current.data_rows.append(row)
            continue
        # Non-data, non-header row ends the current group.
        if bucket is not None and not _is_data_row(row):
            # Only reset on rows with visible header-like content (avoid noise rows)
            if any(not _is_empty_cell(c) for c in row[:3]):
                current = None
                bucket = None

    return answer_key, students


def table_to_student_answers(rows: list[list]) -> dict:
    _, student_groups = _group_rows(rows)
    students: list[dict] = []
    for g in student_groups:
        name = _extract_name(g.header)
        if not name:
            continue
        letters = _collect_letters(g)
        if not letters:
            continue
        students.append(
            {
                "name": name,
                "id": _extract_student_id(g.header) or "",
                "answers": [
                    {"question_number": i + 1, "answer": letter}
                    for i, letter in enumerate(letters)
                ],
            }
        )
    return {"students": students}


def table_to_teacher_answers(rows: list[list]) -> dict:
    key_group, _ = _group_rows(rows)
    if key_group is None:
        return {"answers": []}
    letters = _collect_letters(key_group)
    return {
        "answers": [
            {
                "question_number": i + 1,
                "correct_answer": letter,
                "explanation": None,
            }
            for i, letter in enumerate(letters)
            if letter is not None
        ]
    }


def table_to_data(rows: list[list], data_type: str) -> dict:
    if data_type == "student_answers":
        return table_to_student_answers(rows)
    if data_type == "teacher_answers":
        return table_to_teacher_answers(rows)
    raise ValueError(f"Native table parser does not support data_type={data_type}")