import io
from functions import *
from PyPDF2 import PdfReader
import pandas as pd
from fastapi import FastAPI, File, UploadFile
from fastapi.middleware.cors import CORSMiddleware
from langchain_community.document_loaders import UnstructuredURLLoader



app = FastAPI(title = "ConversAI", root_path = "/api/v1")
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

@app.post("/signup")
async def signup(username: str, password: str):
    response = createUser(username = username, password = password)
    return response


@app.post("/login")
async def login(username: str, password: str):
    response = matchPassword(username = username, password = password)
    return response


@app.post("/newChatbot")
async def newChatbot(chatbotName: str, username: str):
    client.table("ConversAI_ChatbotInfo").insert({"username": username, "chatbotname": chatbotName}).execute()
    chatbotName = f"convai-{username}-{chatbotName}"
    return createTable(tablename = chatbotName)


@app.post("/addPDF")
async def addPDFData(vectorstore: str, pdf: UploadFile = File(...)):
    pdf = await pdf.read()
    reader = PdfReader(io.BytesIO(pdf))
    text = ""
    for page in reader.pages:
        text += page.extract_text()
    username, chatbotname = vectorstore.split("-")[1], vectorstore.split("-")[2]
    df = pd.DataFrame(client.table("ConversAI_ChatbotInfo").select("*").execute().data)
    currentCount = df[(df["username"] == username) & (df["chatbotname"] == chatbotname)]["charactercount"].iloc[0]
    newCount = currentCount + len(text)
    if newCount < 400000:
        client.table("ConversAI_ChatbotInfo").update({"charactercount": str(newCount)}).eq("username", username).eq("chatbotname", chatbotname).execute()
        return addDocuments(text = text, vectorstore = vectorstore)
    else:
        return {
            "output": "DOCUMENT EXCEEDING LIMITS, PLEASE TRY WITH A SMALLER DOCUMENT."
        }


@app.post("/addText")
async def addText(vectorstore: str, text: str):
    username, chatbotname = vectorstore.split("-")[1], vectorstore.split("-")[2]
    df = pd.DataFrame(client.table("ConversAI_ChatbotInfo").select("*").execute().data)
    currentCount = df[(df["username"] == username) & (df["chatbotname"] == chatbotname)]["charactercount"].iloc[0]
    newCount = currentCount + len(text)
    if newCount < 400000:
        client.table("ConversAI_ChatbotInfo").update({"charactercount": str(newCount)}).eq("username", username).eq("chatbotname", chatbotname).execute()
        return addDocuments(text = text, vectorstore = vectorstore)
    else:
        return {
            "output": "WEBSITE EXCEEDING LIMITS, PLEASE TRY WITH A SMALLER DOCUMENT."
        }


@app.post("/addWebsite")
async def addWebsite(vectorstore: str, websiteUrl: str):
    urls = getLinks(websiteUrl)
    loader = UnstructuredURLLoader(urls=urls)
    docs = loader.load()
    text = "\n\n\n\n".join([f"Metadata:\n{docs[doc].metadata} \nPage Content:\n {docs[doc].page_content}" for doc in range(len(docs))])
    username, chatbotname = vectorstore.split("-")[1], vectorstore.split("-")[2]
    df = pd.DataFrame(client.table("ConversAI_ChatbotInfo").select("*").execute().data)
    currentCount = df[(df["username"] == username) & (df["chatbotname"] == chatbotname)]["charactercount"].iloc[0]
    newCount = currentCount + len(text)
    if newCount < 400000:
        client.table("ConversAI_ChatbotInfo").update({"charactercount": str(newCount)}).eq("username", username).eq("chatbotname", chatbotname).execute()
        return addDocuments(text = text, vectorstore = vectorstore)
    else:
        return {
            "output": "WEBSITE EXCEEDING LIMITS, PLEASE TRY WITH A SMALLER DOCUMENT."
        }

@app.post("/answerQuery")
async def answerQuestion(query: str, vectorstore: str, llmModel: str = "llama3-70b-8192"):
    return answerQuery(query=query, vectorstore=vectorstore, llmModel=llmModel)


@app.post("/deleteChatbot")
async def delete(chatbotName: str):
    username, chatbotName = chatbotName.split("-")[1], chatbotName.split("-")[2]
    client.table('ConversAI_ChatbotInfo').delete().eq('username', username).eq('chatbotname', chatbotName).execute()
    return deleteTable(tableName=chatbotName)

@app.post("/listChatbots")
async def delete(username: str):
    return listTables(username=username)

@app.post("/getLinks")
async def crawlUrl(baseUrl: str):
    return {
        "urls": getLinks(url=baseUrl, timeout=30)
        }

@app.post("/getCurrentCount")
async def getCount(vectorstore: str):
    username, chatbotName = chatbotName.split("-")[1], chatbotName.split("-")[2]
    df = pd.DataFrame(client.table("ConversAI_ChatbotInfo").select("*").execute().data)
    return {
        "currentCount": df[(df['username'] == username) & (df['chatbotname'] == chatbotName)]['charactercount'].iloc[0]
        }