"""
create_agent.py — Create an AI Agent with Tools
================================================
Demonstrates creating an Azure AI Agent with:
  - Code Interpreter (run Python in a sandbox)
  - File Search (search over uploaded files)
  - Function Calling (call your own functions)

Prerequisites:
    pip install azure-ai-projects azure-identity python-dotenv

Usage:
    python create_agent.py
"""

import json
import os
import sys
import time
from dotenv import load_dotenv
from azure.identity import DefaultAzureCredential
from azure.ai.projects import AIProjectClient
from azure.ai.projects.models import (
    CodeInterpreterTool,
    FileSearchTool,
    FunctionTool,
    FunctionDefinition,
    ToolSet,
)

load_dotenv()


def get_env(key: str) -> str:
    value = os.getenv(key)
    if not value:
        print(f"ERROR: '{key}' is not set. See .env.example.")
        sys.exit(1)
    return value


def get_project_client() -> AIProjectClient:
    """Build an AIProjectClient from environment variables."""
    credential = DefaultAzureCredential()
    location = get_env("AZURE_LOCATION")
    conn_str = (
        f"{location}.api.azureml.ms;"
        f"{get_env('AZURE_SUBSCRIPTION_ID')};"
        f"{get_env('AZURE_RESOURCE_GROUP')};"
        f"{get_env('AZURE_AI_FOUNDRY_PROJECT_NAME')}"
    )
    return AIProjectClient.from_connection_string(credential=credential, conn_str=conn_str)


# ---------------------------------------------------------------------------
# Custom function the agent can call
# ---------------------------------------------------------------------------
def lookup_product_price(product_name: str) -> str:
    """Simulate a product price lookup."""
    prices = {
        "azure openai": "$0.005 per 1K tokens (GPT-4o-mini)",
        "azure ai search": "$75.61/month (Basic tier)",
        "azure ai foundry": "No additional charge — pay for underlying services",
    }
    result = prices.get(product_name.lower(), f"Price not found for '{product_name}'")
    return json.dumps({"product": product_name, "price": result})


PRODUCT_PRICE_FUNCTION = FunctionDefinition(
    name="lookup_product_price",
    description="Look up the pricing for an Azure AI product or service.",
    parameters={
        "type": "object",
        "properties": {
            "product_name": {
                "type": "string",
                "description": "The name of the Azure AI product, e.g. 'Azure OpenAI'",
            },
        },
        "required": ["product_name"],
    },
)


def main() -> None:
    project_client = get_project_client()
    agents_client = project_client.agents

    deployment = get_env("AZURE_OPENAI_CHAT_DEPLOYMENT")

    # --- Build the tool set --------------------------------------------------
    toolset = ToolSet()

    # Code Interpreter — lets the agent write and run Python
    toolset.definitions.append(CodeInterpreterTool())

    # Function Calling — lets the agent call lookup_product_price
    func_tool = FunctionTool(definitions=[PRODUCT_PRICE_FUNCTION])
    toolset.definitions.append(func_tool)

    print("🔧 Creating agent with Code Interpreter + Function Calling...")

    # --- Create the agent ----------------------------------------------------
    agent = agents_client.create_agent(
        model=deployment,
        name="azure-ai-helper",
        instructions=(
            "You are an Azure AI assistant. You can:\n"
            "1. Write and run Python code using Code Interpreter.\n"
            "2. Look up Azure AI product pricing using the lookup_product_price function.\n"
            "Answer questions accurately and concisely."
        ),
        toolset=toolset,
    )
    print(f"   ✅ Agent created: {agent.id}")

    # --- Create a thread and send a message ----------------------------------
    thread = agents_client.create_thread()
    print(f"   📎 Thread created: {thread.id}")

    user_message = "What is the price of Azure OpenAI? Also, write Python code to calculate 2^10."
    agents_client.create_message(
        thread_id=thread.id,
        role="user",
        content=user_message,
    )
    print(f"   💬 User: {user_message}")

    # --- Run the agent -------------------------------------------------------
    run = agents_client.create_run(thread_id=thread.id, assistant_id=agent.id)
    print(f"   🚀 Run started: {run.id}")

    # Poll until complete
    while run.status in ("queued", "in_progress", "requires_action"):
        time.sleep(2)

        # Handle function calls
        if run.status == "requires_action":
            tool_calls = run.required_action.submit_tool_outputs.tool_calls
            tool_outputs = []
            for tc in tool_calls:
                if tc.function.name == "lookup_product_price":
                    args = json.loads(tc.function.arguments)
                    result = lookup_product_price(**args)
                    tool_outputs.append({"tool_call_id": tc.id, "output": result})

            agents_client.submit_tool_outputs_to_run(
                thread_id=thread.id, run_id=run.id, tool_outputs=tool_outputs
            )

        run = agents_client.get_run(thread_id=thread.id, run_id=run.id)
        print(f"   ⏳ Status: {run.status}")

    # --- Display results -----------------------------------------------------
    if run.status == "completed":
        messages = agents_client.list_messages(thread_id=thread.id)
        for msg in reversed(messages.data):
            role = msg.role.upper()
            for content_block in msg.content:
                if hasattr(content_block, "text"):
                    print(f"\n   [{role}] {content_block.text.value}")
    else:
        print(f"   ❌ Run ended with status: {run.status}")

    # --- Clean up ------------------------------------------------------------
    agents_client.delete_agent(agent.id)
    print("\n   🧹 Agent deleted. Done.")


if __name__ == "__main__":
    main()