#!/usr/bin/env python3
"""
Simple extractor for STATdx snapshot HTML files.
Scans `xhr_captured/` for .html snapshot files (or a single file) and
extracts document sections and subsection points into JSON Lines.

Usage:
  python scrapers/extract_sections.py --input-dir xhr_captured --out-file xhr_captured/extracted_topics.jsonl

The script looks for <section class="document-page__section"> blocks and
collects H1 section titles and H2 subsection headers and <li class="text"> points.
"""
from bs4 import BeautifulSoup
import argparse
import json
import glob
import os
from pathlib import Path


def extract_from_html(html_text):
    soup = BeautifulSoup(html_text, "html.parser")
    result = []
    # The page seems to use <section class="document-page__section"> for large groups
    for sec in soup.find_all("section", class_="document-page__section"):
        section = {}
        # group headline (TERMINOLOGY, KEY FACTS, etc.)
        headline = sec.find(lambda tag: tag.name in ("div", "header") and tag.get("class") and any("headline3" in c for c in tag.get("class")))
        if headline:
            h1 = headline.find("h1")
            if h1:
                section_title = h1.get_text(strip=True)
            else:
                section_title = headline.get_text(strip=True)
            section["section_title"] = section_title
            section["section_id"] = headline.get("id") or None
        else:
            # fallback: try to find any h1
            h1 = sec.find("h1")
            if h1:
                section["section_title"] = h1.get_text(strip=True)
                section["section_id"] = h1.get("id")
            else:
                # skip empty sections
                continue

        subsections = []
        # common structure: <li class="section-title"><h2>Subheader</h2><ul class="section-points"><li class="text">Point</li></ul></li>
        for li in sec.find_all("li", class_="section-title"):
            sub = {}
            h2 = li.find("h2")
            if h2:
                sub["subheader"] = h2.get_text(strip=True)
            else:
                sub["subheader"] = None
            points = [p.get_text(strip=True) for p in li.find_all("li", class_="text")]
            # sometimes the structure nests <ul class="section-points"> directly inside the section
            if not points:
                # find any <li class="text"> inside this li
                points = [p.get_text(strip=True) for p in li.find_all("li") if "text" in (p.get("class") or [])]
            sub["points"] = points
            subsections.append(sub)

        # If no <li class=section-title>, collect direct <li class="text"> in section
        if not subsections:
            points = [p.get_text(strip=True) for p in sec.find_all("li", class_="text")]
            if points:
                subsections.append({"subheader": None, "points": points})

        section["subsections"] = subsections
        result.append(section)

    return result


def process_file(path):
    with open(path, "r", encoding="utf-8") as f:
        html = f.read()
    sections = extract_from_html(html)
    return sections


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--input-dir", default="xhr_captured", help="Directory with captured files")
    parser.add_argument("--input-file", help="Single HTML file to process (optional)")
    parser.add_argument("--out-file", default="xhr_captured/extracted_topics.jsonl", help="JSONL output file")
    args = parser.parse_args()

    paths = []
    if args.input_file:
        paths = [args.input_file]
    else:
        # prefer snapshot_after_capture files, fall back to any .html
        pattern = os.path.join(args.input_dir, "snapshot_after_capture_*.html")
        paths = glob.glob(pattern)
        if not paths:
            paths = glob.glob(os.path.join(args.input_dir, "*.html"))

    out_path = Path(args.out_file)
    out_path.parent.mkdir(parents=True, exist_ok=True)

    with out_path.open("w", encoding="utf-8") as out:
        for p in sorted(paths):
            sections = process_file(p)
            item = {
                "source_file": os.path.relpath(p),
                "sections": sections,
            }
            out.write(json.dumps(item, ensure_ascii=False) + "\n")

    print(f"Wrote {len(paths)} documents to {out_path}")


if __name__ == "__main__":
    main()