"""Step 1: Search DocumentCloud for checkbox-bearing documents."""

import json
import logging
from collections import defaultdict
from typing import Dict, List, Set

from .config import (
    ALL_QUERIES,
    DOCUMENTS_JSONL,
    MAX_DOCS_PER_ORG,
    MAX_DOCS_PER_QUERY,
    MAX_PAGE_COUNT,
    MIN_PAGE_COUNT,
    check_rate_limit,
    get_dc_client,
)
from .state import PipelineState

logger = logging.getLogger(__name__)


def search(state: PipelineState, queries: List[str] = None) -> int:
    """Search DocumentCloud and write candidate documents to documents.jsonl.

    Returns the number of new documents found.
    """
    queries = queries or ALL_QUERIES
    client = get_dc_client()

    seen_ids: Set[str] = set()
    org_counts: Dict[str, int] = defaultdict(int)
    documents: List[Dict] = []

    # Load existing documents if resuming
    if DOCUMENTS_JSONL.exists():
        with open(DOCUMENTS_JSONL) as f:
            for line in f:
                doc = json.loads(line)
                seen_ids.add(str(doc["id"]))
                org_counts[doc.get("organization", "")] += 1
        logger.info("Loaded %d existing documents", len(seen_ids))

    new_count = 0
    per_page = min(MAX_DOCS_PER_QUERY, 100)  # DC API caps at 100 per page

    for query in queries:
        logger.info("Searching: %s", query)

        query_hits = 0
        docs_examined = 0

        try:
            # DC client returns a lazy paginator — iterate without list()
            results = client.documents.search(
                query,
                per_page=per_page,
                page=1,
            )
        except Exception as e:
            check_rate_limit(e)
            logger.warning("Search failed for query %r: %s", query, e)
            continue

        for doc in results:
            docs_examined += 1
            if docs_examined > MAX_DOCS_PER_QUERY:
                break

            doc_id = str(doc.id)

            # Dedup
            if doc_id in seen_ids:
                continue

            # Page count filter
            page_count = getattr(doc, "page_count", None) or 0
            if page_count < MIN_PAGE_COUNT or page_count > MAX_PAGE_COUNT:
                continue

            # Per-org diversity cap (use organization_id to avoid lazy API fetch)
            org_id = getattr(doc, "organization_id", None)
            org_name = str(org_id) if org_id else "unknown"
            if org_counts[org_name] >= MAX_DOCS_PER_ORG:
                continue

            seen_ids.add(doc_id)
            org_counts[org_name] += 1
            query_hits += 1

            record = {
                "id": doc_id,
                "title": getattr(doc, "title", ""),
                "organization": org_name,
                "page_count": page_count,
                "canonical_url": getattr(doc, "canonical_url", ""),
                "pdf_url": getattr(doc, "pdf_url", ""),
                "contributing_queries": [query],
            }
            documents.append(record)

            # Register in state
            state.set_batch(doc_id, {"searched": True})
            new_count += 1

        logger.info("  Query yielded %d new documents (%d examined)", query_hits, docs_examined)

    # Load existing records, merge new ones, rewrite entire file (no duplicates)
    DOCUMENTS_JSONL.parent.mkdir(parents=True, exist_ok=True)
    existing = []
    existing_ids: Set[str] = set()
    if DOCUMENTS_JSONL.exists():
        with open(DOCUMENTS_JSONL) as f:
            for line in f:
                line = line.strip()
                if line:
                    rec = json.loads(line)
                    existing_ids.add(str(rec["id"]))
                    existing.append(rec)

    new_records = [d for d in documents if str(d["id"]) not in existing_ids]
    all_records = existing + new_records

    tmp = DOCUMENTS_JSONL.with_suffix(".tmp")
    with open(tmp, "w") as f:
        for doc in all_records:
            f.write(json.dumps(doc) + "\n")
    tmp.replace(DOCUMENTS_JSONL)

    logger.info("Search complete. %d new documents (%d total)", new_count, len(seen_ids))
    return new_count


def load_documents() -> List[Dict]:
    """Load all document records from documents.jsonl."""
    docs = []
    if DOCUMENTS_JSONL.exists():
        with open(DOCUMENTS_JSONL) as f:
            for line in f:
                line = line.strip()
                if line:
                    docs.append(json.loads(line))
    return docs


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