from fastapi import FastAPI

import os
from dotenv import load_dotenv

import pandas as pd
from langchain.document_loaders import DirectoryLoader
from agents import master_agent, plant_agent, eda_agent, rag_agent

load_dotenv()

os.environ['OPENAI_API_KEY'] = os.environ.get("OPENAI_API_KEY")
os.environ['SERPAPI_API_KEY'] = os.getenv('SERPAPI_API_KEY')

master = master_agent.init_config()

print("init master agent")

plant = plant_agent.init_config()

print("init plant agent")

df = pd.read_csv('data/csv/plant_syn.csv')

eda = eda_agent.init_config(df)

print("init eda agent")

loader = DirectoryLoader("data/txt", glob="*.txt")

rag = rag_agent.init_config(loader)

loader = DirectoryLoader("data/txt", glob="*.txt")

documents = loader.load()

print("init rag agent")

app = FastAPI()

@app.get("/hello")
def hello():
    return {"message": "Hello World"}

@app.post("/ask")
def ask(question: str):
    category = eval(master_agent.answer_question(master, question))

    temp_question = f"Referring to the Blue Indigo False Plant, {question}"

    print(question)
    print(temp_question)
    print(category)

    if category['category_number'] == 1:
        response = eval(plant_agent.answer_question(plant, temp_question))

        category.update(response)

    elif category['category_number'] == 2:
        response = eda_agent.answer_question(eda, temp_question)

        category['response'] = response

    elif category['category_number'] == 3:
        response = rag_agent.answer_question(rag, temp_question)

        category['response'] = response

    category['question'] = question
    
    return category