from __future__ import annotations

from pathlib import Path

from Mama.models import ExecutionContext, IngestResult
from Mama.ports import RuntimePorts

try:
    from langchain_core.documents import Document  # type: ignore
except Exception:  # pragma: no cover
    from langchain.schema import Document  # type: ignore

try:
    from langchain_community.document_loaders import CSVLoader, PyPDFDirectoryLoader, PyPDFLoader  # type: ignore
except Exception:  # pragma: no cover
    from langchain.document_loaders import PyPDFDirectoryLoader, PyPDFLoader  # type: ignore
    from langchain.document_loaders.csv_loader import CSVLoader  # type: ignore

try:
    from langchain_text_splitters import RecursiveCharacterTextSplitter  # type: ignore
except Exception:  # pragma: no cover
    from langchain.text_splitter import RecursiveCharacterTextSplitter  # type: ignore


def _split_documents(documents: list[Document], chunk_size: int, chunk_overlap: int) -> list[Document]:
    splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap)
    return splitter.split_documents(documents)


def load_pdf_directory(
    source_dir: str,
    *,
    chunk_size: int = 500,
    chunk_overlap: int = 0,
) -> list[Document]:
    docs = PyPDFDirectoryLoader(source_dir).load()
    return _split_documents(docs, chunk_size=chunk_size, chunk_overlap=chunk_overlap)


def load_single_pdf(
    file_path: str,
    *,
    chunk_size: int = 500,
    chunk_overlap: int = 0,
) -> list[Document]:
    docs = PyPDFLoader(file_path).load()
    return _split_documents(docs, chunk_size=chunk_size, chunk_overlap=chunk_overlap)


def load_csv(
    file_path: str,
    *,
    field_names: list[str] | None = None,
    delimiter: str = ",",
    quote_char: str = '"',
    source_column: str | None = None,
) -> list[Document]:
    csv_args = {
        "delimiter": delimiter,
        "quotechar": quote_char,
    }
    if field_names:
        csv_args["fieldnames"] = field_names

    loader = CSVLoader(file_path=file_path, csv_args=csv_args, source_column=source_column)
    return loader.load()


def ingest_documents(
    ctx: ExecutionContext,
    documents: list[Document],
    runtime: RuntimePorts,
) -> IngestResult:
    if not documents:
        raise ValueError("documents cannot be empty")

    embeddings = runtime.embeddings_provider.get_embeddings(ctx)
    store = runtime.vector_store_factory.get_store(ctx, embeddings)
    chunks_added = store.add_documents(documents)

    return IngestResult(
        kb_id=ctx.kb_id,
        backend=ctx.storage.backend,
        chunks_added=chunks_added,
        index_identifier=store.identifier(),
        metadata={"tenant_id": ctx.tenant_id},
    )


def ingest_pdf_directory(ctx: ExecutionContext, source_dir: str, runtime: RuntimePorts) -> IngestResult:
    if not Path(source_dir).exists():
        raise FileNotFoundError(f"source directory not found: {source_dir}")
    docs = load_pdf_directory(source_dir)
    return ingest_documents(ctx, docs, runtime)


def ingest_single_pdf(ctx: ExecutionContext, file_path: str, runtime: RuntimePorts) -> IngestResult:
    if not Path(file_path).exists():
        raise FileNotFoundError(f"source file not found: {file_path}")
    docs = load_single_pdf(file_path)
    return ingest_documents(ctx, docs, runtime)


def ingest_csv(ctx: ExecutionContext, file_path: str, runtime: RuntimePorts, **kwargs) -> IngestResult:
    if not Path(file_path).exists():
        raise FileNotFoundError(f"source file not found: {file_path}")
    docs = load_csv(file_path, **kwargs)
    return ingest_documents(ctx, docs, runtime)
