# backend/src/agent/langgraph_setup.py
# #region AgentChat.LangGraph.Setup [C:4] [TYPE Module] [SEMANTICS agent-chat,langgraph,agent]
# @defgroup AgentChat LangGraph agent setup: create_react_agent with PostgresSaver.
# @PRE LLM provider configured via backend API /api/agent/llm-config.
# @POST Compiled StateGraph ready for astream_events().
# @SIDE_EFFECT Initializes checkpointer and message history tables on first call.
# @RELATION DEPENDS_ON -> [AgentChat.Tools]
# @RATIONALE LangGraph create_react_agent provides built-in tool calling + checkpointing + interrupt/resume.
# @REJECTED Using only environment variables for LLM config was rejected — FastAPI API-based config allows runtime switching without restart.
#            RunnableWithMessageHistory wrapper is NOT used — PostgresSaver handles history natively.

# ── Monkey-patch: OpenAI SDK for Pydantic BaseModel classes ──
# LangChain BaseTool objects carry an ``args_schema`` field that is a Pydantic
# BaseModel *class* reference (not an instance). When the OpenAI SDK recursively
# transforms the request body, it hits ``isinstance(data, pydantic.BaseModel)``
# which is True for both instances AND classes. It then calls ``model_dump()``
# on the class, which fails with:
#
#   PydanticSerializationError: Unable to serialize unknown type: ModelMetaclass
#
# The fix: skip model_dump for classes, only dump instances.
import inspect as _inspect
import os

import httpx
from langchain_openai import ChatOpenAI
from langgraph.checkpoint.memory import InMemorySaver
from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
from langgraph.prebuilt import create_react_agent
import openai._utils._transform as _openai_transform
import psycopg
from psycopg.rows import dict_row
import pydantic as _pydantic
import pydantic_core as _pydantic_core

from src.agent._config import AGENT_CONFIRM_TOOLS, AGENT_INTERRUPT_BEFORE as _INTERRUPT_BEFORE, FASTAPI_URL
from src.agent._llm_params import chat_openai_kwargs
from src.core.logger import logger

_original_transform = _openai_transform._async_transform_recursive

async def _patched_transform(data, *, annotation, inner_type=None):
    if isinstance(data, _pydantic.BaseModel):
        if _inspect.isclass(data):
            # BaseModel CLASS (not instance) — skip model_dump
            return data
        # BaseModel INSTANCE — intercept PydanticSerializationError
        try:
            return await _original_transform(data, annotation=annotation, inner_type=inner_type)
        except _pydantic_core.PydanticSerializationError:
            print(f"[PATCH] Caught PydanticSerializationError on {type(data).__name__}")
            # Fallback: dump with exclude of type-ref fields
            serializable = {}
            for field_name in data.model_fields_set:
                val = getattr(data, field_name)
                if isinstance(val, type) and issubclass(val, _pydantic.BaseModel):
                    serializable[field_name] = val.model_json_schema()
                else:
                    serializable[field_name] = val
            return serializable
    return await _original_transform(data, annotation=annotation, inner_type=inner_type)

_openai_transform._async_transform_recursive = _patched_transform

# ── Postgres checkpointer (FR-004/FR-012/FR-027) ──
_CHECKPOINTER: AsyncPostgresSaver | None = None
_CHECKPOINTER_INIT = False
_CHECKPOINTER_CONN = None


async def init_checkpointer() -> None:
    """Initialize the PostgreSQL checkpointer (FR-004/FR-012/FR-027).

    Called once at agent startup. Creates a persistent psycopg async connection
    and passes it to AsyncPostgresSaver. Runs .setup() to create checkpoint tables.
    Connection stays open for the lifetime of the agent process.
    """
    global _CHECKPOINTER, _CHECKPOINTER_INIT, _CHECKPOINTER_CONN
    if _CHECKPOINTER_INIT:
        return
    db_url = os.getenv("DATABASE_URL")
    # Convert SQLAlchemy-style URL to psycopg format
    pg_url = db_url.replace("postgresql+psycopg2://", "postgres://").replace("postgresql://", "postgres://")
    _CHECKPOINTER_CONN = await psycopg.AsyncConnection.connect(pg_url, autocommit=True, row_factory=dict_row)
    _CHECKPOINTER = AsyncPostgresSaver(_CHECKPOINTER_CONN)
    await _CHECKPOINTER.setup()
    _CHECKPOINTER_INIT = True

# ── LLM config (no cache — fetched on each create_agent call) ──
_llm_config: dict | None = None


def configure_from_api(llm_config: dict) -> None:
    """Update LLM config from FastAPI response. Called at startup."""
    global _llm_config
    _llm_config = llm_config


