from collections.abc import Callable, Sequence
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass

from techdocs_graphrag.extraction import GraphExtractor
from techdocs_graphrag.ingestion import ParagraphSplitter, SourceDocument, chunk_documents
from techdocs_graphrag.ports import KnowledgeGraph, LanguageModel

ProgressCallback = Callable[[int, int], None]


@dataclass(frozen=True, slots=True)
class IndexingReport:
    documents: int
    chunks: int
    entities: int
    relations: int


class IndexingPipeline:
    def __init__(
        self,
        *,
        splitter: ParagraphSplitter,
        llm: LanguageModel,
        extractor: GraphExtractor,
        store: KnowledgeGraph,
        workers: int,
    ) -> None:
        self._splitter = splitter
        self._llm = llm
        self._extractor = extractor
        self._store = store
        self._workers = workers

    def run(
        self,
        documents: Sequence[SourceDocument],
        on_progress: ProgressCallback | None = None,
    ) -> IndexingReport:
        chunks = chunk_documents(documents, self._splitter)
        if not chunks:
            return IndexingReport(documents=len(documents), chunks=0, entities=0, relations=0)

        self._store.ensure_schema()
        self._store.upsert_chunks(chunks, self._llm.embed([chunk.text for chunk in chunks]))

        entities = relations = 0
        with ThreadPoolExecutor(max_workers=self._workers) as executor:
            extractions = executor.map(self._extractor.extract, chunks)
            for completed, (chunk, extraction) in enumerate(
                zip(chunks, extractions, strict=True), start=1
            ):
                self._store.upsert_extraction(chunk, extraction)
                entities += len(extraction.entities)
                relations += len(extraction.relations)
                if on_progress:
                    on_progress(completed, len(chunks))

        self._store.await_indexes()
        return IndexingReport(
            documents=len(documents),
            chunks=len(chunks),
            entities=entities,
            relations=relations,
        )