"""PydanticAI Agent Definitions for DataBridge AI.

Provides type-safe, multi-step reasoning agents:

- **WorkflowPlannerAgent** — Upgrades single-shot PlannerAgent to an
  iterative planner with tool calls for validation.
- **SQLAgent** — Combines Vanna RAG + DuckDB for intelligent SQL assistance.

Both agents fall back gracefully when ``pydantic-ai`` is not installed.
"""
from __future__ import annotations

import logging
import os
from typing import Any, Dict, List, Optional

from pydantic import BaseModel, ConfigDict, Field

logger = logging.getLogger(__name__)

# ------------------------------------------------------------------
# Conditional PydanticAI import
# ------------------------------------------------------------------

try:
    from pydantic_ai import Agent

    PYDANTIC_AI_AVAILABLE = True
except ImportError:
    PYDANTIC_AI_AVAILABLE = False
    Agent = None  # type: ignore[misc,assignment]


def _get_sql_agent_duckdb_connection() -> Any:
    """Return the shared DuckDB connection used by DataBridge tools."""
    try:
        try:
            from src.turbo import get_duckdb
        except ImportError:
            from turbo import get_duckdb  # type: ignore[no-redef]
        return get_duckdb()
    except Exception as exc:
        raise RuntimeError(f"DuckDB session unavailable: {exc}") from exc


# ------------------------------------------------------------------
# Structured output models
# ------------------------------------------------------------------


class PlanStepOutput(BaseModel):
    """A single step in a workflow plan."""

    id: str = ""
    name: str = ""
    description: str = ""
    agent_type: str = ""
    capability: str = ""
    input_mapping: Dict[str, str] = Field(default_factory=dict)
    output_key: str = ""
    depends_on: List[str] = Field(default_factory=list)
    reasoning: str = ""


class PlanOutput(BaseModel):
    """Structured output for workflow planning."""

    name: str = ""
    description: str = ""
    steps: List[PlanStepOutput] = Field(default_factory=list)
    confidence_score: float = 0.0
    estimated_duration: str = ""
    warnings: List[str] = Field(default_factory=list)
    alternatives: List[str] = Field(default_factory=list)


class SQLOutput(BaseModel):
    """Structured output for SQL generation."""

    sql: str = ""
    explanation: str = ""
    tables_used: List[str] = Field(default_factory=list)
    confidence: float = 0.0


# ------------------------------------------------------------------
# Agent context (dependencies injected at run-time)
# ------------------------------------------------------------------


class PlannerDeps(BaseModel):
    """Dependencies injected into the planner agent at run-time."""

    model_config = ConfigDict(arbitrary_types_allowed=True)

    available_agents: List[Dict[str, Any]] = Field(default_factory=list)


class SQLDeps(BaseModel):
    """Dependencies injected into the SQL agent at run-time."""

    model_config = ConfigDict(arbitrary_types_allowed=True)

    vanna_available: bool = False


# ------------------------------------------------------------------
# System prompts
# ------------------------------------------------------------------

