# Copyright (c) Microsoft. All rights reserved.

import asyncio
import contextlib

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

# Load environment variables from .env file
load_dotenv()

"""
Foundry Chat Client with File Search Example

This sample demonstrates using get_file_search_tool() with FoundryChatClient
for direct document-based question answering and information retrieval.

Prerequisites:
- Set environment variables:
  - FOUNDRY_PROJECT_ENDPOINT: Your Foundry project endpoint URL
  - FOUNDRY_MODEL: Your Responses API deployment name
- Authenticate via 'az login' for AzureCliCredential
"""

# Helper functions


async def create_vector_store(client: FoundryChatClient) -> tuple[str, str]:
    """Create a vector store with sample documents."""
    file = await client.client.files.create(
        file=("todays_weather.txt", b"The weather today is sunny with a high of 75F."), purpose="assistants"
    )
    vector_store = await client.client.vector_stores.create(
        name="knowledge_base",
        expires_after={"anchor": "last_active_at", "days": 1},
    )
    result = await client.client.vector_stores.files.create_and_poll(vector_store_id=vector_store.id, file_id=file.id)
    if result.last_error is not None:
        raise Exception(f"Vector store file processing failed with status: {result.last_error.message}")

    return file.id, vector_store.id


async def delete_vector_store(client: FoundryChatClient, file_id: str, vector_store_id: str) -> None:
    """Delete the vector store after using it."""
    with contextlib.suppress(Exception):
        await client.client.vector_stores.delete(vector_store_id=vector_store_id)
    with contextlib.suppress(Exception):
        await client.client.files.delete(file_id=file_id)


async def main() -> None:
    print("=== Foundry Chat Client with File Search Example ===\n")

    # Initialize the Foundry chat client
    # Make sure you're logged in via 'az login' before running this sample
    client = FoundryChatClient(credential=AzureCliCredential())

    file_id, vector_store_id = await create_vector_store(client)

    # Create file search tool using instance method
    file_search_tool = client.get_file_search_tool(vector_store_ids=[vector_store_id])

    agent = Agent(
        client=client,
        instructions="You are a helpful assistant that can search through files to find information.",
        tools=[file_search_tool],
    )

    query = "What is the weather today? Do a file search to find the answer."
    print(f"User: {query}")
    result = await agent.run(query)
    print(f"Agent: {result}\n")

    await delete_vector_store(client, file_id, vector_store_id)


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