"""Batch feature extraction for all labeled pages.

Reads the VLM label CSV, extracts features for every page, saves to parquet.

Usage:
    agentic/.venv/bin/python -m classifier.extract_all
    agentic/.venv/bin/python -m classifier.extract_all --limit 100
    agentic/.venv/bin/python -m classifier.extract_all --resume
"""

from __future__ import annotations

import argparse
import logging
import sys
import time
from pathlib import Path

import pandas as pd

# Natural PDF lives one level up
sys.path.insert(0, str(Path(__file__).parent.parent.parent))

from classifier.features import FEATURE_NAMES, extract_features

logger = logging.getLogger(__name__)

LABELS_CSV = Path(__file__).parent.parent / "output" / "all_pages_labels.csv"
DATA_DIR = Path(__file__).parent / "data"

# The 21 original boolean label columns from the VLM CSV
BOOL_LABEL_COLS = [
    "has_handwriting",
    "has_hand_filled_form",
    "has_rotated_90",
    "has_low_quality_scan",
    "has_boxed_fields",
    "has_organizing_boxes",
    "has_checkboxes",
    "has_column_headers",
    "table_has_h_ruling_lines",
    "table_has_h_zebra_stripes",
    "table_has_h_whitespace",
    "table_has_v_ruling_lines",
    "table_has_v_whitespace",
    "has_key_value_pairs",
    "has_section_headers",
    "has_narrative",
    "has_multi_column_text",
    "has_empty_form",
    "has_fill_in_blanks",
    "has_visual_media",
    "has_multiple_records",
]

# All 23 labels (21 originals + 2 derived)
ALL_LABEL_COLS = BOOL_LABEL_COLS + ["has_table", "has_multiple_tables"]


def load_labels() -> pd.DataFrame:
    """Load and prepare the label CSV."""
    df = pd.read_csv(LABELS_CSV, dtype={"page_number": int})

    # Filter ERROR rows
    error_mask = df["page_description"].astype(str).str.startswith("ERROR")
    n_errors = error_mask.sum()
    if n_errors:
        logger.info("Filtering %d ERROR rows", n_errors)
        df = df[~error_mask].copy()

    # Derive binary table labels
    df["has_table"] = (df["table_count"].fillna(0).astype(int) > 0)
    df["has_multiple_tables"] = (df["table_count"].fillna(0).astype(int) > 1)

    # Convert bool columns to int (True->1, False->0)
    for col in ALL_LABEL_COLS:
        df[col] = df[col].map(
            {"True": 1, "False": 0, True: 1, False: 0}
        ).fillna(0).astype(int)

    return df[["pdf_path", "page_number"] + ALL_LABEL_COLS]


def main():
    parser = argparse.ArgumentParser(description="Extract features for labeled pages")
    parser.add_argument("--limit", type=int, help="Max pages to process")
    parser.add_argument("--resume", action="store_true", help="Skip already-processed pages")
    args = parser.parse_args()

    logging.basicConfig(level=logging.INFO, format="%(message)s")

    DATA_DIR.mkdir(parents=True, exist_ok=True)

    features_path = DATA_DIR / "features.parquet"
    labels_path = DATA_DIR / "labels.parquet"

    # Load labels
    labels_df = load_labels()
    logger.info("Loaded %d labeled pages", len(labels_df))

    # Resume support
    done_keys = set()
    existing_features = []
    if args.resume and features_path.exists():
        prev = pd.read_parquet(features_path)
        done_keys = set(zip(prev["pdf_path"], prev["page_number"]))
        existing_features = [prev]
        logger.info("Resuming: %d pages already done", len(done_keys))

    # Group by PDF for efficient open/close
    grouped = labels_df.groupby("pdf_path", sort=True)
    total_pdfs = len(grouped)

    from natural_pdf import PDF

    results = []
    processed = 0
    errors = 0
    start = time.monotonic()

    for pdf_idx, (pdf_path, group) in enumerate(grouped):
        # Skip if all pages for this PDF already done
        if args.resume:
            pending = group[
                ~group.apply(lambda r: (r["pdf_path"], r["page_number"]) in done_keys, axis=1)
            ]
            if pending.empty:
                continue
        else:
            pending = group

        try:
            pdf = PDF(str(pdf_path))
        except Exception as e:
            logger.warning("FAIL open: %s — %s", pdf_path, e)
            errors += len(pending)
            continue

        for _, row in pending.iterrows():
            if args.limit and processed >= args.limit:
                break

            page_num = int(row["page_number"])
            try:
                page = pdf.pages[page_num]
                feats = extract_features(page)
                feats["pdf_path"] = str(pdf_path)
                feats["page_number"] = page_num
                results.append(feats)
            except Exception as e:
                logger.warning("FAIL %s p%d: %s", Path(pdf_path).name, page_num, e)
                # Write zeros for failed pages
                feats = {k: 0.0 for k in FEATURE_NAMES}
                feats["pdf_path"] = str(pdf_path)
                feats["page_number"] = page_num
                results.append(feats)
                errors += 1

            processed += 1

        pdf.close()

        if args.limit and processed >= args.limit:
            break

        # Progress every 50 PDFs
        if (pdf_idx + 1) % 50 == 0 or pdf_idx == total_pdfs - 1:
            elapsed = time.monotonic() - start
            rate = processed / elapsed if elapsed > 0 else 0
            logger.info(
                "[%d/%d PDFs] %d pages, %d errors, %.1f pg/s, %.0fs elapsed",
                pdf_idx + 1,
                total_pdfs,
                processed,
                errors,
                rate,
                elapsed,
            )

        # Periodic save every 200 PDFs
        if (pdf_idx + 1) % 200 == 0 and results:
            _save_intermediate(results, existing_features, features_path)

    # Final save
    if results or existing_features:
        all_features = existing_features + [pd.DataFrame(results)]
        features_df = pd.concat(all_features, ignore_index=True)
        features_df.to_parquet(features_path, index=False)
        logger.info("Saved %d rows to %s", len(features_df), features_path)

    # Save labels (matching the features)
    labels_df.to_parquet(labels_path, index=False)
    logger.info("Saved %d label rows to %s", len(labels_df), labels_path)

    elapsed = time.monotonic() - start
    logger.info(
        "\nDone: %d pages, %d errors, %.0fs (%.1f pg/s)",
        processed,
        errors,
        elapsed,
        processed / elapsed if elapsed > 0 else 0,
    )


def _save_intermediate(results, existing, path):
    """Save intermediate results to parquet."""
    all_dfs = existing + [pd.DataFrame(results)]
    combined = pd.concat(all_dfs, ignore_index=True)
    combined.to_parquet(path, index=False)
    logger.info("  Intermediate save: %d rows", len(combined))


if __name__ == "__main__":
    main()