_PLANNER_SYSTEM_PROMPT = """\
You are a workflow planning agent for DataBridge AI, a data engineering platform.

Your job is to analyse user requests and create optimal workflow plans by selecting
the right agents and ordering steps based on dependencies.

## DataBridge Workflow Knowledge — E2E Assessment + BLCE Pipeline

When a user asks about creating a data model, data warehouse, or data mapping:

### Phase 1: E2E Assessment (Initial Data Model & Warehouse)
The E2E Assessment Pipeline creates the initial data model and warehouse from raw source
system tables. It runs 22 phases: schema discovery → relationship discovery → AI classification
→ dimension/fact detection → hierarchy creation → dbt generation → DDL deployment → quality
validation → observability setup.

After Phase 1 completes, ALWAYS suggest Phase 2 in the alternatives:
"Once the initial warehouse is built, you can share existing reports, SQL queries, formulas,
and system requirements so we can refine the model to match your actual business needs using
the BLCE Engine."

### Phase 2: BLCE Refinement (Business Logic → Data Marts)
The BLCE Engine (72 tools, 21 phases) takes the E2E output warehouse and enriches it:
1. Gather context: existing reports, SQL, formulas, system requirements, knowledge base docs
2. Extract business logic: parse SQL CASE statements, calculations, business rules
3. Identify configurations, cross-references, and hierarchies the user can build
4. Generate data marts on top of the BLCE output warehouse
5. Iterate: users can keep refining with more hierarchies and mart layers

### GraphRAG Context Persistence
After E2E builds the warehouse, DDLs, schema, and model are stored in GraphRAG. This means:
- Future suggestions reference what exists in the current warehouse
- Wright pipeline generation for new hierarchies knows existing tables/columns
- BLCE cross-references and configurations are indexed for retrieval
- The AI can safely suggest edits and changes knowing the full warehouse state

### Workflow Execution & Progress Tracking
When the user wants to execute a plan, the platform can:
1. Generate a Python wrapper script that orchestrates the workflow steps
2. Connect back to the UI to show real-time progress (data flowing, tools being used)
3. Support custom workflows — E2E, BLCE, or hybrid with user-requested edits/enhancements
4. The workflow page shows live phase status updates (like the E2E Assessment runner)

When suggesting plan execution, mention that the user can:
- Execute the plan as-is through the workflow runner
- Customize individual steps before execution
- Monitor progress live in the Workflow page (Flow/Flowchart tabs)

### Connection Handling
If the user hasn't specified connection details, help them configure the connection step by
step like a normal LLM assistant would — ask about database type, host, credentials, schema.
Don't just warn about missing connections; guide them through setup.

## Planning Guidelines
1. Identify the user's goal.
2. Select the minimum set of agents needed.
3. Order steps respecting data-flow dependencies.
4. Steps without inter-dependencies can run in parallel (leave depends_on empty).
5. Provide reasoning for each step so humans can review the plan.
6. Flag ambiguities or missing information as warnings.
7. Suggest alternative approaches when relevant.
8. For data warehouse requests, plan E2E Assessment and mention BLCE as next step.
9. After any warehouse change, note that DDLs/schema persist in GraphRAG.

Use the tools available to you to validate your plan before returning it.
"""

_SQL_SYSTEM_PROMPT = """\
You are a SQL expert for DataBridge AI.

Given a natural-language question, generate a correct SQL query.
Use the tools available to inspect table schemas and optionally run test queries.
If a Vanna RAG suggestion is available, incorporate it but verify correctness.
Return a clear explanation alongside the SQL.
"""


# ------------------------------------------------------------------
# Agent factories (only created when PydanticAI is available)
# ------------------------------------------------------------------


def _create_planner_agent(
    available_agents: List[Dict[str, Any]],
) -> Any:
    """Create a PydanticAI WorkflowPlannerAgent.

    Returns the Agent instance, or ``None`` if pydantic-ai is missing.
    """
    if not PYDANTIC_AI_AVAILABLE or Agent is None:
        return None

    model = os.getenv("PYDANTIC_AI_MODEL", "anthropic:claude-sonnet-4-20250514")

    agent = Agent(
        model,
        output_type=PlanOutput,
        system_prompt=_PLANNER_SYSTEM_PROMPT,
        deps_type=PlannerDeps,
    )

    # ------ Tools the agent can call during planning ------

    @agent.tool_plain
    def list_available_agents() -> List[Dict[str, Any]]:
        """List all available agents and their capabilities."""
        return available_agents

    @agent.tool_plain
    def validate_step_dependency(step_id: str, depends_on: List[str], all_step_ids: List[str]) -> Dict[str, Any]:
        """Validate that step dependencies reference existing steps.

        Args:
            step_id: The step being validated.
            depends_on: List of step IDs this step depends on.
            all_step_ids: All known step IDs in the plan so far.
        """
        missing = [d for d in depends_on if d not in all_step_ids]
        has_cycle = step_id in depends_on
        return {
            "valid": len(missing) == 0 and not has_cycle,
            "missing_deps": missing,
            "has_cycle": has_cycle,
        }

    @agent.tool_plain
    def check_agent_capability(agent_name: str, capability: str) -> Dict[str, Any]:
        """Check if an agent has a specific capability.

        Args:
            agent_name: Name of the agent to check.
            capability: Capability to look for.
        """
        for ag in available_agents:
            if ag.get("name") == agent_name:
                caps = ag.get("capabilities", [])
                return {
                    "found": True,
                    "has_capability": capability in caps,
                    "all_capabilities": caps,
                }
        return {"found": False, "has_capability": False, "all_capabilities": []}

    return agent


