chatkit/backend/seed_from_table.py

#!/usr/bin/env python3
"""Seed questions from Google Sheets HTML table exports into the question_bank.

Supports three column layouts detected automatically:

  Format A — year-based, positional (103-114 會考單題.html):
    Col 0: year  Col 2: stem  Col 3: choices  Col 4: answer letter
    → GPT classifies main_category + tags

  Format B — pre-labelled, positional (總表.html):
    Col 2: main_category  Col 4: stem  Col 5: choices  no answer

  Format C — labelled by header text, any column order (e.g. 年度/題號/
  段落文章/題目/選項/答案/詳解/本題關鍵/題型分類 sheets): columns are found by
  matching header cells like "題型分類"/"題目"/"選項"/"答案"/"段落文章" instead
  of a fixed position, so exports with extra columns (explanations, key
  points, etc.) still parse correctly.

  For B and C: category taken directly from the sheet (no GPT call for
  main_category). Rows with no valid category in the sheet fall back to full
  GPT classification, same as Format A. Rows whose (sheet-given or
  GPT-assigned) category has sub-tags in the taxonomy still get a GPT call
  for tags only.

Usage:
    cd chatkit/backend
    python seed_from_table.py file1.html [file2.html ...]
"""

from __future__ import annotations

import asyncio
import html as html_lib
import json
import re
import sys
from html.parser import HTMLParser
from pathlib import Path

sys.path.insert(0, str(Path(__file__).parent))

from openai import AsyncOpenAI

from app.config import get_settings
from app.database import save_question_bank_batch
from app.question_rewriter import PASSAGE_CATEGORIES
from app.taxonomy import ALL_TAGS, MAIN_CATEGORIES, TAXONOMY, taxonomy_prompt_block

BATCH_SIZE = 25
_VALID_CATS = set(MAIN_CATEGORIES)

# Questions that reference an image or an external passage that isn't inline
_SKIP_RE = re.compile(
    r"look at the picture"
    r"|in the picture"
    r"|the picture shows"
    r"|the following picture"
    r"|according to the (?:passage|article|reading|text|graph|chart|table)"
    r"|based on the (?:passage|article|reading|text)"
    r"|the (?:passage|article|reading) (?:above|below|following)"
    r"|the following (?:passage|article|reading)",
    re.I,
)


def _should_skip(question_text: str) -> bool:
    """Return True for questions referencing a non-existent image or passage."""
    return bool(_SKIP_RE.search(question_text))

CLASSIFY_SYSTEM = f"""You are ClassLens, an assistant for Taiwanese 國中 English exam analysis.

Classify each question into the pre-defined taxonomy. Return ONLY a JSON array:
[{{"idx": 0, "main_category": "語法結構", "tags": ["時態", "過去簡單式"]}}, ...]

Rules:
- main_category should be one of the 6 categories below, or "" if no category clearly fits.
- tags must come directly from that category's tag list. Use [] if none fit or if main_category is "".
- Classify by the PRIMARY grammar or vocabulary skill the question tests.
- Do NOT force a category if the question doesn't clearly belong to one — prefer "" over a weak match.

{taxonomy_prompt_block()}"""


# ---------------------------------------------------------------------------
# HTML table parser
# ---------------------------------------------------------------------------

class _TableParser(HTMLParser):
    def __init__(self) -> None:
        super().__init__()
        self.rows: list[list[str]] = []
        self._row: list[str] = []
        self._cell: list[str] = []
        self._in_cell = False
        self._in_thead = False
        self._skip_row = False
        self._skip_cell = False

    def handle_starttag(self, tag: str, attrs: list) -> None:
        if tag == "thead":
            self._in_thead = True
        elif tag == "tr":
            self._row = []
            # Google Sheets wraps its own spreadsheet column-letter row
            # (A, B, C, ...) in <thead> — that's never real data or the
            # sheet's own header row (which is the first row of <tbody>), so
            # it must be dropped the same way a freezebar row is, or
            # everything downstream mistakes it for the header.
            self._skip_row = self._in_thead or any(
                k == "class" and v and "freezebar" in v for k, v in attrs
            )
        elif tag in ("td", "th"):
            cls = dict(attrs).get("class", "") or ""
            # Skip row-number header cells (the <th> on the left of each row)
            if "row-headers-background" in cls:
                self._skip_cell = True
                self._in_cell = False
            else:
                self._skip_cell = False
                self._cell = []
                self._in_cell = True

    def handle_endtag(self, tag: str) -> None:
        if tag == "thead":
            self._in_thead = False
        elif tag in ("td", "th"):
            if not self._skip_cell:
                self._row.append("".join(self._cell).strip())
            self._in_cell = False
            self._skip_cell = False
        elif tag == "tr" and not self._skip_row and self._row:
            self.rows.append(self._row)

    def handle_data(self, data: str) -> None:
        if self._in_cell:
            self._cell.append(data)

    def handle_entityref(self, name: str) -> None:
        if self._in_cell:
            self._cell.append(html_lib.unescape(f"&{name};"))

    def handle_charref(self, name: str) -> None:
        if self._in_cell:
            self._cell.append(html_lib.unescape(f"&#{name};"))


