"""
LlamaIndex AG-UI Agent

Uses llama-index-protocols-ag-ui to expose a LlamaIndex workflow as an
AG-UI compatible FastAPI router. The router handles all four demo
scenarios (agentic-chat, tool-rendering, hitl, gen-ui-tool-based) through
a single endpoint since LlamaIndex's get_ag_ui_workflow_router builds
the full AG-UI protocol surface automatically.

NOTE: Uses FixedAGUIChatWorkflow from hitl_in_chat_agent to fix three
upstream library bugs (duplicate tool-call rendering, missing
parent_message_id, and incorrect tool-result message roles). See
hitl_in_chat_agent.py module docstring for details.
"""

import json
import os
from typing import Annotated

from llama_index.llms.openai import OpenAI
from llama_index.protocols.ag_ui.router import get_ag_ui_workflow_router

from agents.hitl_in_chat_agent import FixedAGUIChatWorkflow

# Import shared tool implementations
from tools import (
    get_weather_impl,
    query_data_impl,
    manage_sales_todos_impl,
    get_sales_todos_impl,
    schedule_meeting_impl,
    search_flights_impl,
    build_a2ui_operations_from_tool_call,
)


# --- Frontend tools (executed client-side, agent just returns a confirmation) ---

def change_background(
    background: Annotated[str, "CSS background value. Prefer gradients."],
) -> str:
    """Change the background color/gradient of the chat area."""
    return f"Background changed to {background}"


def generate_haiku(
    japanese: Annotated[list[str], "3 lines of haiku in Japanese"],
    english: Annotated[list[str], "3 lines of haiku translated to English"],
    image_name: Annotated[str, "One relevant image name from the valid set"],
    gradient: Annotated[str, "CSS Gradient color for the background"],
) -> str:
    """Generate a haiku with Japanese text, English translation, and a background image."""
    return "Haiku generated!"


def generate_task_steps(
    steps: Annotated[
        list[dict],
        "Array of step objects with 'description' (string) and 'status' ('enabled' or 'disabled')"
    ],
) -> str:
    """Generate a list of task steps for the user to review and approve."""
    return f"Generated {len(steps)} steps for review"


def book_call(
    topic: Annotated[str, "What the call is about (e.g. 'Intro with sales')"],
    attendee: Annotated[str, "Who the call is with (e.g. 'Alice from Sales')"],
) -> str:
    """Ask the user to pick a time slot for a call. The picker UI presents fixed candidate slots; the user's choice is returned to the agent."""
    return f"Booking call about {topic} with {attendee}"


def show_card(
    title: Annotated[str, "Short heading for the card."],
    body: Annotated[str, "Body text for the card."],
) -> str:
    """Display a titled card with a short body of text. Rendered on the frontend via useComponent."""
    return f"Displayed card: {title}"


# --- Backend tools (executed server-side, using shared implementations) ---

# @region[weather-tool-backend]
async def get_weather(
    location: Annotated[str, "The location to get the weather for."],
) -> str:
    """Get the weather for a given location. Returns temperature, conditions, humidity, wind speed, and feels-like temperature."""
    return json.dumps(get_weather_impl(location))
# @endregion[weather-tool-backend]


async def query_data(
    query: Annotated[str, "Natural language query for financial data."],
) -> str:
    """Query financial database for chart data. Always call before showing a chart or graph."""
    return json.dumps(query_data_impl(query))


async def manage_sales_todos(
    todos: Annotated[list[dict], "Complete list of sales todos to replace the current list."],
) -> str:
    """Manage the sales pipeline by replacing the entire list of todos."""
    result = manage_sales_todos_impl(todos)
    return json.dumps({"status": "updated", "count": len(result), "todos": [dict(t) for t in result]})


async def get_sales_todos_tool() -> str:
    """Get the current sales pipeline todos."""
    return json.dumps(get_sales_todos_impl(None))


async def schedule_meeting(
    reason: Annotated[str, "Reason for the meeting."],
) -> str:
    """Schedule a meeting with the user. Requires human approval."""
    return json.dumps(schedule_meeting_impl(reason))


