# Import required libraries
import os  # For environment variables and path operations
import time  # For adding delays if needed
from pathlib import Path  # For cross-platform path handling

# Import Azure and utility libraries
from dotenv import load_dotenv  # For loading environment variables from .env file
from azure.identity import DefaultAzureCredential  # For Azure authentication
from azure.ai.projects import AIProjectClient  # Main client for AI Projects
from azure.ai.projects.models import MessageTextContent  # For handling message content

# Get the path to the .env file which is in the parent directory
notebook_path = Path().absolute()  # Get absolute path of current notebook
parent_dir = notebook_path.parent  # Get parent directory
load_dotenv(parent_dir / '.env')  # Load environment variables from .env file

# Initialize the AI Project Client using connection string from environment variables
try:
    project_client = AIProjectClient.from_connection_string(
        # Use default Azure credentials for authentication
        credential=DefaultAzureCredential(),
        # Get the project connection string from environment variables
        conn_str=os.environ.get("PROJECT_CONNECTION_STRING")
    )
    print("✅ Successfully initialized AIProjectClient")
except Exception as e:
    # Print error message if client initialization fails
    print(f"❌ Error initializing project client: {str(e)}")

def create_health_advisor_agent():
    """Create a health advisor agent with disclaimers and basic instructions."""
    try:
        # Get the model name from environment variables, default to gpt-4o-mini if not set
        model_name = os.environ.get("MODEL_DEPLOYMENT_NAME", "gpt-4o-mini")

        # Create a new agent using the AIProjectClient
        # The agent will use the specified model and follow the given instructions
        agent = project_client.agents.create_agent(
            model=model_name,
            name="fun-fit-health-advisor",
            # Define the agent's behavior and responsibilities
            instructions="""
            You are a friendly AI Health Advisor.
            You provide general health, fitness, and nutrition information, but always:
            1. Include medical disclaimers.
            2. Encourage the user to consult healthcare professionals.
            3. Provide general, non-diagnostic advice around wellness, diet, and fitness.
            4. Clearly remind them you're not a doctor.
            5. Encourage safe and balanced approaches to exercise and nutrition.
            """
        )
        # Log success and return the created agent
        print(f"🎉 Created health advisor agent, ID: {agent.id}")
        return agent
    except Exception as e:
        # Handle any errors during agent creation
        print(f"❌ Error creating agent: {str(e)}")
        return None

# Create an instance of our health advisor agent
health_advisor = create_health_advisor_agent()

# Function to create a new conversation thread for health discussions
def start_health_conversation():
    """Create a new thread for health & fitness discussions."""
    try:
        # Create a new empty thread using the project client
        # A thread stores the back-and-forth messages between user and agent
        thread = project_client.agents.create_thread()
        
        # Print success message with the thread ID for reference
        print(f"📝 Created a new conversation thread, ID: {thread.id}")
        
        # Return the created thread object so we can use it later
        return thread
    except Exception as e:
        # If thread creation fails, print error and return None
        print(f"❌ Error creating thread: {str(e)}")
        return None

# Initialize a new conversation thread that we'll use for our health Q&A
health_thread = start_health_conversation()

def ask_health_question(thread_id, user_question):
    """Add a user message to the conversation thread.
    
    Args:
        thread_id: ID of the conversation thread
        user_question: The health/fitness question from the user
        
    Returns:
        Message object if successful, None if error occurs
    """
    try:
        # Create a new message in the thread from the user's perspective
        # The role="user" indicates this is a user message (vs assistant)
        return project_client.agents.create_message(
            thread_id=thread_id,
            role="user", 
            content=user_question
        )
    except Exception as e:
        print(f"❌ Error adding user message: {e}")
        return None

def process_thread_run(thread_id, agent_id):
    """Ask the agent to process the thread and generate a response.
    
    Args:
        thread_id: ID of the conversation thread
        agent_id: ID of the health advisor agent
        
    Returns:
        Run object if successful, None if error occurs
    """
    try:
        # Create a new run to have the agent process the thread
        run = project_client.agents.create_run(
            thread_id=thread_id,
            assistant_id=agent_id
        )

        # Poll the run status until completion or error
        # Status can be: queued, in_progress, requires_action, completed, failed
        while run.status in ["queued", "in_progress", "requires_action"]:
            time.sleep(1)  # Wait 1 second between status checks
            run = project_client.agents.get_run(
                thread_id=thread_id,
                run_id=run.id
            )

        print(f"🤖 Run completed with status: {run.status}")
        return run
    except Exception as e:
        print(f"❌ Error processing thread run: {str(e)}")
        return None

def view_thread_messages(thread_id):
    """Display all messages in the conversation thread in chronological order.
    
    Args:
        thread_id: ID of the conversation thread to display
    """
    try:
        # Get all messages in the thread
        messages = project_client.agents.list_messages(thread_id=thread_id)
        print("\n🗣️ Conversation so far (oldest to newest):")
        
        # Loop through messages in reverse order to show oldest first
        for m in reversed(messages.data):
            if m.content:
                # Extract the text content from the message
                # We only handle text messages for now (not images etc)
                last_content = m.content[-1]
                if isinstance(last_content, MessageTextContent):
                    print(f"{m.role.upper()}: {last_content.text.value}\n")
        print("-----------------------------------\n")
    except Exception as e:
        print(f"❌ Error viewing thread: {str(e)}")

# First verify that we have valid agent and thread objects before proceeding
if health_advisor and health_thread:
    # 1) Ask about BMI calculation and interpretation
    # This demonstrates how the agent handles technical health questions
    msg1 = ask_health_question(health_thread.id, "How do I calculate my BMI, and what does it mean?")
    # Process the BMI question and wait for agent's response
    run1 = process_thread_run(health_thread.id, health_advisor.id)

    # 2) Ask about personalized meal planning
    # This shows how the agent provides customized nutrition advice
    msg2 = ask_health_question(health_thread.id, "Can you give me a balanced meal plan for someone who exercises 3x a week?")
    # Process the meal plan question and wait for agent's response
    run2 = process_thread_run(health_thread.id, health_advisor.id)

    # Display the full conversation history to see both Q&As
    # This will show the agent's responses including any health disclaimers
    view_thread_messages(health_thread.id)
else:
    # Error handling if agent or thread initialization failed
    print("❌ Could not run example queries because agent or thread is None.")

# Function to clean up and delete the agent when we're done
def cleanup(agent):
    # Only attempt cleanup if we have a valid agent
    if agent:
        try:
            # Delete the agent using the project client
            project_client.agents.delete_agent(agent.id)
            # Print confirmation message with the agent name
            print(f"🗑️ Deleted health advisor agent: {agent.name}")
        except Exception as e:
            # Handle any errors that occur during deletion
            print(f"Error cleaning up agent: {e}")
    else:
        # If no agent was provided, inform the user
        print("No agent to clean up.")

# Call cleanup function to delete our health advisor agent
cleanup(health_advisor)