"""SetOfMarkExtractor (Set-of-Mark): vision identifies tables, Camelot fills them in-region.

Thin wrapper that renders the document's pages once, wires `SetOfMarkDeps`, and
runs the vision-guided extraction graph built in `pipeline`. The graph does the
work: locate every table, capture each in its own region, clean structure under
the no-number guard, and reading-order the result. Emits one `ExtractedTable`
per visual table; a region where Camelot finds nothing is still emitted with
its visual identity so a table is never silently dropped.
"""

from __future__ import annotations

import asyncio
import tempfile
from pathlib import Path
from typing import List, Optional

import fitz

from quber.agents.capture_advisor import CaptureAdvisor, get_capture_advisor
from quber.agents.grid_locator import GridLocator, get_grid_locator
from quber.agents.llm_client import LLMClient, get_llm_client
from quber.agents.status_inspector import StatusInspector, get_status_inspector
from quber.core.extractors.base import ExtractedTable
from quber.core.extractors.camelot.acquire import render_pages
from quber.core.extractors.set_of_mark.pipeline import SetOfMarkDeps, SetOfMarkState, build_set_of_mark_graph


class SetOfMarkExtractor:
    """Set-of-Mark extractor: vision locates, Camelot fills in-region, correction cleans."""

    def __init__(
        self,
        locator: Optional[GridLocator] = None,
        llm: Optional[LLMClient] = None,
        dpi: int = 200,
        page_concurrency: int = 16,
        correct_concurrency: int = 16,
        advisor: Optional[CaptureAdvisor] = None,
        inspector: Optional[StatusInspector] = None,
    ) -> None:
        self.locator = locator or get_grid_locator()
        self.llm = llm or get_llm_client(None)
        self.dpi = dpi
        self.page_concurrency = page_concurrency
        self.correct_concurrency = correct_concurrency
        self.advisor = advisor if advisor is not None else get_capture_advisor()
        self.inspector = inspector if inspector is not None else get_status_inspector()

    async def extract_tables(self, source: Path) -> List[ExtractedTable]:
        doc = fitz.open(str(source))
        page_dims = {i + 1: (doc[i].rect.width, doc[i].rect.height) for i in range(doc.page_count)}
        with tempfile.TemporaryDirectory(prefix="quber-som-") as tmp:
            images = await asyncio.to_thread(render_pages, source, self.dpi, Path(tmp))
            page_images = {i + 1: img for i, img in enumerate(images)}
            deps = SetOfMarkDeps(
                locator=self.locator,
                llm=self.llm,
                source=source,
                page_images=page_images,
                page_dims=page_dims,
                dpi=self.dpi,
                correct_sem=asyncio.Semaphore(self.correct_concurrency),
                page_sem=asyncio.Semaphore(self.page_concurrency),
                advisor=self.advisor,
                inspector=self.inspector,
            )
            return await build_set_of_mark_graph().run(state=SetOfMarkState(), deps=deps, inputs=source)

    def extract_tables_sync(self, source: Path) -> List[ExtractedTable]:
        return asyncio.run(self.extract_tables(source))