def _fix_mojibake(s: str) -> str:
    """Fix UTF-8 bytes that were decoded as Latin-1 (common in Google Sheets exports)."""
    try:
        return s.encode("latin-1").decode("utf-8")
    except (UnicodeEncodeError, UnicodeDecodeError):
        return s


def parse_html_bytes(raw_bytes: bytes) -> list[list[str]]:
    """Parse a Google Sheets HTML export's raw bytes into rows, fixing
    mojibake per cell. Shared by the CLI (_parse_html reads a file first)
    and the in-app upload endpoint (already has the bytes from the upload) —
    a single implementation so both paths always parse identically."""
    raw = raw_bytes.decode("utf-8", errors="replace")
    p = _TableParser()
    p.feed(raw)
    return [[_fix_mojibake(cell) for cell in row] for row in p.rows]


def _parse_html(html_path: Path) -> list[list[str]]:
    return parse_html_bytes(html_path.read_bytes())


# ---------------------------------------------------------------------------
# Format detection + extraction
# ---------------------------------------------------------------------------

def _detect_format(rows: list[list[str]]) -> str:
    """Return 'C' (header-labeled), 'B' (pre-labelled positional), 'A'
    (year-based positional), or 'unknown'."""
    if rows and _detect_header_columns(rows[0]):
        return "C"
    for row in rows[1:6]:  # skip header, sample a few data rows
        if len(row) < 3:
            continue
        col2 = row[2].strip()
        if col2 in _VALID_CATS:
            return "B"
        if row[0].strip().isdigit():
            return "A"
    return "unknown"


def _extract_format_a(rows: list[list[str]]) -> list[dict]:
    """Year-based format: year in col 0, stem col 2, choices col 3, answer col 4."""
    questions: list[dict] = []
    for row_idx, row in enumerate(rows):
        if row_idx == 0 or len(row) < 5:
            continue
        if not row[0].strip().isdigit():
            continue

        stem = row[2].strip()
        opts = row[3].strip()
        ans  = row[4].strip().upper()

        if not stem or len(stem) < 5:
            continue

        question_text = f"{stem} {opts}".strip() if opts else stem
        if _should_skip(question_text):
            continue
        answer = ans[0] if ans and ans[0] in "ABCD" else ""

        questions.append({"question_text": question_text, "answer": answer})
    return questions


def _extract_format_b(rows: list[list[str]]) -> list[dict]:
    """Pre-labelled format: category col 2, stem col 4, choices col 5 (6-col)
    or stem col 3, choices col 4 (5-col single-category sheets).

    A row with no valid category in col 2 is NOT skipped as long as it has a
    real question stem — it's kept with main_category="" so the caller can
    send just those to GPT for full classification, same as Format A."""
    questions: list[dict] = []
    for row_idx, row in enumerate(rows):
        if row_idx == 0:
            continue

        cat = row[2].strip() if len(row) > 2 else ""
        if cat not in _VALID_CATS:
            cat = ""

        # 6-col (總表 layout): col4=stem, col5=choices
        # 5-col (single-category sheets): col3=stem, col4=choices
        if len(row) >= 6:
            stem = row[4].strip()
            opts = row[5].strip()
        else:
            stem = row[3].strip() if len(row) > 3 else ""
            opts = row[4].strip() if len(row) > 4 else ""

        # Skip placeholder rows and image/passage references
        if not stem or len(stem) < 5 or "未提供" in stem:
            continue

        question_text = f"{stem} {opts}".strip() if opts else stem
        if _should_skip(question_text):
            continue

        questions.append({
            "question_text": question_text,
            "answer": "",    # no answer column in this format
            "main_category": cat,  # "" if the sheet didn't give a valid one
            "tags": [],
        })
    return questions


