from huggingface_hub import InferenceClient
import gradio as gr
import json
from datetime import datetime
import os
from google_auth_oauthlib.flow import InstalledAppFlow
from googleapiclient.discovery import build
from google.oauth2 import service_account
from googleapiclient.http import MediaFileUpload
from pathlib import Path
import argparse
from huggingface_hub import snapshot_download
import random 
import string

#TODO : Cleanup and comment


repo_name = 'TheBloke/OpenHermes-2.5-Mistral-7B-GGUF'
model_file = "openhermes-2.5-mistral-7b.Q4_K_M.gguf"

print('Fetching model:', repo_name, model_file)
snapshot_download(repo_id=repo_name, local_dir=".", allow_patterns=model_file)
print('Done fetching model:')

DEFAULT_MODEL_PATH = model_file

chat_history = []
chat_log_history = []
isFirstRun = True
dir_path = os.path.dirname(os.path.realpath(__file__))
directory = '/home/user/app'
if not os.path.exists(directory):
    print("directory dont exist")
else: print("path exists")
    
file_path = '/home/user/app/test.txt'
with open(file_path, 'w') as file:
    # Write "Hello World" to the file
    file.write("Hello World")
    print("file written")

print(dir_path)

SCOPES = ['https://www.googleapis.com/auth/drive']
SERVICE_ACCOUNT_FILE = '/home/user/app/service_account_creds.json'

#initContext = "You are a hospital patient nameed Barry Lawson. You have an injured arm and are in great pain. You are frustrated because you have been waiting for a long time. Every time I prompt you you get more agressive and violent in your responses. Your responses should use excessive profanity. You will provide only one response to each of my prompts. I am a nurse and you are angry at me. Nurse: "

initContext = """<|im_start|>You are playing the role of an aggressive patient called Barry Lawson. Barry has been waiting for hours at the hospital with severe pain in his chest and hasn't been seen yet. You are incredibly frustrated and angry that things are taking so long, you will only calm down if your concerns are addressed. Waiting will only make you more angry and aggressive towards the staff. At any moment you could break into a RAGE! Make sure you don't censor swear words. Make sure to explain your emotional state and thinking before responding, for example, Barry: (tired of waiting and very angry) What the fuck do I have to do around here to get some treatment!"""
context = ""
timestamp = datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
folder_id = '1Hjbu7FmuVs0Yz8y_veo6SzY_2tc48OWt'

Name = ""
Occupation = ""
Ethnicity = ""
Gender = ""
Age=""
YearsOfExp = ""

chat_log_name =""

from llama_cpp import Llama
llm = Llama(model_path=model_file, model_type="mistral",n_gpu_layers=-1,n_ctx = 2048)

def load_model():
    global llm
    llm = Llama(model_path=model_file, model_type="mistral",n_gpu_layers=-1,n_ctx = 2048)
    return "Model loaded"

def generate_unique_id():
    # Generate a random sequence of 3 letters and 3 digits
    letters = ''.join(random.choices(string.ascii_letters, k=3))
    digits = ''.join(random.choices(string.digits, k=3))
    unique_id = letters + digits
    return unique_id

unique_id = generate_unique_id()

def get_drive_service():
    credentials = service_account.Credentials.from_service_account_file(
        SERVICE_ACCOUNT_FILE, scopes=SCOPES)
    service = build('drive', 'v3', credentials=credentials)
    print("Google Service Created")
    return service

service = get_drive_service()
    
def search_file():
    #Search for a file by name in the specified Google Drive folder.
    query = f"name = '{chat_log_name}' and '{folder_id}' in parents and trashed = false"
    response = service.files().list(q=query, spaces='drive', fields='files(id, name)').execute()
    files = response.get('files', [])
    if not files:
        print(f"Chat log {chat_log_name} does not exist")
    else:
        print(f"Chat log {chat_log_name} exist")
    return files


def strip_special_tokens(text):
    # List of special tokens to be removed
    special_tokens = ["</s>", "<s>", "[INST]", "[/INST]"]
    
    # Iterate over the list of special tokens and replace each with an empty string
    for token in special_tokens:
        text = text.replace(token, "")
    
    return text
    
def upload_to_google_drive():

    existing_files = search_file()
    print(existing_files)
    
    data = {
            "name": Name,
            "occupation": Occupation,
            "years of experience": YearsOfExp,
            "ethnicity": Ethnicity,
            "gender": Gender,
            "age": Age,
            "chat_history": chat_log_history
            }
    
    with open(chat_log_name, "w") as log_file:
            json.dump(data, log_file, indent=4)

    if not existing_files:
        # If the file does not exist, upload it
        file_metadata = {
            'name': chat_log_name,
            'parents': [folder_id],'mimeType': 'application/json'
        }
        media = MediaFileUpload(chat_log_name, mimetype='application/json')
        file = service.files().create(body=file_metadata, media_body=media, fields='id').execute()
        print(f"Uploaded new file with ID: {file.get('id')}")
    else:
        print(f"File '{chat_log_name}' already exists.")
        # Example: Update the file content
        file_id = existing_files[0]['id']
        media = MediaFileUpload(chat_log_name, mimetype='application/json')
        updated_file = service.files().update(fileId=file_id, media_body=media).execute()
        print(f"Updated existing file with ID: {updated_file.get('id')}")
        

