"""
extractor — Production-grade document extraction to versioned RAG-ready JSON/JSONL.

Public API:
    from extractor import extract
    for element in extract("paper.pdf", mode="elements"):
        print(element.element_type, element.section_path, element.text[:80])
"""

from __future__ import annotations

import asyncio
import concurrent.futures
from collections.abc import AsyncIterator, Iterable, Iterator
from importlib.metadata import PackageNotFoundError as _PNF
from importlib.metadata import version as _pkg_version
from pathlib import Path
from typing import Literal, overload

from .errors import ExtractorError, ParserError, QuarantineError, UnsupportedFormatError
from .schema import (
    ATOMIC_TYPES,
    ELEMENT_TYPES,
    SCHEMA_VERSION,
    BaseElement,
    ChunkRecord,
    EntityAnnotation,
    ManifestRecord,
    RunConfig,
)

try:
    __version__ = _pkg_version("coresdk-extractor")
except _PNF:
    __version__ = "0.0.0+dev"

# extractor.analytics is a sub-package for cross-document entity intelligence.
# It is not imported eagerly here. Import explicitly when needed:
#   from extractor.analytics import EntityIndex, cite_entity
__all__ = [
    "extract",
    "extract_batch",
    "aextract",
    "aextract_batch",
    "SCHEMA_VERSION",
    "__version__",
    "ExtractorError",
    "ParserError",
    "QuarantineError",
    "UnsupportedFormatError",
    "BaseElement",
    "ChunkRecord",
    "EntityAnnotation",
    "ManifestRecord",
    "ELEMENT_TYPES",
    "ATOMIC_TYPES",
    "RunConfig",
]


@overload
def extract(
    source: Path | str,
    mode: Literal["elements"] = ...,
    *,
    chunk_size: int = ...,
    overlap: int = ...,
    strategy: str = ...,
    include_types: list[str] | None = ...,
    exclude_types: list[str] | None = ...,
    tokenizer: str = ...,
    inject_context_prefix: bool = ...,
    parent_size: int = ...,
    child_size: int = ...,
    entities: bool = ...,
    entities_model: str = ...,
    entities_types: list[str] | None = ...,
    entities_threshold: float = ...,
) -> Iterator[BaseElement]: ...


@overload
def extract(
    source: Path | str,
    mode: Literal["chunks", "dual-chunks"],
    *,
    chunk_size: int = ...,
    overlap: int = ...,
    strategy: str = ...,
    include_types: list[str] | None = ...,
    exclude_types: list[str] | None = ...,
    tokenizer: str = ...,
    inject_context_prefix: bool = ...,
    parent_size: int = ...,
    child_size: int = ...,
    entities: bool = ...,
    entities_model: str = ...,
    entities_types: list[str] | None = ...,
    entities_threshold: float = ...,
) -> Iterator[ChunkRecord]: ...


def extract(
    source: str | Path,
    mode: str = "elements",
    chunk_size: int = 512,
    overlap: int = 0,
    strategy: str = "fast",
    include_types: list[str] | None = None,
    exclude_types: list[str] | None = None,
    tokenizer: str = "cl100k_base",
    inject_context_prefix: bool = False,
    parent_size: int = 512,
    child_size: int = 128,
    entities: bool = False,
    entities_model: str = "fastino/gliner2-base-v1",
    entities_types: list[str] | None = None,
    entities_threshold: float = 0.50,
    classify_labels: list[str] | None = None,
    relation_types: list[str] | None = None,
    extract_schema: dict | None = None,
    extract_on_types: list[str] | None = None,
) -> Iterator[BaseElement]:
    """
    Extract a document and yield BaseElement records.

    Args:
        source: Path to the source file or URL
        mode: "elements" (raw semantic units), "chunks" (section-aware), or
              "dual-chunks" (paired coarse+fine chunks for ParentDocumentRetriever)
        chunk_size: Maximum tokens per chunk (default 512)
        overlap: Overlap tokens between chunks (default 0, disabled)
        strategy: "fast" | "accurate" | "ocr"
        include_types: Only emit these element types (default: all)
        exclude_types: Suppress these element types
        parent_size: Token budget for coarse (parent) chunks in dual-chunks mode
        child_size: Token budget for fine (child) chunks in dual-chunks mode

    Yields:
        BaseElement records in document order
    """
    from . import pipeline

    yield from pipeline.run(
        path=Path(source),
        mode=mode,
        chunk_size=chunk_size,
        overlap=overlap,
        strategy=strategy,
        include_types=include_types,
        exclude_types=exclude_types,
        tokenizer=tokenizer,
        inject_context_prefix=inject_context_prefix,
        parent_size=parent_size,
        child_size=child_size,
        entities=entities,
        entities_model=entities_model,
        entities_types=entities_types,
        entities_threshold=entities_threshold,
        classify_labels=classify_labels,
        relation_types=relation_types,
        extract_schema=extract_schema,
        extract_on_types=extract_on_types,
    )


def extract_batch(
    sources: Iterable[Path | str],
    *,
    workers: int = 4,
    on_error: Literal["skip", "raise"] = "skip",
    **kwargs,
) -> Iterator[tuple[Path, list[BaseElement] | Exception]]:
    """Extract multiple files in parallel. Yields (path, elements_or_exception) tuples."""
    paths = [Path(s) for s in sources]

    def _extract_one(p: Path) -> tuple[Path, list[BaseElement] | Exception]:
        try:
            return (p, list(extract(p, **kwargs)))
        except Exception as e:
            if on_error == "raise":
                raise
            return (p, e)

    with concurrent.futures.ThreadPoolExecutor(max_workers=workers) as executor:
        futures = {executor.submit(_extract_one, p): p for p in paths}
        for future in concurrent.futures.as_completed(futures):
            yield future.result()


async def aextract(
    source: Path | str,
    **kwargs,
) -> AsyncIterator[BaseElement]:
    """Async streaming extraction — wraps extract() to avoid blocking the event loop."""
    queue: asyncio.Queue = asyncio.Queue(maxsize=64)
    loop = asyncio.get_running_loop()

    async def _enqueue():
        def _run_sync():
            try:
                for el in extract(source, **kwargs):
                    asyncio.run_coroutine_threadsafe(queue.put(el), loop).result()
            finally:
                asyncio.run_coroutine_threadsafe(queue.put(None), loop).result()

        await loop.run_in_executor(None, _run_sync)

    task = asyncio.ensure_future(_enqueue())
    try:
        while True:
            item = await queue.get()
            if item is None:
                break
            yield item
    finally:
        await task  # propagate exceptions from worker


async def aextract_batch(
    sources: Iterable[Path | str],
    *,
    workers: int = 4,
    **kwargs,
) -> AsyncIterator[tuple[Path, list[BaseElement] | Exception]]:
    """Async batch extraction with concurrency control."""
    semaphore = asyncio.Semaphore(workers)
    paths = [Path(s) for s in sources]

    async def _extract_one(p: Path) -> tuple[Path, list[BaseElement] | Exception]:
        async with semaphore:
            try:
                elements = []
                async for el in aextract(p, **kwargs):
                    elements.append(el)
                return (p, elements)
            except Exception as e:
                return (p, e)

    tasks = [asyncio.create_task(_extract_one(p)) for p in paths]
    for coro in asyncio.as_completed(tasks):
        yield await coro