_HEADER_ALIASES: dict[str, tuple[str, ...]] = {
    "category": ("題型分類", "分類", "類別"),
    "stem": ("題目",),
    "options": ("選項",),
    "answer": ("答案",),
    "passage": ("段落文章", "文章", "篇章"),
}


def _detect_header_columns(header_row: list[str]) -> dict[str, int] | None:
    """Map semantic field names to column indexes by matching header text.

    Returns None unless both a category and a stem column are found — those
    are the two fields Format A/B can't locate by fixed position for a sheet
    whose columns aren't in the 5/6-col layouts they assume."""
    col_map: dict[str, int] = {}
    for idx, cell in enumerate(header_row):
        cell = cell.strip()
        for field, aliases in _HEADER_ALIASES.items():
            if field not in col_map and cell in aliases:
                col_map[field] = idx
    if "category" in col_map and "stem" in col_map:
        return col_map
    return None


def _extract_labeled(rows: list[list[str]], col_map: dict[str, int]) -> list[dict]:
    """Extract questions from a sheet whose columns were identified by header
    text (see _detect_header_columns) rather than fixed position. Category is
    taken directly from the sheet when valid, same as Format B."""
    stem_col = col_map["stem"]
    cat_col = col_map["category"]
    opts_col = col_map.get("options")
    ans_col = col_map.get("answer")
    passage_col = col_map.get("passage")

    questions: list[dict] = []
    for row_idx, row in enumerate(rows):
        if row_idx == 0:
            continue

        stem = row[stem_col].strip() if len(row) > stem_col else ""
        if not stem or len(stem) < 5 or "未提供" in stem:
            continue

        opts = row[opts_col].strip() if opts_col is not None and len(row) > opts_col else ""
        passage = row[passage_col].strip() if passage_col is not None and len(row) > passage_col else ""
        question_text = " ".join(p for p in (passage, stem, opts) if p)
        if _should_skip(question_text):
            continue

        ans = row[ans_col].strip().upper() if ans_col is not None and len(row) > ans_col else ""
        answer = ans[0] if ans and ans[0] in "ABCD" else ""

        cat = row[cat_col].strip() if len(row) > cat_col else ""
        if cat not in _VALID_CATS:
            cat = ""
        has_passage = bool(passage)
        if has_passage and cat not in PASSAGE_CATEGORIES:
            # A question with its own passage can only be one of the
            # passage-type categories (文意推論/篇章大意/篇章細節/篇章結構) —
            # a sheet-given category outside that set (e.g. 字詞理解/語法結構,
            # or none at all) contradicts the row's own structure, so
            # re-classify instead of trusting it.
            cat = ""

        questions.append({
            "question_text": question_text,
            "answer": answer,
            "main_category": cat,
            "tags": [],
            "_needs_passage_category": has_passage and not cat,
        })
    return questions


# ---------------------------------------------------------------------------
# GPT classification (Format A only)
# ---------------------------------------------------------------------------

async def _classify_batch(
    client: AsyncOpenAI,
    batch: list[dict],
    model: str,
) -> list[dict]:
    items = "\n".join(
        f'{i}. {q["question_text"][:350]}'
        for i, q in enumerate(batch)
    )
    resp = await client.chat.completions.create(
        model=model,
        messages=[
            {"role": "system", "content": CLASSIFY_SYSTEM},
            {"role": "user", "content": f"Classify these {len(batch)} questions:\n\n{items}"},
        ],
        temperature=0.0,
        max_tokens=2048,
    )
    content = resp.choices[0].message.content or ""
    content = re.sub(r"^```json?\s*|\s*```$", "", content.strip())
    results: list[dict] = json.loads(content)

    class_map: dict[int, dict] = {}
    for r in results:
        idx = r.get("idx")
        if idx is None:
            continue
        cat = r.get("main_category", "")
        if cat not in _VALID_CATS:
            cat = ""
        tags = r.get("tags", [])
        if not isinstance(tags, list):
            tags = []
        class_map[idx] = {"main_category": cat, "tags": tags}

    classified = []
    for i, q in enumerate(batch):
        cls = class_map.get(i, {})
        classified.append({
            "question_text": q["question_text"],
            "answer": q["answer"],
            "main_category": cls.get("main_category", ""),
            "tags": cls.get("tags", []),
        })
    return classified


