from __future__ import annotations

from pathlib import Path
from typing import Any

from ketju.config import DOCS_DIR
from ketju.rag.llama_index.ingestion.simple_pdf import SimplePdfIngestionHandler
from ketju.rag.llama_index.storage.pgvector import PgVectorStorageBackend
from ketju.rag.llama_index.variants.simple import LlamaIndexRag


class PgVectorRag(LlamaIndexRag):
    """Compatibility wrapper. Prefer ``build_pgvector_rag``."""

    def __init__(
        self,
        name: str = "pgvector_rag",
        *,
        database_url: str | None = None,
        table_name: str = "ketju_vectors",
        schema_name: str = "public",
        embed_dim: int | None = None,
        hybrid_search: bool = False,
        text_search_config: str | None = None,
        hnsw_kwargs: dict[str, Any] | None = None,
        persist_dir: Path | None = None,
        docs_dir: Path | None = None,
    ) -> None:
        storage_backend = PgVectorStorageBackend(
            database_url=database_url,
            table_name=table_name,
            schema_name=schema_name,
            embed_dim=embed_dim,
            hybrid_search=hybrid_search,
            text_search_config=text_search_config,
            hnsw_kwargs=hnsw_kwargs,
        )
        super().__init__(
            name=name,
            persist_dir=persist_dir,
            docs_dir=docs_dir or DOCS_DIR,
            ingestion_handler=SimplePdfIngestionHandler(),
            storage_backend=storage_backend,
        )


def build_pgvector_rag(
    *,
    name: str = "pgvector_rag",
    database_url: str | None = None,
    table_name: str = "ketju_vectors",
    schema_name: str = "public",
    embed_dim: int | None = None,
    hybrid_search: bool = False,
    text_search_config: str | None = None,
    hnsw_kwargs: dict[str, Any] | None = None,
    persist_dir: Path | None = None,
    docs_dir: Path | None = None,
) -> LlamaIndexRag:
    """Create a LlamaIndexRag wired with pgvector storage + simple PDF ingestion."""
    storage_backend = PgVectorStorageBackend(
        database_url=database_url,
        table_name=table_name,
        schema_name=schema_name,
        embed_dim=embed_dim,
        hybrid_search=hybrid_search,
        text_search_config=text_search_config,
        hnsw_kwargs=hnsw_kwargs,
    )
    return LlamaIndexRag(
        name=name,
        persist_dir=persist_dir,
        docs_dir=docs_dir or DOCS_DIR,
        ingestion_handler=SimplePdfIngestionHandler(),
        storage_backend=storage_backend,
    )
