# pyright: reportUnusedFunction=false

from __future__ import annotations

import os
import importlib
from ketju.core.logging import configure_example_logging
from ketju.rag import BaseRAG
from pathlib import Path
from pydantic import BaseModel, Field, ConfigDict
from typing import Any, List, Union

try:
    from pydantic_ai import RunContext
except ModuleNotFoundError:
    raise ModuleNotFoundError(
        "pydantic_ai is required for CLI/AG-UI agents.\n"
        "Install with one of:\n"
        "  uv sync --extra agent\n"
        "  uv sync --extra examples\n"
        "  pip install 'ketju[agent]'"
    ) from None

def is_pdf_file(file_path: Path | str) -> bool:
    try:
        fp = Path(file_path)
        return fp.exists() and fp.is_file() and fp.suffix == ".pdf"
    except Exception:
        return False


default_system_prompt = """
    You are a helpful assistant that can answer questions about the document.
    If you do not find answers in ingested documents, don't hallucinate.
    You can ingest documents into the knowledge database by providing the path to the document.
"""

class Deps(BaseModel):
    model_config = ConfigDict(arbitrary_types_allowed=True)

    rag_engine: BaseRAG
    ingested_files: List[str] = Field(default_factory=list)



def create_agent(
    model: str = "gpt-4o-mini",
    system_prompt: str = default_system_prompt,
    deps_type: type[Deps] = Deps,
) -> Any:
    try:
        pydantic_ai = importlib.import_module("pydantic_ai")
    except ModuleNotFoundError:
        raise ModuleNotFoundError(
            "pydantic_ai is required for CLI/AG-UI agents.\n"
            "Install with one of:\n"
            "  uv sync --extra agent\n"
            "  uv sync --extra examples\n"
            "  pip install 'ketju[agent]'"
        ) from None

    agent = pydantic_ai.Agent(
        model=model,
        system_prompt=system_prompt,
        deps_type=deps_type,
    )

    @agent.tool_plain
    def validate_file(file_path: str) -> Union[Path, bool]:
        """Check that the file exists and that it is a PDF file."""
        return is_pdf_file(file_path)

    @agent.tool
    def has_been_ingested(ctx: RunContext[Deps], file_path: str) -> bool:
        """Check that the file has been ingested into the RAG."""
        return file_path in ctx.deps.ingested_files

    @agent.tool
    def ingest_file(ctx: RunContext[Deps], file_path: str) -> str:
        """Ingest the document into the RAG."""

        if not is_pdf_file(file_path):
            return "File is not a valid PDF file"
        if has_been_ingested(ctx, file_path):
            return "File has already been ingested"

        try:
            ctx.deps.rag_engine.ingest(files=[Path(file_path)], max_files=1)
            ctx.deps.ingested_files.append(file_path)
            return f"File {file_path} has been ingested"
        except Exception as e:
            return f"Could not ingest file {file_path}: {e}"

    @agent.tool
    def consult_documents(ctx: RunContext[Deps], question: str):
        """Consult the knowledge database for an answer."""
        return ctx.deps.rag_engine.query(question)  # type: ignore

    return agent


def run_cli_agent(
    rag_engine: BaseRAG,
    model: str = "gpt-4o-mini",
    system_prompt: str = default_system_prompt,
):
    configure_example_logging(
        log_level=os.getenv("KETJU_LOG_LEVEL", "WARNING"),
        quiet_third_party=os.getenv("KETJU_QUIET", "1") != "0",
    )
    if os.getenv("KETJU_INSTRUMENT", "0") == "1":
        # Optional: enable Logfire / OpenTelemetry instrumentation.
        # Install with: uv sync --extra observability
        from ketju.core.instrumentation import instrument_ketju

        instrument_ketju(service_name="ketju-agent")
    agent = create_agent(model, system_prompt, deps_type=Deps)
    agent.to_cli_sync(deps=Deps(rag_engine=rag_engine, ingested_files=[]), prog_name="ketju-agent")


def agui_agent(
    rag_engine: BaseRAG,
    model: str = "gpt-4o-mini",
    system_prompt: str = default_system_prompt,
):
    agent = create_agent(model, system_prompt, deps_type=Deps)
    return agent.to_ag_ui(deps=Deps(rag_engine=rag_engine, ingested_files=[]))