_PASSAGE_CLASSIFY_SYSTEM = f"""You are ClassLens, an assistant for Taiwanese 國中 English exam analysis.

Every question below is a reading-passage question — its stem already includes (or is
paired with) a passage, so it can ONLY belong to one of these 4 categories:
{', '.join(sorted(PASSAGE_CATEGORIES))}

Return ONLY a JSON array: [{{"idx": 0, "main_category": "篇章細節"}}, ...]

Rules:
- main_category MUST be one of the 4 categories above — never a vocabulary/grammar
  category, and never "" (every passage question fits one of these 4).
- Classify by what the question is actually testing: 文意推論 (inference), 篇章大意
  (main idea), 篇章細節 (detail), 篇章結構 (structure/reference/connector)."""


async def _classify_passage_category_batch(
    client: AsyncOpenAI,
    batch: list[dict],
    model: str,
) -> list[dict]:
    """Classify main_category only, constrained to the 4 passage-type
    categories — for rows already known (from the sheet's passage column) to
    depend on a passage, so a full 6-category classification could otherwise
    misfile them under 字詞理解/語法結構. No tags call needed: none of the
    passage categories have sub-tags in the taxonomy."""
    items = "\n".join(
        f'{i}. {q["question_text"][:350]}'
        for i, q in enumerate(batch)
    )
    resp = await client.chat.completions.create(
        model=model,
        messages=[
            {"role": "system", "content": _PASSAGE_CLASSIFY_SYSTEM},
            {"role": "user", "content": f"Classify these {len(batch)} questions:\n\n{items}"},
        ],
        temperature=0.0,
        max_tokens=1024,
    )
    content = resp.choices[0].message.content or ""
    content = re.sub(r"^```json?\s*|\s*```$", "", content.strip())
    results: list[dict] = json.loads(content)

    class_map: dict[int, str] = {}
    for r in results:
        idx = r.get("idx")
        if idx is None:
            continue
        cat = r.get("main_category", "")
        if cat in PASSAGE_CATEGORIES:
            class_map[idx] = cat

    classified = []
    for i, q in enumerate(batch):
        classified.append({
            "question_text": q["question_text"],
            "answer": q["answer"],
            "main_category": class_map.get(i, ""),
            "tags": [],
        })
    return classified


async def _classify_tags_batch(
    client: AsyncOpenAI,
    batch: list[dict],
    model: str,
    category: str,
) -> dict[int, list[str]]:
    """Assign sub-tags only, for questions whose main_category is already
    known (from the sheet or from a prior full classification pass)."""
    valid_tags = set(ALL_TAGS.get(category, []))
    if not valid_tags:
        return {}

    tag_list = ", ".join(sorted(valid_tags))
    system = f"""You are ClassLens, an assistant for Taiwanese 國中 English exam analysis.

Every question below is already classified as "{category}". Assign each question the
best-fitting tag(s) from this list only: {tag_list}

Return ONLY a JSON array: [{{"idx": 0, "tags": ["過去簡單式"]}}, ...]

Rules:
- tags must come directly from the list above. Use [] if none clearly fit — never invent a tag.
- Prefer the single best-fit tag; only include more than one if genuinely both apply."""

    items = "\n".join(
        f'{i}. {q["question_text"][:350]}'
        for i, q in enumerate(batch)
    )
    resp = await client.chat.completions.create(
        model=model,
        messages=[
            {"role": "system", "content": system},
            {"role": "user", "content": f"Assign tags to these {len(batch)} questions:\n\n{items}"},
        ],
        temperature=0.0,
        max_tokens=1024,
    )
    content = resp.choices[0].message.content or ""
    content = re.sub(r"^```json?\s*|\s*```$", "", content.strip())
    results: list[dict] = json.loads(content)

    tag_map: dict[int, list[str]] = {}
    for r in results:
        idx = r.get("idx")
        if idx is None:
            continue
        tags = r.get("tags", [])
        if not isinstance(tags, list):
            tags = []
        tag_map[idx] = [t for t in tags if t in valid_tags]
    return tag_map


