import os
import json

import gradio as gr
from groq import Groq

from memory import add_event, search_events, get_upcoming_events


def build_system_prompt():
    """Assemble the system prompt with upcoming events context."""
    upcoming = get_upcoming_events(days=7)

    base_prompt = (
        "You are a helpful calendar memory agent. You remember the user's "
        "important dates, deadlines, meetings, and events. You can save new "
        "events and retrieve existing ones.\n\n"
        "When the user mentions a date, deadline, or event, use the save_event "
        "tool to store it. When they ask about their schedule, use the "
        "get_events tool to look it up.\n\n"
        "Always confirm what you saved or found. Be concise and helpful."
    )

    if upcoming:
        events_text = "\n".join(
            f"- {e['date']}: {e['description']} ({e['category']})"
            for e in upcoming
        )
        base_prompt += (
            f"\n\nUpcoming events (next 7 days):\n{events_text}"
        )
    else:
        base_prompt += "\n\nNo upcoming events in the next 7 days."

    return base_prompt

def define_tools():
    """Define the tools available to the model."""
    tools = [
        {
            "type": "function",
            "function": {
                "name": "save_event",
                "description": "Save a new calendar event, deadline, or reminder to memory.",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "date": {
                            "type": "string",
                            "description": "The date of the event in YYYY-MM-DD format"
                        },
                        "description": {
                            "type": "string",
                            "description": "A brief description of the event"
                        },
                        "category": {
                            "type": "string",
                            "enum": ["deadline", "meeting", "reminder", "event"],
                            "description": "The category of the event"
                        }
                    },
                    "required": ["date", "description"]
                }
            }
        },
        {
            "type": "function",
            "function": {
                "name": "get_events",
                "description": "Retrieve calendar events from memory by date range or keyword.",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "start_date": {
                            "type": "string",
                            "description": "Start date in YYYY-MM-DD format (optional)"
                        },
                        "end_date": {
                            "type": "string",
                            "description": "End date in YYYY-MM-DD format (optional)"
                        },
                        "keyword": {
                            "type": "string",
                            "description": "Keyword to search in event descriptions (optional)"
                        }
                    }
                }
            }
        }
    ]
    return tools

def handle_tool_call(tool_call):
    """Execute a client-side tool call and return the result."""
    name = tool_call.function.name
    args = json.loads(tool_call.function.arguments)

    if name == "save_event":
        event = add_event(
            date=args["date"],
            description=args["description"],
            category=args.get("category", "event")
        )
        return json.dumps({
            "status": "saved",
            "event": event
        })

    elif name == "get_events":
        events = search_events(
            start_date=args.get("start_date"),
            end_date=args.get("end_date"),
            keyword=args.get("keyword")
        )
        return json.dumps({
            "status": "found",
            "count": len(events),
            "events": events
        })

    return json.dumps({"error": f"Unknown tool: {name}"})

def chat_with_agent(message, history):
    """Process a user message through the Groq agent with tool calling."""
    client = Groq(api_key=os.getenv("GROQ_API_KEY"))
    tools = define_tools()
    system_prompt = build_system_prompt()

    # Build the messages array with system prompt and user message
    messages = [
        {"role": "system", "content": system_prompt},
        {"role": "user", "content": message}
    ]

    while True:
        response = client.chat.completions.create(
            model="llama-3.3-70b-versatile",
            messages=messages,
            tools=tools,
        )

        choice = response.choices[0]

        # If the model wants to call tools, execute them
        if choice.finish_reason == "tool_calls":
            # Append the assistant's message (with tool_calls) to history
            messages.append(choice.message)

            # Execute each tool call and append results
            for tc in choice.message.tool_calls:
                result = handle_tool_call(tc)
                messages.append({
                    "role": "tool",
                    "tool_call_id": tc.id,
                    "name": tc.function.name,
                    "content": result,
                })
        else:
            # No tool calls — return the final text response
            return choice.message.content

demo = gr.ChatInterface(
    fn=chat_with_agent,
    title="Calendar Memory Agent",
    description="I remember your dates, deadlines, and events. Tell me about upcoming events or ask what's on your schedule!",
    examples=[
        "I have a dentist appointment on 2026-06-15",
        "What events do I have coming up?",
        "My project deadline is 2026-06-20",
    ],
)

if __name__ == "__main__":
    demo.launch()