async def search_flights(
    flights: Annotated[list[dict], "List of flight objects to search and display as rich cards. Return exactly 2 flights."],
) -> str:
    """Search for flights and display the results as rich A2UI cards.

    Each flight must have: airline, airlineLogo, flightNumber, origin, destination,
    date, departureTime, arrivalTime, duration, status, statusColor, price, currency.
    """
    result = search_flights_impl(flights)
    return json.dumps(result)


async def generate_a2ui(
    context: Annotated[str, "Conversation context to generate UI from."],
) -> str:
    """Generate dynamic A2UI components based on the conversation.

    A secondary LLM designs the UI schema and data. The result is
    returned as an a2ui_operations container for the middleware to detect.
    """
    from openai import OpenAI

    client = OpenAI()
    tool_schema = {
        "type": "function",
        "function": {
            "name": "render_a2ui",
            "description": "Render a dynamic A2UI v0.9 surface.",
            "parameters": {
                "type": "object",
                "properties": {
                    "surfaceId": {"type": "string"},
                    "catalogId": {"type": "string"},
                    "components": {"type": "array", "items": {"type": "object"}},
                    "data": {"type": "object"},
                },
                "required": ["surfaceId", "catalogId", "components"],
            },
        },
    }

    response = client.chat.completions.create(
        model="gpt-4.1",
        messages=[
            {"role": "system", "content": context or "Generate a useful dashboard UI."},
            {"role": "user", "content": "Generate a dynamic A2UI dashboard based on the conversation."},
        ],
        tools=[tool_schema],
        tool_choice={"type": "function", "function": {"name": "render_a2ui"}},
    )

    if not response.choices[0].message.tool_calls:
        return json.dumps({"error": "LLM did not call render_a2ui"})

    tool_call = response.choices[0].message.tool_calls[0]
    args = json.loads(tool_call.function.arguments)
    result = build_a2ui_operations_from_tool_call(args)
    return json.dumps(result)


_openai_kwargs = {}
if os.environ.get("OPENAI_BASE_URL"):
    _openai_kwargs["api_base"] = os.environ["OPENAI_BASE_URL"]

_AGENT_SYSTEM_PROMPT = (
    "You are a polished, professional demo assistant for CopilotKit. "
    "Keep responses brief and clear -- 1 to 2 sentences max.\n\n"
    "You can:\n"
    "- Chat naturally with the user\n"
    "- Change the UI background when asked (via frontend tool)\n"
    "- Query data and render charts (via query_data tool)\n"
    "- Get weather information (via get_weather tool)\n"
    "- Schedule meetings with the user (via schedule_meeting tool)\n"
    "- Manage sales pipeline todos (via manage_sales_todos / get_sales_todos tools)\n"
    "- Search flights and display rich A2UI cards (via search_flights tool)\n"
    "- Generate dynamic A2UI dashboards from conversation context (via generate_a2ui tool)\n"
    "- Generate step-by-step plans for user review (human-in-the-loop)\n"
    "- Book calls with people (via book_call frontend tool)\n"
    "- Show titled cards with a body of text (via show_card frontend tool)\n"
    "When asked about weather, always use the get_weather tool. "
    "When asked about financial data or charts, use query_data first. "
    "When asked to book a call, use the book_call tool with topic and name."
)


async def _agent_workflow_factory():
    wf = FixedAGUIChatWorkflow(
        llm=OpenAI(model="gpt-4.1", **_openai_kwargs),
        frontend_tools=[change_background, generate_haiku, generate_task_steps, book_call, show_card, get_weather],
        backend_tools=[query_data, manage_sales_todos, get_sales_todos_tool, schedule_meeting, search_flights, generate_a2ui],
        system_prompt=_AGENT_SYSTEM_PROMPT,
        initial_state={
            "todos": [],
        },
    )
    # Tools that use useRenderTool on the frontend — emit
    # TOOL_CALL_RESULT so the render transitions to "complete".
    wf.render_only_tool_names = {"get_weather"}
    return wf


agent_router = get_ag_ui_workflow_router(
    workflow_factory=_agent_workflow_factory,
)
