"""Tool for marking samples as valid or invalid. Takes in JSONL and writes a
file of with lines of this form:

    ^(?P<sample_id>[0-9]+): (?P<validity>valid|invalid)$

You do this after running create_{hijacking,extraction}_robustness_dataset."""

import argparse
import json
import tkinter as tk
from pathlib import Path
from tkinter import messagebox, ttk
from typing import Dict, Literal

ValidOrInvalid = Literal["valid", "invalid"]


def load_results_into_dict(results_filename: str | Path) -> Dict[str, ValidOrInvalid]:
    """Parse saved valid/invalid annotations into a dictionary."""
    results_dict: Dict[str, ValidOrInvalid] = {}
    with open(results_filename, "r") as file:
        for line in file:
            sample_id, result = line.strip().split(":")
            result = result.strip()
            assert result == "valid" or result == "invalid"
            # The same annotation may appear twice. This is fine; we just take
            # the later one.
            results_dict[sample_id.strip()] = result
    return results_dict


class SampleReviewer(tk.Tk):
    def __init__(
        self, samples_filename: str | Path, results_filename: str | Path
    ) -> None:
        super().__init__()

        self.title("Sample Reviewer")
        self.geometry("1750x950")

        # Load the samples first
        self.samples = self.load_samples(samples_filename)
        self.results_filename = results_filename
        self.index = 0

        # Create the UI widgets
        self.create_widgets()

        # Load the results into the dictionary
        self.results_dict = self.load_results_into_dict()

        # Setup key bindings
        self.bind("<Left>", lambda _: self.move_prev())
        self.bind("<Right>", lambda _: self.move_next())
        self.bind("v", lambda _: self.mark_sample("valid"))
        self.bind("i", lambda _: self.mark_sample("invalid"))

        # Display the first sample
        self.display_sample()

    def load_results_into_dict(self) -> Dict[str, ValidOrInvalid]:
        return load_results_into_dict(self.results_filename)

    def check_result(self, sample_id):
        return self.results_dict.get(str(sample_id), "unmarked")

    def save_result(self, sample_id, result):
        # Append the result to 'results.txt' file
        with open(self.results_filename, "a") as file:
            file.write(f"{sample_id}: {result}\n")
        # Update the results dictionary
        self.results_dict[str(sample_id)] = result

    @staticmethod
    def load_samples(file_name):
        with open(file_name, "r") as file:
            return [json.loads(line) for line in file]

    def display_sample(self):
        if 0 <= self.index < len(self.samples):
            sample = self.samples[self.index]
            self.update_textbox(self.pre_prompt_textbox, sample["pre_prompt"])
            self.update_textbox(self.access_code_textbox, sample["access_code"])
            self.update_textbox(self.post_prompt_textbox, sample["post_prompt"])
            self.update_textbox(self.attack_textbox, sample["attack"])
            status = self.check_result(sample["sample_id"])
            self.update_status_label(status)
        else:
            messagebox.showinfo("Done", "You've reached the end of the samples.")

    def update_status_label(self, status):
        progress_pct = self.index / max(1, len(self.samples)) * 100
        num_marker = f" ({self.index+1}/{len(self.samples)}, {progress_pct:.2f}% done)"
        if status == "valid":
            self.status_label.config(text=f"Valid {num_marker}", bg="green", fg="white")
        elif status == "invalid":
            self.status_label.config(text=f"Invalid {num_marker}", bg="red", fg="white")
        else:
            self.status_label.config(
                text=f"Unmarked {num_marker}", bg="yellow", fg="black"
            )

    @staticmethod
    def update_textbox(textbox, text):
        textbox.delete("1.0", tk.END)
        textbox.insert(tk.END, text)

    def mark_sample(self, result):
        self.save_result(self.samples[self.index]["sample_id"], result)
        self.move_next()
        self.display_sample()

    def move_next(self):
        self.index += 1
        self.display_sample()

    def move_prev(self):
        self.index -= 1
        self.display_sample()

    def create_labelled_textbox(self, parent, label_text):
        frame = ttk.Frame(parent)

        # Label for the textbox
        label = ttk.Label(frame, text=label_text)
        label.pack(anchor=tk.W)  # Using pack here for simplicity

        # Textbox
        textbox = tk.Text(frame, height=120, width=40, wrap=tk.WORD)
        textbox.pack(fill=tk.BOTH, expand=True)

        return frame, textbox

    def create_widgets(self):
        # Create the status label at the top
        self.status_label = tk.Label(self, font=("Arial", 16))
        self.status_label.grid(row=0, columnspan=4, pady=20, sticky=tk.W + tk.E)

        # Creating labelled text boxes for each section using grid
        frame1, self.pre_prompt_textbox = self.create_labelled_textbox(
            self, "Pre-Prompt"
        )
        frame1.grid(row=1, column=0, sticky=tk.EW)

        frame2, self.access_code_textbox = self.create_labelled_textbox(
            self, "Access Code"
        )
        frame2.grid(row=1, column=1, sticky=tk.EW)

        frame3, self.post_prompt_textbox = self.create_labelled_textbox(
            self, "Post-Prompt"
        )
        frame3.grid(row=1, column=2, sticky=tk.EW)

        frame4, self.attack_textbox = self.create_labelled_textbox(self, "Attack")
        frame4.grid(row=1, column=3, sticky=tk.EW)

        # Buttons
        control_frame = ttk.Frame(self)
        control_frame.grid(row=3, column=0, columnspan=4, pady=20)  # Changed row to 3

        prev_button = ttk.Button(
            control_frame, text="Previous (left)", command=self.move_prev
        )
        prev_button.pack(side=tk.LEFT, padx=10)

        next_button = ttk.Button(
            control_frame, text="Next (right)", command=self.move_next
        )
        next_button.pack(side=tk.LEFT, padx=10)

        valid_button = ttk.Button(
            control_frame, text="Valid (V)", command=lambda: self.mark_sample("valid")
        )
        valid_button.pack(side=tk.LEFT, padx=10)

        invalid_button = ttk.Button(
            control_frame,
            text="Invalid (I)",
            command=lambda: self.mark_sample("invalid"),
        )
        invalid_button.pack(side=tk.LEFT, padx=10)

        # Adjust the grid weights for rows and columns to make components expand properly.
        self.grid_columnconfigure(0, weight=1)
        self.grid_columnconfigure(1, weight=1)
        self.grid_columnconfigure(2, weight=1)
        self.grid_columnconfigure(3, weight=1)
        self.grid_rowconfigure(
            0, weight=0
        )  # Ensure the status label doesn't expand vertically
        self.grid_rowconfigure(1, weight=1)
        self.grid_rowconfigure(3, weight=0)


# Example showing usage
if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Start the SampleReviewer application."
    )

    # Add arguments for the samples file and the results file
    parser.add_argument(
        "samples_file", help="Path to the samples file (e.g., your_filename.jsonl)."
    )
    parser.add_argument(
        "results_file", help="Path to the results file (e.g., results.txt)."
    )

    # Parse the arguments
    args = parser.parse_args()

    # Create the reviewer instance with the provided arguments and start the main loop
    reviewer = SampleReviewer(args.samples_file, args.results_file)
    reviewer.mainloop()
