# ============================================================
# SPID — extraction engine
#
# Iterates enriched sections, runs the chosen extractor
# (rules or gpt-4o-mini), filters by confidence, and loads
# mentions into DuckDB.  LLM calls are threaded (IO-bound).
# ============================================================
from __future__ import annotations

from concurrent.futures import ThreadPoolExecutor, as_completed

from spid.config import get_logger, get_settings
from spid.extract.base import get_extractor
from spid.store import Store

log = get_logger("pipeline")


def _ctx_from_rows(rows: list[dict]) -> dict:
    ctx = {}
    for r in rows:
        ctx.setdefault(r["accession_number"], {
            "ticker": r.get("ticker"),
            "company_name": r.get("company_name") or r.get("filing_company"),
            "filing_company": r.get("filing_company"),
            "form_type": r.get("form_type"),
            "filing_date": r.get("filing_date"),
            "fiscal_year": r.get("fiscal_year"),
            "sector": r.get("sector"),
        })
    return ctx


def run_extraction(store: Store, *, backend: str | None = None,
                   ciks: list[str] | None = None, forms: list[str] | None = None,
                   limit_sections: int | None = None, min_words: int = 30,
                   workers: int = 8, min_confidence: float | None = None,
                   flush_every: int = 500) -> dict:
    s = get_settings()
    backend = backend or s.backend
    min_conf = s.min_confidence if min_confidence is None else min_confidence
    extractor = get_extractor(backend)

    rows = store.sections_for(ciks=ciks, forms=forms, limit=limit_sections, min_words=min_words)
    ctx = _ctx_from_rows(rows)
    log.info("Extracting from %d sections | backend=%s | workers=%d", len(rows), backend, workers)

    kept = 0
    buf = []
    done_rows: list[tuple] = []

    def work(r: dict):
        return extractor.extract(text=r["text"], accession_number=r["accession_number"],
                                 cik=r["cik"], section_id=r["section_id"],
                                 section_name=r["section_name"])

    def flush():
        nonlocal kept, buf, done_rows
        if buf:
            kept += store.upsert_mentions(buf, ctx)
            buf = []
        if done_rows:
            store.mark_extracted(done_rows)
            done_rows = []

    done = 0
    if backend == "llm" and workers > 1:
        with ThreadPoolExecutor(max_workers=workers) as ex:
            futs = {ex.submit(work, r): r for r in rows}
            for fut in as_completed(futs):
                done += 1
                r = futs[fut]
                try:
                    ms = [m for m in fut.result() if m.confidence >= min_conf]
                    done_rows.append((r["section_id"], r["accession_number"], backend, len(ms)))
                except Exception as e:  # noqa: BLE001
                    log.warning("section failed: %s", e); ms = []
                buf.extend(ms)
                if len(buf) >= flush_every:
                    flush()
                if done % 200 == 0:
                    log.info("  ... %d/%d sections, %d mentions kept", done, len(rows), kept + len(buf))
    else:
        for r in rows:
            done += 1
            ms = [m for m in work(r) if m.confidence >= min_conf]
            done_rows.append((r["section_id"], r["accession_number"], backend, len(ms)))
            buf.extend(ms)
            if len(buf) >= flush_every:
                flush()
            if done % 2000 == 0:
                log.info("  ... %d/%d sections, %d mentions kept", done, len(rows), kept + len(buf))
    flush()
    log.info("Extraction complete: %d mentions from %d sections.", kept, len(rows))
    return {"sections": len(rows), "mentions": kept, "backend": backend}