import spaces
import gradio as gr
import torch
from transformers import AutoTokenizer
import json
from datetime import datetime
from uuid import uuid4
import os
from pathlib import Path

from huggingface_hub import CommitScheduler

# TODO make it so that feedback is only saved on prev. example if user makes another obfuscation
# and changes slider but doesn't hit obfuscate

tokenizer = AutoTokenizer.from_pretrained("gpt2")
MODEL_PATHS = {
    "length_more": "hallisky/lora-length-long-llama-3-8b",
    "length_less": "hallisky/lora-length-long-llama-3-8b",
    "function_more": "hallisky/lora-function-more-llama-3-8b",
    "function_less": "hallisky/lora-function-less-llama-3-8b",
    "grade_more": "hallisky/lora-grade-highschool-llama-3-8b", 
    "grade_less": "hallisky/lora-grade-elementary-llama-3-8b",
}

# Global variable to store the latest obfuscation result
latest_obfuscation = {}
user_id = str(uuid4())  # Generate a unique session-specific user ID

JSON_DATASET_DIR = Path("json_dataset")
JSON_DATASET_DIR.mkdir(parents=True, exist_ok=True)

JSON_DATASET_PATH = JSON_DATASET_DIR / f"train-{user_id}.json"

scheduler = CommitScheduler(
    repo_id="authorship-obfuscation-demo-data",
    repo_type="dataset",
    folder_path=JSON_DATASET_DIR,
    path_in_repo="data",
    every=0.5
)

@spaces.GPU
def temp(text):
    response = tokenizer(text, return_tensors="pt")
    return response

def save_data(data):
    with scheduler.lock:
        with JSON_DATASET_PATH.open("a") as f:
            json.dump(data, f)
            f.write("\n")

def save_feedback(feedback_rating, feedback_text):
    global latest_obfuscation
    if latest_obfuscation:
        feedback_data = latest_obfuscation.copy()
        feedback_data['feedback'] = {
            "rating": feedback_rating,
            "text": feedback_text
        }
        save_data(feedback_data)
    return "No Feedback Selected", ""

def greet(input_text, length, function_words, grade_level, sarcasm, formality, voice, persuasive, descriptive, narrative, expository):
    global latest_obfuscation, user_id
    current_time = datetime.now().isoformat()
    
    response = (
        f"Hello!\n"
        f"Input Text: {input_text}\n"
        f"Length: {length}\n"
        f"Function Words: {function_words}\n"
        f"Grade Level: {grade_level}\n"
        f"Sarcasm: {sarcasm}\n"
        f"Formality: {formality}\n"
        f"Voice: {voice}\n"
        f"Persuasive: {persuasive}\n"
        f"Descriptive: {descriptive}\n"
        f"Narrative: {narrative}\n"
        f"Expository: {expository}"
    )
    
    # Save the new obfuscation result and reset feedback
    latest_obfuscation = {
        "datetime": current_time,
        "user_id": user_id,
        "input_text": input_text,
        "sliders": {
            "length": length,
            "function_words": function_words,
            "grade_level": grade_level,
            "sarcasm": sarcasm,
            "formality": formality,
            "voice": voice,
            "persuasive": persuasive,
            "descriptive": descriptive,
            "narrative": narrative,
            "expository": expository
        },
        "output": response,
        "feedback": {
            "rating": "No Feedback Selected",
            "text": ""
        }
    }
    
    # Save the obfuscation result
    save_data(latest_obfuscation)
    
    return response, gr.update(interactive=True), gr.update(interactive=True)

def reset_sliders():
    return [0.5] * 7 + [0] * 3

def toggle_slider(checked, value):
    if checked:
        return gr.update(value=value, interactive=True)
    else:
        return gr.update(value=0, interactive=False)

def reset_writing_type_sliders(selected_type):
    reset_values = [gr.update(value=0, interactive=False) for _ in range(4)]
    if selected_type != "None":
        index = ["Persuasive", "Descriptive", "Narrative", "Expository"].index(selected_type)
        reset_values[index] = gr.update(value=0, interactive=True)
    return reset_values

def update_save_feedback_button(feedback_rating, feedback_text):
    if feedback_rating != "No Feedback Selected" or feedback_text.strip() != "":
        return gr.update(interactive=True), gr.update(visible=False)
    else:
        return gr.update(interactive=False), gr.update(visible=True)

def update_obfuscate_button(input_text):
    if input_text.strip() == "":
        return gr.update(interactive=False), gr.update(visible=True)
    else:
        return gr.update(interactive=True), gr.update(visible=False)

def check_initial_feedback_state(feedback_rating, feedback_text):
    return update_save_feedback_button(feedback_rating, feedback_text)

demo = gr.Blocks()

