print("Loading models...", end="")

from data_loader import load_data
from rag import find_similar_examples
from prompter import system_prompt, user_prompt_for
from llm import llm_request

from dotenv import load_dotenv
import os

load_dotenv()

from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware

from datetime import datetime
import pytz

vietnam_tz = pytz.timezone("Asia/Ho_Chi_Minh")

now_vietnam = datetime.now(vietnam_tz)

df = load_data()
print("df shape:")
print(df.shape)
print("-> Done")

def reload_data():
    global df
    now_vietnam = datetime.now(vietnam_tz).strftime("%Y-%m-%d %H:%M:%S")
    print(f"[{now_vietnam}] Reloading data...")
    df = load_data() 
    print(f"-> Done | Update: df_shape = {df.shape}")

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

@app.get('/')
def get_():
    return { "Bye": "World!" }

@app.get('/check')
def check():
    LLM_URL = os.getenv("LLM_URL", "https://generativelanguage.googleapis.com/v1beta/openai/")
    LLM_API_KEY = os.getenv("LLM_API_KEY")
    LLM_MODEL = os.getenv("LLM_MODEL")
    LLM_MODEL_LEARN_LM = os.getenv("LLM_MODEL_LEARN_LM")
    print(f"url: {LLM_URL}")
    print(f"key: {LLM_API_KEY}")
    print(f"model: {LLM_MODEL}")
    print(f"model: {LLM_MODEL_LEARN_LM}")
    return { "url":  LLM_URL, "model": LLM_MODEL, "model": LLM_MODEL_LEARN_LM}

@app.post('/q')
def solve_q(q: str):
    print(q)
    print("Finding examples....", end="")
    examples = find_similar_examples(q, df, 5)
    print("-> Done\nSolving...", end="")
    print(examples)
    user_prompt = user_prompt_for(q, examples)
    result = llm_request(system_prompt, user_prompt)
    print("-> Done")
    print(result)
    print("\n\n---------------------\n\n")
    return { "r": result }

@app.post('/reload')
def reload():
    """ API để reload dữ liệu """
    reload_data()
    return {"status": "Data reloaded", "df_shape": df.shape}

print("-> Done\nLoading data...", end="")

def main():
    df = load_data()
    print("-> Done")
    while True:
        q = input("Input question: ")
        if q == "exit()" or q == "exit":
            print("FQA AI exit!")
            return
        print("Finding examples....", end="")
        examples = find_similar_examples(q, df, 8)
        print("-> Done\nSolving...", end="")
        print(examples)
        user_prompt = user_prompt_for(q, examples)
        result = llm_request(system_prompt, user_prompt)
        print("-> Done")
        print(result)
        print("\n\n---------------------\n\n")



# if __name__ == "__main__":
#     main()