from transformers import GPT2Tokenizer, GPT2LMHeadModel
from dataset import clean
import re
import gdown
import os
cache_dir = os.path.expanduser("~/.cache/gdown")
os.makedirs(cache_dir, exist_ok=True)

def load_tokenizer_model():
    folder_url = "https://drive.google.com/drive/folders/1DDJ9t-HfMrf6OLYim5bVrP20QgyOZahc?usp=drive_link"
    local_folder="ChatbotCheckpoint"
    if not os.path.exists(local_folder) or not os.listdir(local_folder):
        print("Downloading model and tokenizer...")
        gdown.download_folder(folder_url, output=local_folder, quiet=False)
    else:
        print("Model and tokenizer already cached, skipping download.")
    model_name = "ChatbotCheckpoint"
    model = GPT2LMHeadModel.from_pretrained(model_name)
    model.eval()
    tokenizer = GPT2Tokenizer.from_pretrained(model_name)
    tokenizer.pad_token = tokenizer.eos_token
    return tokenizer, model


def generate_answer(query):
    tokenizer,model=load_tokenizer_model()
    inputs = tokenizer.encode(query, return_tensors='pt')
    attention_mask = (inputs != tokenizer.pad_token_id).long() 
    outputs =model.generate(inputs, attention_mask=attention_mask, max_length=100, num_return_sequences=1)
    return tokenizer.decode(outputs[0], skip_special_tokens=True)


def inference(query):
    query=clean(query)
    if "power wash" in query:
        answer="Power wash Ctrl + Alt + Shift + r . On the logging Screen."
        return answer
    elif "change password" in query  or "changing password" in query :
       answer="""
    Password Change Process
    To Change Your Password – you must complete these steps on the RCS network. Note - preferably at
    the beginning of your work day – not at the end of the day when you are rushing out.
    1- If you use the RCS WIFI on your cell phone, you need to “Forget This Network”. On the iPhone: Settings – WIFI – RCS – “Forget This Network”
    2- Log in to your computer like you normally would.
    3- Hold down the Ctrl+Alt+Del keys all at once.
    4- Select Change a password.
    5- Type your current password in the Old Password field.
    6- Type your new password in the New Password field and in the Confirm password field.
    7- Press the Enter Key to complete the reset of your password.
    8- Once your password has been changed – please open Outlook and let the emails update and also log into your Google Drive.
    9-  When you are done using your computer for the day, please “shutdown” your computer before you leave for the day.
    10-  Lower left corner   click then then “Shutdown”.

    """
       return answer
    
    elif "battery life" in query:
        answer="Battery life test: 1- Control, Alt, T . 2- Battery_test on the script screen"
        return answer
    else:
        answer = generate_answer(query)
        modified_answer = re.sub(query, "", answer)
        answer=modified_answer.strip()
        return answer