import datetime
from json import loads
import time
from typing import List
from fastapi import Body, 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
from semantic_search import semantic_cache

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=["*"],
)
class ChatHistoryItem(BaseModel):
    q: str
    a: str
    
@app.post("/agents")
def agent(question: str, chat_history: List[ChatHistoryItem] = Body(...)):
    start = time.time()
    
    with get_openai_callback() as cb:
        cache_question, cache_answer = semantic_cache(supabase, question)
        if cache_answer:
            processing_time = time.time() - start
            save_history(question, cache_answer, cb, processing_time)

            return {"Answer": cache_answer}
    
        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, 
                ssl_keyfile="/etc/ssl_file/key.pem", 
                ssl_certfile="/etc/ssl_file/cert.pem")