from __future__ import annotations

from Mama.ingestion import (
    ingest_csv,
    ingest_documents,
    ingest_pdf_directory,
    ingest_single_pdf,
    load_csv,
    load_pdf_directory,
    load_single_pdf,
)
from Mama.models import ExecutionContext, IngestResult
from Mama.ports import RuntimePorts
from Mama.runtime import default_runtime


def index_documents(source_dir: str, *, chunk_size: int = 500, chunk_overlap: int = 0):
    return load_pdf_directory(source_dir, chunk_size=chunk_size, chunk_overlap=chunk_overlap)


def train_on_documents(
    ctx: ExecutionContext,
    documents,
    runtime: RuntimePorts | None = None,
) -> IngestResult:
    active_runtime = runtime or default_runtime()
    return ingest_documents(ctx, documents, active_runtime)


def train_pdf_directory(ctx: ExecutionContext, source_dir: str, runtime: RuntimePorts | None = None) -> IngestResult:
    active_runtime = runtime or default_runtime()
    return ingest_pdf_directory(ctx, source_dir, active_runtime)


def train_single_doc(ctx: ExecutionContext, file_path: str, runtime: RuntimePorts | None = None) -> IngestResult:
    active_runtime = runtime or default_runtime()
    return ingest_single_pdf(ctx, file_path, active_runtime)


def train_on_csv(ctx: ExecutionContext, file_path: str, runtime: RuntimePorts | None = None, **kwargs) -> IngestResult:
    active_runtime = runtime or default_runtime()
    return ingest_csv(ctx, file_path, active_runtime, **kwargs)


__all__ = [
    "index_documents",
    "load_csv",
    "load_pdf_directory",
    "load_single_pdf",
    "train_on_documents",
    "train_pdf_directory",
    "train_single_doc",
    "train_on_csv",
]
