from fastapi import FastAPI, Request, staticfiles, Depends, HTTPException, status
from fastapi.responses import HTMLResponse, StreamingResponse, JSONResponse

from fastapi.security import HTTPBasic, HTTPBasicCredentials
import os, json, re, requests
from datetime import datetime, timedelta
import time
import jwt
import openai
from openai import error as openai_error
from hashlib import sha256
import tiktoken
from supported import supp_langs

app = FastAPI()
security = HTTPBasic()
tokenizer = tiktoken.encoding_for_model("gpt-3.5-turbo")
app.mount("/static", staticfiles.StaticFiles(directory="static"), name="static")
users = json.loads(str(os.getenv("USER_KEYS")).replace("\n", ""))
for key in users:
    if key == "master": continue
    password = key+users[key]+users["master"]
    users[key] = sha256(password.encode('UTF-8')).hexdigest()

def write_line(line):
    with open("log.txt", "a") as f:
        f.write(line)
        f.write("\n")

def write_multi_line(list_text):
    for line in list_text:
        write_line(line)
    write_line("---------------------------")


write_multi_line(["inicio"])

fecha_unix = str(int(time.time()))

JWT_SECRET = users["master"]
JWT_ALGORITHM = "HS256"
JWT_EXPIRATION_TIME_MINUTES = 30

def create_jwt_token(data):
    to_encode = {"data": data}
    expire = datetime.utcnow() + timedelta(minutes=JWT_EXPIRATION_TIME_MINUTES)
    to_encode.update({"exp": expire})
    encoded_jwt = jwt.encode(to_encode, JWT_SECRET, algorithm=JWT_ALGORITHM)
    return encoded_jwt

async def validate_token(request: Request):
    data = {}
    try:
        data = await request.json()
        token = data.pop("token")
        payload = jwt.decode(token, JWT_SECRET, algorithms=[JWT_ALGORITHM])
        data["token_data"] = payload["data"]
    except: 
        raise HTTPException(status_code=404, detail="Datos no válidos")    
    return data

def authenticate_user(credentials: HTTPBasicCredentials):

    password = credentials.username+credentials.password+users["master"]
    password = sha256(password.encode('UTF-8')).hexdigest()

    if credentials.username not in users or password != users[credentials.username]:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED,
            detail="Incorrect username or password",
            headers={"WWW-Authenticate": "Basic"},
        )

    return True

    
@app.get("/", response_class=HTMLResponse)
async def root(request: Request, credentials: HTTPBasicCredentials = Depends(security)):
    if authenticate_user(credentials):
        token = create_jwt_token({"user":credentials.username})
        with open(os.path.join("static", "main.html")) as f:
            return HTMLResponse(f.read().replace("{% token %}", token).replace("{% version %}", fecha_unix))

@app.get("/tabs", response_class=HTMLResponse)
async def root_tabs(request: Request, credentials: HTTPBasicCredentials = Depends(security)):
    if authenticate_user(credentials):
        token = create_jwt_token({"user":credentials.username})
        with open(os.path.join("static", "main2.html")) as f:
            return HTMLResponse(f.read().replace("{% token %}", token).replace("{% version %}", fecha_unix))

    
@app.post("/chat_stream")
async def chat_stream(data = Depends(validate_token)):
    messages = data.get("messages", "")
    if not messages:
        print("Empty message")
        error = "What??"
        raise HTTPException(
            status_code=status.HTTP_418_IM_A_TEAPOT,
            detail= error
        )
    try:
        token_length = len(tokenizer.encode(messages[-1]["content"]))
    except:
        token_length = len(messages[-1]["content"])
        print("Error in token length")
        print("Message:", messages[-1]["content"])

    config = {
        "temperature": float(data.get("config", []).get("temperature", 1)),
        "frequency_penalty": float(data.get("config", []).get("frequency_penalty", 1)),
        "presence_penalty": float(data.get("config", []).get("presence_penalty", 1))
    }
 
    try:
        response = openai.ChatCompletion.create(
            model="gpt-3.5-turbo",
            messages=messages,
            temperature=config["temperature"],
            frequency_penalty=config["frequency_penalty"],
            presence_penalty=config["presence_penalty"],
            request_timeout = 25,
            stream=True
        )
    except requests.exceptions.RequestException as e:
        print("Timeout (requests)")
        print(e)
        raise HTTPException(
            status_code=status.HTTP_408_REQUEST_TIMEOUT,
            detail="Los servidores tardaron mucho en responder, puede haber sobrecarga en OpenAI, reintenta luego (error 1)"
        )
    except openai_error.APIConnectionError as e:
        print("APIConnectionError")
        print(e)
        raise HTTPException(
            status_code=status.HTTP_408_REQUEST_TIMEOUT,
            detail="El servidor no respondió, puede haber sobrecarga en OpenAI, reintenta luego (error 2)"
        )
    except openai_error.Timeout as e:
        print("Timeout (openai)")
        print(e)
        raise HTTPException(
            status_code=status.HTTP_408_REQUEST_TIMEOUT,
            detail="El servidor no respondió, puede haber sobrecarga en OpenAI, reintenta luego (error 3)"
        )
    except openai_error.InvalidRequestError as e:
        print("Timeout (openai)")
        print(e)
        error = "El servidor no respondió, puede haber sobrecarga en OpenAI, reintenta luego (error 3)"
        if "This model's maximum context length is 4097 tokens" in e.message:
            error = "ChatGPT se gomitó 🤮, limpia el chat y reintenta."
        raise HTTPException(
            status_code=status.HTTP_408_REQUEST_TIMEOUT,
            detail= error
        )
    
    token = create_jwt_token(data.pop("token_data"))
    
    async def __streamer():
        yield json.dumps({"object": "chat.token",  "token": token})
        #yield json.dumps({"object": "chat.user.length", "length": token_length})
        tokens = 0
        for chunk in response:
            tokens += 1
            yield json.dumps(chunk)
        #yield({"object": "chat.assistant.length", "length": tokens})
            
    return StreamingResponse(__streamer(), media_type="application/json")

@app.get("/read_log", response_class=HTMLResponse)
async def read_log(request: Request, credentials: HTTPBasicCredentials = Depends(security)):
    if authenticate_user(credentials):
        with open("log.txt", "r") as f:
            return HTMLResponse(f.read())


@app.get("/mfdfastapi")
async def webhook_get(request: Request):
    
    write_multi_line([
        "GET mfdfastapi:",
        "headers: {}".format(json.dumps(dict(request.headers.items()))),
        "params : {}".format(json.dumps(dict(request.query_params.items()))),
        
    ])

    mode = request.query_params.get('hub.mode', "")
    if mode=="subscribe":
        challenge = request.query_params.get('hub.challenge', "")
        verify_token = request.query_params.get('hub.verify_token', "")

        if not mode or not challenge or not verify_token:
            raise HTTPException(status_code=404, detail="Datos no válidos")  

        if verify_token != users["master"]:
            raise HTTPException(status_code=404, detail="Datos no válidos")    
        return HTMLResponse(str(challenge))

@app.post("/mfdfastapi")
async def webhook_post(request: Request):
    parametros_post = await request.json()
    write_multi_line([
        "POST mfdfastapi:",
        "headers: {}".format(json.dumps(dict(request.headers.items()))),
        "params : {}".format(json.dumps(dict(request.query_params.items()))),
        "post   : {}".format(json.dumps(parametros_post))
    ])

    return HTMLResponse("ok")