import datetime
from json import loads
import time
from typing import List
from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware

import pandas as pd
from pydantic import BaseModel
import uvicorn

from dotenv import load_dotenv
import os
from supabase.client import Client, create_client

from langchain.callbacks import get_openai_callback

from ai_agent import main

load_dotenv()
URI = os.getenv("SUPABASE_URI")
supabase_url = os.environ.get("SUPABASE_URL")
supabase_key = os.environ.get("SUPABASE_KEY")
supabase: Client = create_client(supabase_url, supabase_key)


app = FastAPI()
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

@app.get("/agents")
def agent(question: str):
    start = time.time()
    with get_openai_callback() as cb:
        answer = main(question)
    processing_time = time.time() - start
    save_history(question, answer, cb, processing_time)
    return {"answer": answer}  

def save_history(question, answer, cb, processing_time):
    # reference = [doc.dict() for doc in reference]
    record = {
        'Question': question,
        'Answer': answer,
        'Total_Tokens': cb.total_tokens,
        'Total_Cost': cb.total_cost,
        'Processing_time': processing_time,
    }
    response = (
        supabase.table("agent_records")
        .insert(record)
        .execute()
    )

class history_output(BaseModel):
    Question: str
    Answer: str
    Total_Tokens: int
    Total_Cost: float
    Processing_time: float
    Time: datetime.datetime
    
@app.get('/history', response_model=List[history_output])
async def get_history():
    response = supabase.table("agent_records").select("*").execute()
    df = pd.DataFrame(response.data)

    # engine = create_engine(URI, echo=True)

    # df = pd.read_sql_table("systex_records", engine.connect())  
    # df.fillna('', inplace=True)
    result = df.to_json(orient='index', force_ascii=False)
    result = loads(result)
    return result.values()  

if __name__ == "__main__":
    
    uvicorn.run("systex_app:app", host='0.0.0.0', reload=True, port=8080)