with demo:
    with gr.Row():
        with gr.Column(variant="panel"):
            gr.Markdown("# 1) Input Text\n### Enter the text to be obfuscated.")
            input_text = gr.Textbox(
                label="Input Text",
                placeholder="The quick brown fox jumped over the lazy dogs."
            )
            gr.Markdown("# 2) Style Element Sliders\n### Adjust the style element sliders to the desired levels to steer the obfuscation.")
            
            reset_button = gr.Button("Choose slider values automatically (based on input text)")
            
            sliders = []
            slider_values = [
                ("Length (Shorter \u2192 Longer)", -1, 1, 0),
                ("Function Words (Fewer \u2192 More)", -1, 1, 0),
                ("Grade Level (Lower \u2192 Higher)", -1, 1, 0),
                ("Formality (Less \u2192 More)", -1, 1, 0),
                ("Sarcasm (Less \u2192 More)", -1, 1, 0),
                ("Voice (Passive \u2192 Active)", -1, 1, 0),
                ("Writing Type: Persuasive (None \u2192 More)", 0, 1, 0),
                ("Writing Type: Descriptive (None \u2192 More)", 0, 1, 0),
                ("Writing Type: Narrative (None \u2192 More)", 0, 1, 0),
                ("Writing Type: Expository (None \u2192 More)", 0, 1, 0)
            ]
            
            non_writing_type_sliders = []
            writing_type_sliders = []
            
            for idx, (label, min_val, max_val, default) in enumerate(slider_values):
                if "Writing Type" not in label:
                    with gr.Row():
                        checkbox = gr.Checkbox(label=label)
                        slider = gr.Slider(label=label, minimum=min_val, maximum=max_val, step=0.01, value=default, interactive=False)
                        checkbox.change(fn=toggle_slider, inputs=[checkbox, gr.State(default)], outputs=slider)
                        non_writing_type_sliders.append(slider)
                        sliders.append(slider)
            
            writing_type_radio = gr.Radio(
                label="Writing Type",
                choices=["None", "Persuasive", "Descriptive", "Narrative", "Expository"],
                value="None"
            )
            
            writing_type_radio.change(fn=reset_writing_type_sliders, inputs=writing_type_radio, outputs=writing_type_sliders)
            
            for idx, (label, min_val, max_val, default) in enumerate(slider_values):
                if "Writing Type" in label:
                    with gr.Row():
                        slider = gr.Slider(label=label, minimum=min_val, maximum=max_val, step=0.01, value=default, interactive=False)
                        writing_type_sliders.append(slider)
                        sliders.append(slider)
            
            obfuscate_button = gr.Button("Obfuscate Text", interactive=False)
            warning_message = gr.Markdown(
                "<div style='text-align: center; color: red;'>⚠️ Please enter text before obfuscating. ⚠️</div>", visible=True
            )
            
            reset_button.click(fn=reset_sliders, inputs=[], outputs=sliders)
            input_text.change(fn=update_obfuscate_button, inputs=input_text, outputs=[obfuscate_button, warning_message])
            # Initialize the button and warning message state on page load
            demo.load(fn=update_obfuscate_button, inputs=input_text, outputs=[obfuscate_button, warning_message])

        with gr.Column(variant="panel"):
            gr.Markdown("# 3) Obfuscated Output")

            output = gr.Textbox(label="Output", lines=3)
            
            gr.Markdown("## Feedback [Optional]")

            # Add thumbs up / thumbs down
            gr.Markdown("### Is the response good or bad?")
            feedback_rating = gr.Radio(choices=["No Feedback Selected", "Good 👍", "Bad 👎"], value="No Feedback Selected", interactive=False, label="Rate the Response")
            
            # Add feedback box
            gr.Markdown("### Provide any feedback on the obfuscation")
            feedback_text = gr.Textbox(label="Feedback", lines=3, interactive=False)

            obfuscate_button.click(fn=greet, inputs=[input_text] + sliders, outputs=[output, feedback_rating, feedback_text])

            save_feedback_button = gr.Button("Save Feedback", interactive=False)
            feedback_warning_message = gr.Markdown(
                "<div style='text-align: center; color: red;'>⚠️ Please provide feedback or a rating before submitting. ⚠️</div>", visible=True
            )
            
            # Update the interactivity of the save_feedback_button based on feedback_rating and feedback_text
            feedback_rating.change(fn=update_save_feedback_button, inputs=[feedback_rating, feedback_text], outputs=[save_feedback_button, feedback_warning_message])
            feedback_text.change(fn=update_save_feedback_button, inputs=[feedback_rating, feedback_text], outputs=[save_feedback_button, feedback_warning_message])

            save_feedback_button.click(
                fn=save_feedback,
                inputs=[feedback_rating, feedback_text],
                outputs=[feedback_rating, feedback_text]
            )
            # Initialize the save feedback button and warning message state on page load
            demo.load(fn=check_initial_feedback_state, inputs=[feedback_rating, feedback_text], outputs=[save_feedback_button, feedback_warning_message])

demo.launch()