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 (
    FileSearchTool,
    MessageTextContent,
    MessageRole
)

# Load environment variables from parent .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.get("PROJECT_CONNECTION_STRING")
    )
    print("✅ Successfully initialized AIProjectClient")
except Exception as e:
    print(f"❌ Error initializing project client: {e}")

def create_sample_files():
    recipes_md = (
        """# Healthy Recipes Database\n\n"
        "## Gluten-Free Recipes\n"
        "1. Quinoa Bowl\n"
        "   - Ingredients: quinoa, vegetables, olive oil\n"
        "   - Instructions: Cook quinoa, add vegetables\n\n"
        "2. Rice Pasta with Vegetables\n"
        "   - Ingredients: rice pasta, mixed vegetables\n"
        "   - Instructions: Boil pasta, sauté vegetables\n\n"
        "## Diabetic-Friendly Recipes\n"
        "1. Low-Carb Stir Fry\n"
        "   - Ingredients: chicken, vegetables, tamari sauce\n"
        "   - Instructions: Cook chicken, add vegetables\n\n"
        "2. Greek Salad\n"
        "   - Ingredients: cucumber, tomatoes, feta, olives\n"
        "   - Instructions: Chop vegetables, combine\n\n"
        "## Heart-Healthy Recipes\n"
        "1. Baked Salmon\n"
        "   - Ingredients: salmon, lemon, herbs\n"
        "   - Instructions: Season salmon, bake\n\n"
        "2. Mediterranean Bowl\n"
        "   - Ingredients: chickpeas, vegetables, tahini\n"
        "   - Instructions: Combine ingredients\n"""
    )

    guidelines_md = (
        """# Dietary Guidelines\n\n"
        "## General Guidelines\n"
        "- Eat a variety of foods\n"
        "- Control portion sizes\n"
        "- Stay hydrated\n\n"
        "## Special Diets\n"
        "1. Gluten-Free Diet\n"
        "   - Avoid wheat, barley, rye\n"
        "   - Focus on naturally gluten-free foods\n\n"
        "2. Diabetic Diet\n"
        "   - Monitor carbohydrate intake\n"
        "   - Choose low glycemic foods\n\n"
        "3. Heart-Healthy Diet\n"
        "   - Limit saturated fats\n"
        "   - Choose lean proteins\n"""
    )

    # Save to local .md files
    with open("recipes.md", "w", encoding="utf-8") as f:
        f.write(recipes_md)
    with open("guidelines.md", "w", encoding="utf-8") as f:
        f.write(guidelines_md)

    print("📄 Created sample resource files: recipes.md, guidelines.md")
    return ["recipes.md", "guidelines.md"]

sample_files = create_sample_files()

def create_vector_store(files, store_name="my_health_resources"):
    try:
        # Upload each file
        uploaded_ids = []
        for fp in files:
            upl = project_client.agents.upload_file_and_poll(file_path=fp, purpose="assistants")
            uploaded_ids.append(upl.id)
            print(f"✅ Uploaded: {fp} -> File ID: {upl.id}")

        # Now create the vector store from these file IDs
        vs = project_client.agents.create_vector_store_and_poll(file_ids=uploaded_ids, name=store_name)
        print(f"🎉 Created vector store '{store_name}', ID: {vs.id}")
        return vs, uploaded_ids
    except Exception as e:
        print(f"❌ Error creating vector store: {e}")
        return None, []

vector_store, file_ids = None, []
if sample_files:
    vector_store, file_ids = create_vector_store(sample_files, "health_resources_example")

def create_health_resource_agent(vstore_id):
    try:
        # Instantiate the file search tool
        file_search_tool = FileSearchTool(vector_store_ids=[vstore_id])

        # Create the agent
        agent = project_client.agents.create_agent(
            model=os.environ.get("MODEL_DEPLOYMENT_NAME", "gpt-4o-mini"),
            name="health-search-agent",
            instructions="""
                You are a health resource advisor with access to dietary and recipe files.
                You:
                1. Always present disclaimers (you're not a doctor!)
                2. Provide references to the files when possible
                3. Focus on general nutrition or recipe tips.
                4. Encourage professional consultation for more detailed advice.
            """,
            tools=file_search_tool.definitions,
            tool_resources=file_search_tool.resources
        )
        print(f"🎉 Created health resource agent, ID: {agent.id}")
        return agent
    except Exception as e:
        print(f"❌ Error creating health resource agent: {e}")
        return None

health_agent = None
if vector_store:
    health_agent = create_health_resource_agent(vector_store.id)

def create_search_thread(agent):
    try:
        # Create a new conversation thread
        thread = project_client.agents.create_thread()
        print(f"📝 Created new search thread, ID: {thread.id}")
        return thread
    except Exception as e:
        print(f"❌ Error creating search thread: {e}")
        return None

def ask_search_question(thread_id, agent_id, user_question):
    try:
        # Add a user message to the thread with the search query
        message = project_client.agents.create_message(
            thread_id=thread_id,
            role="user",
            content=user_question
        )
        print(f"🔎 Searching: '{user_question}'")

        # Now create_and_process_run to let the agent answer
        run = project_client.agents.create_and_process_run(
            thread_id=thread_id,
            assistant_id=agent_id
        )
        print(f"🤖 Run finished with status: {run.status}")
        if run.last_error:
            print(f"Error details: {run.last_error}")
        return run
    except Exception as e:
        print(f"❌ Error searching question: {e}")
        return None

# Let's do a couple of sample queries!
if health_agent:
    search_thread = create_search_thread(health_agent)

    if search_thread:
        queries = [
            "Could you suggest a gluten-free lunch recipe?",
            "Show me some heart-healthy meal ideas.",
            "What guidelines do you have for someone with diabetes?"
        ]

        for q in queries:
            ask_search_question(search_thread.id, health_agent.id, q)

def display_thread_messages(thread_id):
    try:
        messages = project_client.agents.list_messages(thread_id=thread_id)
        print("\n🗣️ Conversation so far:")
        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")

        # Print any file citations
        print("\n📎 Checking for citations...")
        for c in messages.file_citation_annotations:
            # The file ID is nested in c.file_citation, as a dict with key 'file_id'
            print(f"- Citation snippet: '{c.text}' from file ID: {c.file_citation['file_id']}")

    except Exception as e:
        print(f"❌ Error displaying messages: {e}")

def cleanup_all():
    try:
        # Delete vector store
        if 'vector_store' in globals() and vector_store:
            project_client.agents.delete_vector_store(vector_store.id)
            print("🗑️ Deleted vector store.")

        # Delete uploaded files from service
        if 'file_ids' in globals() and file_ids:
            for fid in file_ids:
                project_client.agents.delete_file(fid)
            print("🗑️ Deleted uploaded files from the service.")

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

        # Delete local sample files
        if 'sample_files' in globals() and sample_files:
            for sf in sample_files:
                if os.path.exists(sf):
                    os.remove(sf)
            print("🗑️ Deleted local sample files.")

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


cleanup_all()