"""High-level pipeline helpers for fast user onboarding."""

from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path
from typing import Any, Dict, Optional

from .kb_manager import KBManager
from .pdf_parser import parse_pdfs
from .vector_kb import KBBuildResult, create_vector_kb


@dataclass
class PipelineResult:
    kb_name: str
    kb_path: str
    parse_stats: Dict[str, Any]
    build_result: KBBuildResult
    manager: Optional[KBManager] = None


def build_kb_from_pdfs(
    *,
    name: str,
    input_dir: str,
    output_dir: str,
    embedding_model: Optional[str] = "all-MiniLM-L6-v2",
    parser_backend: str = "auto",
    image_backend: str = "auto",
    docling_artifacts_path: Optional[str] = None,
    mode: str = "rebuild",
    overwrite: bool = False,
    chunk_size: int = 800,
    chunk_overlap: int = 80,
    register_in_manager: bool = False,
    manager: Optional[KBManager] = None,
) -> PipelineResult:
    """
    One-shot pipeline:
    1) parse PDFs to structured markdown/json
    2) build vector DB
    3) optionally register into KBManager
    """
    out_root = Path(output_dir)
    parsed_dir = out_root / "parsed" / name
    vector_dir = out_root / "vector_db" / name

    parse_stats = parse_pdfs(
        input_dir=input_dir,
        output_dir=str(parsed_dir),
        parser_backend=parser_backend,  # type: ignore[arg-type]
        image_backend=image_backend,  # type: ignore[arg-type]
        docling_artifacts_path=docling_artifacts_path,
    )

    build_result = create_vector_kb(
        name=name,
        input_dir=str(parsed_dir),
        output_dir=str(vector_dir),
        overwrite=overwrite,
        embedding_model=embedding_model,
        chunk_size=chunk_size,
        chunk_overlap=chunk_overlap,
        mode=mode,  # type: ignore[arg-type]
        return_result=True,
    )
    assert isinstance(build_result, KBBuildResult)

    result_manager = manager
    if register_in_manager:
        result_manager = manager or KBManager(default_embedding_model=embedding_model)
        result_manager.add_kb(name, build_result.chroma_path, embedding_model=embedding_model)

    return PipelineResult(
        kb_name=name,
        kb_path=build_result.chroma_path,
        parse_stats=parse_stats,
        build_result=build_result,
        manager=result_manager,
    )


def load_kb(
    *,
    kb_path: str,
    kb_name: str = "default",
    manager: Optional[KBManager] = None,
    embedding_model: Optional[str] = None,
) -> KBManager:
    """Load an existing vector DB path into KBManager."""
    mgr = manager or KBManager(default_embedding_model=embedding_model)
    mgr.add_kb(kb_name, kb_path, embedding_model=embedding_model)
    return mgr


__all__ = ["PipelineResult", "build_kb_from_pdfs", "load_kb"]
