"""Step 2: Text-based page screening via DocumentCloud API."""

import logging
from typing import Dict, List

from .config import CHECKBOX_SIGNALS, MIN_TEXT_SCORE, check_rate_limit, get_dc_client
from .search import load_documents
from .state import PipelineState

logger = logging.getLogger(__name__)


def score_page(text: str) -> int:
    """Score a page's text for checkbox signals."""
    text_lower = text.lower()
    return sum(1 for sig in CHECKBOX_SIGNALS if sig in text_lower)


def prefilter(state: PipelineState) -> Dict[str, List[int]]:
    """Screen pages by text signals using DocumentCloud's text API.

    Returns dict mapping doc_id -> list of qualifying page indices.
    """
    client = get_dc_client()
    documents = load_documents()
    results: Dict[str, List[int]] = {}
    skipped = 0
    total_pages_kept = 0

    try:
        from tqdm import tqdm
        doc_iter = tqdm(documents, desc="Pre-filtering documents", unit="doc")
    except ImportError:
        doc_iter = documents

    for doc_record in doc_iter:
        doc_id = str(doc_record["id"])

        if state.is_done(doc_id, "prefiltered"):
            # Already done — load cached result
            cached = state.get(doc_id, "candidate_pages", [])
            if cached:
                results[doc_id] = cached
            continue

        page_count = doc_record.get("page_count", 0)
        if page_count == 0:
            state.set_batch(doc_id, {"prefiltered": True, "candidate_pages": []})
            skipped += 1
            continue

        try:
            dc_doc = client.documents.get(doc_id)
        except Exception as e:
            check_rate_limit(e)
            logger.warning("Failed to fetch document %s: %s", doc_id, e)
            skipped += 1
            continue

        # Fetch full text in one API call, then split by page
        candidate_pages = []
        try:
            full_text = dc_doc.full_text
            # DocumentCloud separates pages with form-feed characters
            page_texts = full_text.split("\f") if full_text else []
        except Exception as e:
            check_rate_limit(e)
            # Fall back to per-page fetching if full_text fails
            page_texts = []
            for page_idx in range(page_count):
                try:
                    page_texts.append(dc_doc.get_page_text(page_idx + 1))
                except Exception as e2:
                    check_rate_limit(e2)
                    page_texts.append("")

        for page_idx, text in enumerate(page_texts):
            if score_page(text) >= MIN_TEXT_SCORE:
                candidate_pages.append(page_idx)  # Store as 0-indexed

        state.set_batch(doc_id, {
            "prefiltered": True,
            "candidate_pages": candidate_pages,
        })

        if candidate_pages:
            results[doc_id] = candidate_pages
            total_pages_kept += len(candidate_pages)
            logger.debug("Doc %s: %d/%d pages pass text filter", doc_id, len(candidate_pages), page_count)
        else:
            logger.debug("Doc %s: no pages pass text filter", doc_id)

    total_docs = len(documents)
    docs_with_pages = len(results)
    logger.info(
        "Pre-filter complete. %d/%d docs have qualifying pages (%d total pages). %d skipped.",
        docs_with_pages, total_docs, total_pages_kept, skipped,
    )
    return results


if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)
    st = PipelineState()
    prefilter(st)