def generate(prompt, history):

    global isFirstRun,initContext,Name,Occupation,Ethnicity,Gender,Age,context,YearsOfExp
    
    if not len(Name) == 0 and not len(Occupation) == 0 and not len(Ethnicity) == 0 and not len(Gender) == 0 and not len(Age) == 0 and not len(YearsOfExp):
        firstmsg =""
        if isFirstRun:
            context = initContext
            isFirstRun = False
            firstmsg = prompt
            

        context += """
                  <|im_start|>nurse
                  Nurse:"""+prompt+"""
                  <|im_start|>barry
                  Barry:
                  """

        response = ""
        history2 = []

        while(len(response) < 1):
            output = llm(context, max_tokens=400, stop=["Nurse:"], echo=False)
            response = output["choices"][0]["text"]
            response = response.strip()
            yield response


      #  for output in llm(input, stream=True, max_tokens=100, ):
        #    piece = output['choices'][0]['text']
       #     response += piece
       #     chatbot[-1] = (chatbot[-1][0], response)

       #     yield  response
            
        history.append((prompt,response))
        if  not isFirstRun:
            chat_log_history.append({"user": prompt, "bot": response})
            upload_to_google_drive()
        else:
            chat_log_history.append({"user": firstmsg, "bot": reponse})
        
        context += response
        
        print (context)
        return history
        
    else:
        output = "Did you forget to enter your Details? Please go to the User Info Tab and Input your data. "
        yield output


def name_interface(name,occupation,yearsofexp,ethnicity,gender,age,reset_button):
    global Name, Occupation,Ethnicity,Gender,Age,chat_log_name,YearsOfExp
    print(reset_button)
    Name = name
    Occupation = occupation
    Ethnicity=ethnicity
    Gender=gender
    Age=age
    YearsOfExp = yearsofexp
    
    if name and occupation and ethnicity and gender and age and yearsofexp: 
        chat_log_name =  f'chat_log_for_{Name}_{datetime.now().strftime("%Y-%m-%d_%H-%M-%S")}.json'
        return f"You can start chatting now {Name}"    
    else:
        return "Enter ALL the details to start chatting"
        
def reset_chat_interface():
    global chat_history, chat_log_history, isFirstRun
    chat_history = []
    chat_log_history = []
    isFirstRun = True
    return "Chat has been reset."

def reset_name_interface():
    global Name, Occupation, Ethnicity, Gender, Age,YearsOfExp, chat_log_name
    Name = ""
    Occupation = ""
    YearsOfExp = ""
    Ethnicity = ""
    Gender = ""
    Age = ""
    chat_log_name = ""
    return "User info has been reset."

def reset_all():
    global unique_id
    message1 = reset_chat_interface()
    message2 = reset_name_interface()
    message3 = load_model()
    unique_id = generate_unique_id()
    return f"All Chat components have been rest.  Uniqe ID for this session is, {unique_id}. Please note this down."

chat_bot=gr.ChatInterface(
    fn=generate,
    chatbot=gr.Chatbot(show_label=False, show_share_button=False, show_copy_button=True, likeable=True, layout="panel"),
    title="""AI Chat Bot - ECU IVADE"""
)

name_interface = gr.Interface(
    fn=name_interface,
    inputs=[
        gr.Textbox(label="Unique ID", value = unique_id ,container = True, scale = 2 ,interactive = False, show_copy_button = True),
        gr.Textbox(label="Name", placeholder="Enter your name here..."),
        gr.Textbox(label="Occupation", placeholder="Enter your occupation here..."),
        gr.Textbox(label="Years of Experience", placeholder="Enter the years of experience"),
        gr.Textbox(label="Ethnicity", placeholder="Enter your Ethnicity here..."),
        gr.Dropdown(choices=["Male", "Female","Other","Prefer Not To Say"], label="Gender"),
        gr.Textbox(label="Age", placeholder="Enter your Age here..."),
    ],outputs="text",
    title="ECU-IVADE : User Information",
    description="Please enter your name and occupation."
)


with gr.Blocks() as reset_section:
    gr.Markdown("Press the button below to restart the chat Interface.")
    with gr.Row():
        out = gr.Textbox(label="Output")
    btn = gr.Button("Reset ChatBot Instance")
    btn.click(fn=reset_all, inputs=None, outputs=out)


tabs = gr.TabbedInterface([name_interface, chat_bot,reset_section], ["User Info", "Chat Bot","Reset Chatbot"])

if __name__ == "__main__":
    tabs.launch(debug=True,share=False,inbrowser=True)