# Copyright (c) Microsoft. All rights reserved.

import asyncio
import os
import tempfile

from agent_framework import Agent
from agent_framework.foundry import FoundryChatClient
from azure.identity import AzureCliCredential
from dotenv import load_dotenv
from openai import AsyncOpenAI

load_dotenv()

"""
Foundry Chat Client with Code Interpreter and Files Example

This sample demonstrates using get_code_interpreter_tool() with Responses on Foundry
for Python code execution and data analysis with uploaded files.

Environment variables:
    FOUNDRY_PROJECT_ENDPOINT — Foundry project endpoint
    FOUNDRY_MODEL            — Foundry model to use (e.g. "gpt-4o-mini")
"""

# Helper functions


async def create_sample_file_and_upload(openai_client: AsyncOpenAI) -> tuple[str, str]:
    """Create a sample CSV file and upload it for Foundry code interpreter use."""
    csv_data = """name,department,salary,years_experience
Alice Johnson,Engineering,95000,5
Bob Smith,Sales,75000,3
Carol Williams,Engineering,105000,8
David Brown,Marketing,68000,2
Emma Davis,Sales,82000,4
Frank Wilson,Engineering,88000,6
"""

    # Create temporary CSV file
    with tempfile.NamedTemporaryFile(mode="w", suffix=".csv", delete=False) as temp_file:
        temp_file.write(csv_data)
        temp_file_path = temp_file.name

    # Upload file for the code interpreter tool
    print("Uploading file for code interpreter...")
    with open(temp_file_path, "rb") as file:
        uploaded_file = await openai_client.files.create(
            file=file,
            purpose="assistants",  # Required for code interpreter
        )

    print(f"File uploaded with ID: {uploaded_file.id}")
    return temp_file_path, uploaded_file.id


async def cleanup_files(openai_client: AsyncOpenAI, temp_file_path: str, file_id: str) -> None:
    """Clean up both local temporary file and uploaded file."""
    # Clean up: delete the uploaded file
    await openai_client.files.delete(file_id)
    print(f"Cleaned up uploaded file: {file_id}")

    # Clean up temporary local file
    os.unlink(temp_file_path)
    print(f"Cleaned up temporary file: {temp_file_path}")


async def main() -> None:
    print("=== Foundry Chat Client with Code Interpreter and File Upload ===")

    # Create the FoundryChatClient
    client = FoundryChatClient(
        project_endpoint=os.getenv("FOUNDRY_PROJECT_ENDPOINT"),
        model=os.getenv("FOUNDRY_MODEL"),
        credential=AzureCliCredential(),
    )
    # use the openai client from the foundry client to upload files for the code interpreter tool
    openai_client = client.project_client.get_openai_client()
    temp_file_path, file_id = await create_sample_file_and_upload(openai_client)
    # Create agent with code interpreter tool with file access
    agent = Agent(
        client=client,
        instructions="You are a helpful assistant that can analyze data files using Python code.",
        tools=FoundryChatClient.get_code_interpreter_tool(file_ids=[file_id]),
    )
    try:
        # Test the code interpreter with the uploaded file
        query = "Analyze the employee data in the uploaded CSV file. Calculate average salary by department."
        print(f"User: {query}")
        result = await agent.run(query)
        print(f"Agent: {result.text}")
    finally:
        await cleanup_files(openai_client, temp_file_path, file_id)


if __name__ == "__main__":
    asyncio.run(main())