async def _resolve_categories_and_tags(
    raw_entries: list[dict],
    model: str,
    client: AsyncOpenAI | None,
    settings,
) -> tuple[list[dict], AsyncOpenAI | None]:
    """Fill in categories/tags for entries extracted with main_category
    possibly already set from the sheet (Format B/C):
    - Entries with no valid category ("") get a full GPT classification pass.
    - Entries whose sheet-given category has sub-tags in the taxonomy get a
      cheaper, tags-only GPT pass — their main_category is never re-guessed.
    Returns (all_entries, client) so the caller can reuse the lazily-created client.
    """
    labelled = [e for e in raw_entries if e["main_category"]]
    unlabelled = [e for e in raw_entries if not e["main_category"]]
    print(
        f"  Parsed {len(raw_entries)} questions "
        f"({len(labelled)} pre-labelled, {len(unlabelled)} with no valid category in the sheet)."
    )
    all_entries = list(labelled)

    def ensure_client() -> AsyncOpenAI:
        nonlocal client
        if client is None:
            if not settings.openai_api_key:
                raise RuntimeError("OPENAI_API_KEY is required to classify these questions.")
            client = AsyncOpenAI(api_key=settings.openai_api_key)
        return client

    # Rows already known (from the sheet's passage column) to depend on a
    # passage are classified separately, constrained to just the 4
    # passage-type categories — a full 6-category pass could otherwise
    # misfile a passage question under 字詞理解/語法結構.
    unlabelled_passage = [e for e in unlabelled if e.get("_needs_passage_category")]
    unlabelled_other = [e for e in unlabelled if not e.get("_needs_passage_category")]

    if unlabelled_passage:
        print(f"  Classifying {len(unlabelled_passage)} passage question(s) via GPT (passage categories only)...")
        c = ensure_client()
        for start in range(0, len(unlabelled_passage), BATCH_SIZE):
            batch = unlabelled_passage[start : start + BATCH_SIZE]
            try:
                all_entries.extend(await _classify_passage_category_batch(c, batch, model))
            except Exception as exc:
                print(f"  ERROR classifying passage batch: {exc}", file=sys.stderr)
                all_entries.extend({**q, "main_category": "", "tags": []} for q in batch)

    if unlabelled_other:
        print(f"  Classifying {len(unlabelled_other)} unlabelled question(s) via GPT...")
        c = ensure_client()
        for start in range(0, len(unlabelled_other), BATCH_SIZE):
            batch = unlabelled_other[start : start + BATCH_SIZE]
            try:
                all_entries.extend(await _classify_batch(c, batch, model))
            except Exception as exc:
                print(f"  ERROR classifying unlabelled batch: {exc}", file=sys.stderr)
                all_entries.extend({**q, "main_category": "", "tags": []} for q in batch)

    # Pre-labelled questions whose category has sub-tags still need a
    # (cheaper, tags-only) GPT call — the sheet only gives main_category.
    needs_tags = [e for e in labelled if TAXONOMY.get(e["main_category"])]
    if needs_tags:
        print(f"  Assigning tags via GPT for {len(needs_tags)} pre-labelled question(s)...")
        c = ensure_client()
        by_cat: dict[str, list[dict]] = {}
        for e in needs_tags:
            by_cat.setdefault(e["main_category"], []).append(e)
        for cat, entries in by_cat.items():
            for start in range(0, len(entries), BATCH_SIZE):
                batch = entries[start : start + BATCH_SIZE]
                try:
                    tag_map = await _classify_tags_batch(c, batch, model, cat)
                    for i, e in enumerate(batch):
                        e["tags"] = tag_map.get(i, [])
                except Exception as exc:
                    print(f"  ERROR assigning tags for {cat}: {exc}", file=sys.stderr)

    return all_entries, client


# ---------------------------------------------------------------------------
# Shared entrypoint — the CLI (seed, below) and the in-app upload endpoint
# (main.py's /api/admin/seed-questions) both call this so format detection
# and classification logic exists in exactly one place. Never duplicate this
# logic at a call site — that's what silently drifted out of sync before.
# ---------------------------------------------------------------------------

