from __future__ import annotations

# pyright: reportMissingTypeStubs=false

import argparse
import os

from ketju.config import DATA_DIR
from ketju.examples._cli import import_agent_module, missing_optional_deps_message, parse_args


def _database_connection_error_message(database_url: str, error: Exception) -> str:
    details = str(error)
    base = (
        "Could not connect to Postgres for the pgvector example.\n"
        f"Database URL: {database_url}\n\n"
        "Check:\n"
        "  1) Postgres is running and reachable\n"
        "     - pg_isready -h 127.0.0.1 -p 5432\n"
        "  2) Credentials/database are correct\n"
        "     - psql \"<DATABASE_URL>\" -c \"select 1\"\n"
        "  3) pgvector extension exists in the target database\n"
        "     - psql \"<DATABASE_URL>\" -c \"create extension if not exists vector;\""
    )
    if "pg_filenode.map" in details:
        base += (
            "\n\n"
            "Your Postgres instance reports a storage-level file error "
            "(pg_filenode.map). This is a database server/data-dir issue, "
            "not an application dependency issue."
        )
    return f"{base}\n\nOriginal error: {details}"


def main(argv: list[str] | None = None) -> None:
    def _configure_parser(parser: argparse.ArgumentParser) -> None:
        parser.add_argument(
            "--database-url",
            default=os.getenv("DATABASE_URL"),
            help="Postgres connection string (defaults to env DATABASE_URL).",
        )

    args = parse_args(
        description="Pgvector-backed RAG example.",
        argv=argv,
        configure=_configure_parser,
    )

    try:
        from llama_index.core import Settings
        from llama_index.embeddings.openai import OpenAIEmbedding
        from llama_index.llms.openai import OpenAI
        from ketju.rag.llama_index.variants.pgvector import build_pgvector_rag
        run_cli_agent = import_agent_module().run_cli_agent
    except ModuleNotFoundError as e:
        raise SystemExit(
            missing_optional_deps_message(
                extras="--extra rag --extra pgvector --extra agent",
                run_as_module="ketju.examples.pgvector_example",
                run_tail="--database-url <DB_URL>",
            )
        ) from e

    # Configure models (choose what you actually use)
    Settings.embed_model = OpenAIEmbedding(model="text-embedding-3-small")
    Settings.llm = OpenAI(model="gpt-4o-mini")

    # IMPORTANT: set embed_dim to match your embed model output dimension
    # text-embedding-3-small is 1536
    embed_dim = 1536

    if not args.database_url:
        raise SystemExit(
            "Missing database URL.\n"
            "Pass --database-url or set DATABASE_URL, e.g.\n"
            "  postgresql://postgres:password@127.0.0.1:5432/ketju_db"
        )

    try:
        rag = build_pgvector_rag(
            name="pgvector_demo",
            persist_dir=DATA_DIR / "pgvector_demo",  # still used for docstore/other local state
            table_name="ketju_pgvector_demo",
            embed_dim=embed_dim,
            database_url=args.database_url,
        )
    except Exception as e:
        raise SystemExit(_database_connection_error_message(args.database_url, e)) from e

    print("Starting the agent...")
    print(
        "You can ask questions about contents in pdf files, "
        "just paste the path to the file in the chat."
    )
    run_cli_agent(rag)


if __name__ == "__main__":
    main()