async def _fetch_llm_config() -> dict | None:
    """Fetch LLM config from FastAPI.

    Called on every create_agent() to pick up Admin UI changes immediately.
    Falls back to cached config if fetch fails.
    """
    global _llm_config
    try:
        fastapi_url = FASTAPI_URL
        async with httpx.AsyncClient(timeout=5) as client:
            resp = await client.get(f"{fastapi_url}/api/agent/llm-config")
            if resp.status_code == 200:
                config = resp.json()
                if config.get("configured"):
                    _llm_config = config
                    return config
    except Exception as e:
        logger.explore("Failed to fetch LLM config from FastAPI", error=str(e),
                       extra={"src": "AgentChat.LangGraph.Setup"})
    return _llm_config


def _interrupt_before_from_env() -> list[str]:
    """Return LangGraph node names that must pause for HITL confirmation."""
    if AGENT_CONFIRM_TOOLS:
        return ["tools"]
    raw = _INTERRUPT_BEFORE
    if not raw:
        return []
    return [name.strip() for name in raw.split(",") if name.strip()]


async def create_agent(
    tools: list,
    env_id: str | None = None,
    interrupt_before: list[str] | None = None,
):
    """Create the LangGraph agent with PostgreSQL checkpointer and message history.

    LLM configuration source:
        llm_config from FastAPI /api/agent/llm-config (fetched on every call).
        If backend has no configured provider, agent raises an error.

    Returns a compiled StateGraph ready for astream_events().
    interrupt_before is set from AGENT_CONFIRM_TOOLS (or AGENT_INTERRUPT_BEFORE env var)
    to enable HITL guardrails — when pending tools are detected, the graph pauses before
    executing them and yields confirm_required metadata to the frontend.
    Checkpointer is AsyncPostgresSaver (survives container restarts).
    """
    # Fetch fresh LLM config from FastAPI on every call
    config = await _fetch_llm_config()

    if config and config.get("configured"):
        api_key = config["api_key"]
        base_url = config.get("base_url")
        model = config.get("default_model")
        config_source = "FastAPI"
    else:
        raise RuntimeError(
            "No LLM provider configured in backend. "
            "Configure one via Settings → AI Providers in the web UI."
        )

    logger.reason(
        "Creating LangGraph agent",
        payload={"model": model, "config_source": config_source, "tools_count": len(tools), "env_id": env_id},
        extra={"src": "AgentChat.LangGraph.Setup"},
    )

    llm = ChatOpenAI(**chat_openai_kwargs(
        model=model,
        base_url=base_url,
        api_key=api_key,
        max_tokens=2048,
    ))

    # System prompt — env_id injected deterministically, not in user message
    prompt = (
        "You are a Superset Tools assistant. You have access to tools for searching "
        "dashboards, managing maintenance, running migrations and backups, "
        "executing SQL and exploring databases, auditing permissions, "
        "managing Git operations (branch/commit/deploy), running LLM validation "
        "and documentation, creating and copying dashboards and datasets, "
        "and checking system health, environments, and task status. "
        "You handle all intent detection — multi-intent queries, negations (\"don't run\"), "
        "synonyms (\"панели\" = \"дашборды\"), and typos are your responsibility. "
        "Call the right tool(s) for the job. If data is already provided in context, "
        "use it directly rather than calling redundant tools. "
        "For maintenance requests, use the RUNTIME CONTEXT current datetime when the user says "
        "\"start\", \"run\", \"now\", \"запусти\", or \"сейчас\" without an explicit start time. "
        "Convert user durations into end_time. Do not ask for ISO datetime in that case. "
        "If a user asks for dashboard maintenance, resolve the dashboard from provided context or tools, "
        "then infer affected tables when possible; ask for table names only after resolution fails."
    )
    if env_id:
        prompt += f"\n\nCurrent environment: '{env_id}'. When calling tools that accept env_id, use this value."

    # Checkpointer — AsyncPostgresSaver. Fallback to InMemory if Postgres init failed
    if _CHECKPOINTER is not None:
        checkpointer = _CHECKPOINTER
    else:
        checkpointer = InMemorySaver()
        logger.explore(
            "Postgres checkpointer unavailable, falling back to InMemorySaver",
            error="_CHECKPOINTER is None — checkpoints will be lost on restart",
            extra={"src": "AgentChat.LangGraph.Setup"},
        )

    graph = create_react_agent(
        model=llm,
        tools=tools,
        prompt=prompt,
        version="v2",
        checkpointer=checkpointer,
        interrupt_before=_interrupt_before_from_env() if interrupt_before is None else interrupt_before,
    )

    logger.reflect(
        "LangGraph agent created",
        payload={"model": model, "checkpointer_type": type(checkpointer).__name__, "tools_count": len(tools)},
        extra={"src": "AgentChat.LangGraph.Setup"},
    )
    return graph
# #endregion AgentChat.LangGraph.Setup