async def classify_rows(
    rows: list[list[str]],
    model: str,
    client: AsyncOpenAI | None,
    settings,
) -> tuple[list[dict], str, AsyncOpenAI | None]:
    """Detect the sheet's format, extract questions, and classify/tag them.

    Returns (entries, format, client) — format is "A"/"B"/"C"/"unknown".
    Raises RuntimeError if GPT classification is needed but no API key is
    configured (callers decide how to surface that — CLI exits, the web
    endpoint turns it into an HTTP error).
    """
    fmt = _detect_format(rows)

    if fmt == "A":
        raw_questions = _extract_format_a(rows)
        if not raw_questions:
            return [], fmt, client

        if client is None:
            if not settings.openai_api_key:
                raise RuntimeError("OPENAI_API_KEY is required for Format A files.")
            client = AsyncOpenAI(api_key=settings.openai_api_key)

        all_entries: list[dict] = []
        for start in range(0, len(raw_questions), BATCH_SIZE):
            batch = raw_questions[start : start + BATCH_SIZE]
            try:
                all_entries.extend(await _classify_batch(client, batch, model))
            except Exception as exc:
                print(f"  ERROR classifying batch: {exc}", file=sys.stderr)
                all_entries.extend({**q, "main_category": "", "tags": []} for q in batch)
        return all_entries, fmt, client

    if fmt == "B":
        raw_entries = _extract_format_b(rows)
        if not raw_entries:
            return [], fmt, client
        all_entries, client = await _resolve_categories_and_tags(raw_entries, model, client, settings)
        return all_entries, fmt, client

    if fmt == "C":
        col_map = _detect_header_columns(rows[0])
        assert col_map is not None  # guaranteed by _detect_format returning "C"
        raw_entries = _extract_labeled(rows, col_map)
        if not raw_entries:
            return [], fmt, client
        all_entries, client = await _resolve_categories_and_tags(raw_entries, model, client, settings)
        return all_entries, fmt, client

    return [], fmt, client


# ---------------------------------------------------------------------------
# Main seed routine (CLI)
# ---------------------------------------------------------------------------

async def seed(html_paths: list[Path], model: str = "gpt-4o-mini") -> None:
    settings = get_settings()
    client: AsyncOpenAI | None = None
    total = 0

    for path in html_paths:
        print(f"\n[{path.name}]")
        rows = _parse_html(path)

        try:
            all_entries, fmt, client = await classify_rows(rows, model, client, settings)
        except RuntimeError as exc:
            print(f"ERROR: {exc}", file=sys.stderr)
            sys.exit(1)

        print(f"  Detected format: {fmt}")
        if fmt == "unknown":
            print("  ERROR: Could not detect table format — skipping.", file=sys.stderr)
            continue
        if not all_entries:
            print("  No questions found.")
            continue

        print(f"  Seeding {len(all_entries)} questions into question_bank...")
        try:
            result = await save_question_bank_batch(None, None, all_entries)  # type: ignore[arg-type]
            inserted = result["inserted"]
            if result["duplicates"]:
                print(f"  Skipped {result['duplicates']} duplicates. Inserted {inserted} new.")
        except Exception as exc:
            print(f"  ERROR seeding: {exc}", file=sys.stderr)
            continue

        from collections import Counter
        cats = Counter(e["main_category"] or "uncategorized" for e in all_entries)
        for cat, n in sorted(cats.items()):
            print(f"    {cat}: {n}")

        total += inserted

    print(f"\nDone. Total seeded: {total} questions.")


def main() -> None:
    args = sys.argv[1:]
    if not args:
        print("Usage: python seed_from_table.py file1.html [file2.html ...]")
        sys.exit(1)

    paths: list[Path] = []
    for a in args:
        p = Path(a)
        if not p.exists():
            print(f"WARNING: {p} not found — skipping.", file=sys.stderr)
        elif p.suffix.lower() not in (".html", ".htm"):
            print(f"WARNING: {p} not HTML — skipping.", file=sys.stderr)
        else:
            paths.append(p)

    if not paths:
        print("No valid HTML files.", file=sys.stderr)
        sys.exit(1)

    asyncio.run(seed(paths))


if __name__ == "__main__":
    main()