import os
import time
from pathlib import Path

from dotenv import load_dotenv
from azure.identity import DefaultAzureCredential
from azure.ai.projects import AIProjectClient
from azure.ai.projects.models import BingGroundingTool, MessageTextContent

# Load environment variables
notebook_path = Path().absolute()
parent_dir = notebook_path.parent
load_dotenv(parent_dir / '.env')

# Initialize AIProjectClient
try:
    project_client = AIProjectClient.from_connection_string(
        credential=DefaultAzureCredential(),
        conn_str=os.environ.get("PROJECT_CONNECTION_STRING")
    )
    print("✅ Successfully initialized AIProjectClient")
except Exception as e:
    print(f"❌ Error initializing project client: {e}")

def create_bing_grounded_agent():
    """Create an agent that can use Bing to ground queries with up-to-date info."""
    try:
        # 1. Retrieve Bing connection from your AI Foundry project
        bing_conn_name = os.environ.get("BING_CONNECTION_NAME")
        if not bing_conn_name:
            raise ValueError("BING_CONNECTION_NAME not set in .env")

        bing_connection = project_client.connections.get(connection_name=bing_conn_name)
        conn_id = bing_connection.id
        print(f"Bing Connection ID: {conn_id}")

        # 2. Initialize Bing grounding tool
        bing_tool = BingGroundingTool(connection_id=conn_id)

        # 3. Create an agent that can search with Bing
        agent = project_client.agents.create_agent(
            model=os.environ.get("MODEL_DEPLOYMENT_NAME", "gpt-4o-mini"),
            name="health-bing-agent",
            instructions="""
                You are a health and fitness assistant with Bing search capabilities.
                Always:
                1. Provide disclaimers that you are not a medical professional.
                2. Encourage professional consultation.
                3. Use Bing for real-time references.
                4. Provide brief, helpful answers.
            """,
            tools=bing_tool.definitions,
            # Must pass special preview header to use Bing grounding (subject to change)
            headers={"x-ms-enable-preview": "true"},
        )

        print(f"🎉 Created Bing-grounded agent, ID: {agent.id}")
        return agent
    except Exception as e:
        print(f"❌ Error creating Bing-grounded agent: {e}")
        return None

# Create our Bing-based agent
bing_agent = create_bing_grounded_agent()

def ask_bing_question(agent, user_query):
    try:
        # Create a thread if we haven't yet
        thread = project_client.agents.create_thread()
        print(f"📝 Created a conversation thread, ID: {thread.id}")

        # Post user query as a message
        user_message = project_client.agents.create_message(
            thread_id=thread.id,
            role="user",
            content=user_query
        )
        print(f"📨 Created user message with query: '{user_query}'")

        # Process the query with the agent
        run = project_client.agents.create_and_process_run(
            thread_id=thread.id,
            agent_id=agent.id
        )
        print(f"🤖 Run finished with status: {run.status}")
        if run.last_error:
            print(f"Error detail: {run.last_error}")

        return thread, run
    except Exception as e:
        print(f"❌ Error asking Bing question: {e}")
        return None, None

if bing_agent:
    # Let's ask a few fun questions!
    questions = [
        "What are some new HIIT workout trends I should know about?",
        "What's the current WHO recommendation for sugar intake?",
        "Any news on intermittent fasting for weight management?"
    ]

    for q in questions:
        ask_bing_question(bing_agent, q)

def view_bing_responses(thread_id):
    try:
        messages = project_client.agents.list_messages(thread_id=thread_id)
        print("\n🗣️ Conversation so far (oldest to newest):")
        for m in reversed(messages.data):
            if m.content:
                last_content = m.content[-1]
                if hasattr(last_content, "text"):
                    print(f"{m.role.upper()}: {last_content.text.value}\n")
    except Exception as e:
        print(f"❌ Error viewing Bing responses: {e}")

def cleanup_bing_agent(agent):
    if agent:
        try:
            project_client.agents.delete_agent(agent.id)
            print(f"🗑️ Deleted Bing-grounded agent: {agent.name}")
        except Exception as e:
            print(f"❌ Error cleaning up agent: {e}")
    else:
        print("No agent to clean up.")

# Uncomment if you want to remove the agent now
# cleanup_bing_agent(bing_agent)