import asyncio
import os
import dotenv
import pandas as pd

from agents import Agent, CodeInterpreterTool, Runner, trace

dotenv.load_dotenv()
os.environ["OPENAI_API_KEY"] = os.getenv("OPENAI_API_KEY")

CSV_PATH = r"C:\Users\Dell\Documents\MR-AI\openai_agents\healthcare-data-30.csv"

async def main():
    if not os.path.exists(CSV_PATH):
        raise FileNotFoundError(f"❌ CSV file not found at {CSV_PATH}")

    df = pd.read_csv(CSV_PATH)
    print(f"📄 Loaded CSV data (first 5 rows):\n{df.head()}\n")

    # Prepare textual context for the model
    df_preview = df.head(10).to_markdown()
    data_context = (
        f"The dataset preview is:\n{df_preview}\n\n"
        f"Column headers: {', '.join(df.columns)}"
    )

    # Define the agent
    agent = Agent(
        name="CSV Code Interpreter",
        model="gpt-4.1",
        instructions=(
            "You are a data analyst who loves working with CSV data. "
            "You can analyze and visualize data from pandas DataFrames using Python code. "
            "Use the provided CSV data to answer questions.\n"
        ),
        tools=[
            CodeInterpreterTool(
                tool_config={"type": "code_interpreter", "container": {"type": "auto"}}
            )
        ],
    )

    query = input("💬 Ask a question about the CSV file: ")

    with trace("CSV code interpreter example"):
        print("\n🔍 Processing query...")

        # ✅ Combine the CSV context + user query in one go
        result = Runner.run_streamed(
            agent,
            f"Dataset context:\n{data_context}\n\nUser query: {query}"
        )

        async for event in result.stream_events():
            if (
                event.type == "run_item_stream_event"
                and event.item.type == "tool_call_item"
                and event.item.raw_item.type == "code_interpreter_call"
            ):
                print(f"\n💻 Code executed:\n```\n{event.item.raw_item.code}\n```\n")
            elif event.type == "run_item_stream_event":
                print(f"Other event: {event.item.type}")

        print(f"\n✅ Final output: {result.final_output}")

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