# Import required libraries
import os
import time
from pathlib import Path

import pandas as pd
from dotenv import load_dotenv
from azure.identity import DefaultAzureCredential
from azure.ai.projects import AIProjectClient
from azure.ai.projects.models import CodeInterpreterTool, FilePurpose, MessageTextContent

# Load environment variables from the parent directory's .env
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["PROJECT_CONNECTION_STRING"],
    )
    print("✅ Successfully initialized AIProjectClient")
except Exception as e:
    print(f"❌ Error initializing client: {str(e)}")

# Create sample CSV data for demonstration
def create_sample_data():
    try:
        data = {
            'Date': pd.date_range(start='2024-01-01', periods=7),
            'Calories': [2100, 1950, 2300, 2050, 1900, 2200, 2150],
            'Protein_g': [80, 75, 85, 78, 72, 82, 79],
            'Carbs_g': [250, 230, 270, 245, 225, 260, 255],
            'Fat_g': [70, 65, 75, 68, 63, 73, 71],
            'Fiber_g': [25, 22, 28, 24, 21, 26, 23]
        }
        df = pd.DataFrame(data)
        filename = "nutrition_data.csv"
        df.to_csv(filename, index=False)
        print(f"📄 Created sample data file: {filename}")
        return filename
    except Exception as e:
        print(f"❌ Error creating sample data: {e}")
        return None

sample_file = create_sample_data()

def create_health_calculator(file_path):
    """Create an agent with code interpreter for health/nutrition calculations."""
    try:
        uploaded_file = project_client.agents.upload_file_and_poll(
            file_path=file_path,
            purpose=FilePurpose.AGENTS
        )
        print(f"✅ Uploaded CSV file, ID: {uploaded_file.id}")

        # Create a Code Interpreter tool referencing the uploaded file
        code_tool = CodeInterpreterTool(file_ids=[uploaded_file.id])

        # Create the agent with instructions to do basic calculations
        agent = project_client.agents.create_agent(
            model=os.environ.get("MODEL_DEPLOYMENT_NAME", "gpt-4o-mini"),
            name="health-calculator-agent",
            instructions="""
            You are a health calculator agent that can:
            1. Calculate and interpret BMI
            2. Analyze provided nutrition data
            3. Generate charts/plots
            4. Include disclaimers that you are not a medical professional
            """,
            tools=code_tool.definitions,
            tool_resources=code_tool.resources
        )
        print(f"🎉 Created health calculator agent, ID: {agent.id}")
        return agent, uploaded_file
    except Exception as e:
        print(f"❌ Error creating health calculator agent: {e}")
        return None, None

health_agent, uploaded_file = None, None
if sample_file:
    health_agent, uploaded_file = create_health_calculator(sample_file)

def calculate_bmi_with_agent(agent, height_inches, weight_pounds):
    """Calculate BMI using the code interpreter agent."""
    try:
        # Create a new conversation thread
        thread = project_client.agents.create_thread()
        print(f"📝 Created thread for BMI calculation, ID: {thread.id}")

        # Construct user message requesting BMI calculation
        user_text = (
            f"Calculate BMI for \n"
            f"Height: {height_inches} inches\n"
            f"Weight: {weight_pounds} pounds\n"
            "Please: \n"
            "1. Show calculation \n"
            "2. Interpret the result \n"
            "3. Include disclaimers \n"
        )

        msg = project_client.agents.create_message(
            thread_id=thread.id,
            role="user",
            content=user_text
        )
        print(f"➕ Created BMI request message, ID: {msg.id}")

        # Create and process the run, letting the agent handle code
        run = project_client.agents.create_and_process_run(
            thread_id=thread.id,
            assistant_id=agent.id
        )
        print(f"🤖 BMI run finished with status: {run.status}")
        return thread, run
    except Exception as e:
        print(f"❌ Error during BMI calculation: {e}")
        return None, None

if health_agent:
    bmi_thread, bmi_run = calculate_bmi_with_agent(health_agent, 70, 180)  # example: 5'10" and 180 lbs

def analyze_nutrition_data(agent):
    """Ask the agent to analyze the uploaded nutrition data."""
    try:
        thread = project_client.agents.create_thread()
        print(f"📝 Created thread for nutrition analysis, ID: {thread.id}")

        user_text = (
            "Analyze the CSV file with daily nutrition data.\n"
            "1. Compute average daily macros (calories, protein, carbs, fat, fiber).\n"
            "2. Create a chart to show trends.\n"
            "3. Discuss any insights or disclaimers.\n"
        )

        msg = project_client.agents.create_message(
            thread_id=thread.id,
            role="user",
            content=user_text
        )
        print(f"➕ Created nutrition request message, ID: {msg.id}")

        run = project_client.agents.create_and_process_run(
            thread_id=thread.id,
            assistant_id=agent.id
        )
        print(f"🤖 Nutrition run finished with status: {run.status}")
        return thread, run
    except Exception as e:
        print(f"❌ Error analyzing nutrition data: {e}")
        return None, None

if health_agent:
    nutrition_thread, nutrition_run = analyze_nutrition_data(health_agent)

def view_agent_responses(thread_id):
    try:
        messages = project_client.agents.list_messages(thread_id=thread_id)
        print("\n🔎 Agent Responses:")
        for msg in messages.data:
            if msg.role == "assistant" and msg.content:
                for c in msg.content:
                    if hasattr(c, "text"):
                        print("Response:", c.text.value, "\n")

        # If images were generated, let's save them
        for img in messages.image_contents:
            file_id = img.image_file.file_id
            outname = f"chart_{file_id}.png"
            project_client.agents.save_file(file_id=file_id, file_name=outname)
            print(f"🖼️ Saved image output: {outname}")

    except Exception as e:
        print(f"❌ Error viewing agent responses: {e}")

# Display BMI calculations
if bmi_thread and bmi_run:
    print("\n=== BMI Calculation Results ===")
    view_agent_responses(bmi_thread.id)

# Display nutrition analyses
if nutrition_thread and nutrition_run:
    print("\n=== Nutrition Analysis Results ===")
    view_agent_responses(nutrition_thread.id)

def cleanup_all():
    try:
        # Delete the uploaded CSV file from the service
        if 'uploaded_file' in globals() and uploaded_file:
            project_client.agents.delete_file(uploaded_file.id)
            print("🗑️ Deleted uploaded file from agent service.")

        # Delete the agent if we created one
        if 'health_agent' in globals() and health_agent:
            project_client.agents.delete_agent(health_agent.id)
            print("🗑️ Deleted health calculator agent.")

        # Delete local CSV file
        if 'sample_file' in globals() and sample_file and os.path.exists(sample_file):
            os.remove(sample_file)
            print("🗑️ Deleted local sample CSV file.")

    except Exception as e:
        print(f"❌ Error during cleanup: {e}")

cleanup_all()