def _create_sql_agent() -> Any:
    """Create a PydanticAI SQLAgent.

    Returns the Agent instance, or ``None`` if pydantic-ai is missing.
    """
    if not PYDANTIC_AI_AVAILABLE or Agent is None:
        return None

    model = os.getenv("PYDANTIC_AI_MODEL", "anthropic:claude-sonnet-4-20250514")

    agent = Agent(
        model,
        output_type=SQLOutput,
        system_prompt=_SQL_SYSTEM_PROMPT,
        deps_type=SQLDeps,
    )

    @agent.tool_plain
    def get_table_schema(table_name: str) -> Dict[str, Any]:
        """Get schema for a registered DuckDB table.

        Args:
            table_name: Name of the table to inspect.
        """
        try:
            con = _get_sql_agent_duckdb_connection()
            result = con.execute(f"DESCRIBE {table_name}").fetchall()
            columns = [{"name": row[0], "type": row[1]} for row in result]
            return {"table": table_name, "columns": columns}
        except Exception as exc:
            return {"table": table_name, "error": str(exc)}

    @agent.tool_plain
    def run_test_query(sql: str) -> Dict[str, Any]:
        """Execute a test query against DuckDB and return up to 5 rows.

        Args:
            sql: SQL statement to execute.
        """
        try:
            con = _get_sql_agent_duckdb_connection()
            result = con.execute(sql).fetchmany(5)
            columns = [desc[0] for desc in con.description or []]
            return {"columns": columns, "rows": [list(r) for r in result], "success": True}
        except Exception as exc:
            return {"error": str(exc), "success": False}

    @agent.tool_plain
    def get_vanna_suggestion(question: str) -> Dict[str, Any]:
        """Get Vanna's RAG-based SQL suggestion (if trained).

        Args:
            question: The natural-language question.
        """
        try:
            from src.vanna_client import get_vanna_client

            vanna = get_vanna_client()
            if not vanna.is_trained():
                return {"available": False, "reason": "Vanna not trained yet"}
            result = vanna.ask(question)
            if result:
                return {"available": True, "suggestion": result}
            return {"available": False, "reason": "No SQL generated"}
        except Exception as exc:
            return {"available": False, "reason": str(exc)}

    return agent


# ------------------------------------------------------------------
# High-level runner functions (called from planner_agent.py / server.py)
# ------------------------------------------------------------------


async def run_planner_agent(
    user_request: str,
    available_agents: List[Dict[str, Any]],
    context: Optional[Dict[str, Any]] = None,
) -> Optional[PlanOutput]:
    """Run the PydanticAI WorkflowPlannerAgent.

    Args:
        user_request: Natural-language request from the user.
        available_agents: List of agent info dicts.
        context: Optional additional context dict.

    Returns:
        PlanOutput if successful, None if PydanticAI unavailable or failed.
    """
    agent = _create_planner_agent(available_agents)
    if agent is None:
        return None

    prompt = user_request
    if context:
        import json

        prompt += f"\n\nAdditional context:\n```json\n{json.dumps(context, indent=2)}\n```"

    try:
        deps = PlannerDeps(available_agents=available_agents)
        result = await agent.run(prompt, deps=deps)
        return result.output
    except Exception as exc:
        logger.error("PydanticAI planner failed: %s", exc)
        return None


async def run_sql_agent(
    question: str,
    tables: Optional[List[str]] = None,
) -> Optional[SQLOutput]:
    """Run the PydanticAI SQLAgent.

    Args:
        question: Natural-language question about the data.
        tables: Optional list of table names for context.

    Returns:
        SQLOutput if successful, None otherwise.
    """
    agent = _create_sql_agent()
    if agent is None:
        return None

    prompt = question
    if tables:
        prompt += f"\n\nAvailable tables: {', '.join(tables)}"

    try:
        from src.vanna_client import VANNA_AVAILABLE as _va

        deps = SQLDeps(vanna_available=_va)
    except ImportError:
        deps = SQLDeps(vanna_available=False)

    try:
        result = await agent.run(prompt, deps=deps)
        return result.output
    except Exception as exc:
        logger.error("PydanticAI SQL agent failed: %s", exc)
        return None
