main.py

import tkinter as tk
from tkinter import messagebox, filedialog, scrolledtext
import json


class TextAnnotator:
    def __init__(self):
        self.root = tk.Tk()
        self.root.title("Text Annotator")
        self.root.geometry("800x600")

        self.text = ""
        self.selections = []
        self.entries = []
        self.colors = [
            "#ffeb3b",
            "#e91e63",
            "#9c27b0",
            "#673ab7",
            "#3f51b5",
            "#2196f3",
            "#03a9f4",
            "#00bcd4",
            "#009688",
            "#4caf50",
            "#8bc34a",
            "#cddc39",
            "#ffc107",
            "#ff9800",
            "#ff5722",
        ]

        self.setup_ui()

    def setup_ui(self):
        # Input
        frame = tk.Frame(self.root)
        frame.pack(fill=tk.X, padx=10, pady=5)

        tk.Label(frame, text="Text:", font=("Arial", 12, "bold")).pack(side=tk.LEFT)
        self.entry = tk.Entry(frame, font=("Arial", 11))
        self.entry.pack(side=tk.LEFT, fill=tk.X, expand=True, padx=5)
        self.entry.bind("<Return>", lambda e: self.load_text())

        tk.Button(
            frame,
            text="Load",
            command=self.load_text,
            bg="#4CAF50",
            fg="black",
            font=("Arial", 10, "bold"),
        ).pack(side=tk.RIGHT)

        # Split view
        paned = tk.PanedWindow(self.root, orient=tk.VERTICAL, sashwidth=5)
        paned.pack(fill=tk.BOTH, expand=True, padx=10, pady=5)

        # Text display
        self.text_widget = scrolledtext.ScrolledText(
            paned,
            wrap=tk.WORD,
            font=("Arial", 14),
            state=tk.DISABLED,
            cursor="arrow",
            height=15,
        )
        self.text_widget.bind("<Button-1>", self.on_click)
        self.text_widget.bind("<ButtonRelease-1>", self.on_release)
        paned.add(self.text_widget, minsize=200)

        # JSONL preview
        self.preview = scrolledtext.ScrolledText(
            paned,
            wrap=tk.WORD,
            font=("Courier", 10),
            state=tk.DISABLED,
            bg="#f5f5f5",
            height=15,
        )
        paned.add(self.preview, minsize=150)

        # Controls
        controls = tk.Frame(self.root)
        controls.pack(fill=tk.X, padx=10, pady=5)

        buttons = [
            ("New Text", self.new_text, "#2196F3"),
            ("Export JSONL", self.export_jsonl, "#FF9800"),
        ]

        for text, cmd, color in buttons:
            tk.Button(
                controls,
                text=text,
                command=cmd,
                bg=color,
                fg="black",
                font=("Arial", 10, "bold"),
                padx=10,
                pady=5,
            ).pack(side=tk.LEFT, padx=5)

        # Status
        self.status = tk.Label(controls, text="Ready", font=("Arial", 10))
        self.status.pack(side=tk.RIGHT)

    def load_text(self):
        self.text = self.entry.get().strip()
        if not self.text:
            return

        self.selections = []
        self.text_widget.config(state=tk.NORMAL)
        self.text_widget.delete(1.0, tk.END)
        self.text_widget.insert(1.0, self.text)
        self.text_widget.config(state=tk.DISABLED)

        # Clear all tags
        for tag in self.text_widget.tag_names():
            if tag.startswith("sel_"):
                self.text_widget.tag_delete(tag)

        self.update_preview()
        self.status.config(text=f"Loaded: {len(self.text)} chars")

    def on_click(self, event):
        pos = self.text_widget.index(f"@{event.x},{event.y}")
        # Check if clicking on existing selection
        for i, sel in enumerate(self.selections):
            start_idx = self.char_to_index(sel["start"])
            end_idx = self.char_to_index(sel["end"])
            if self.text_widget.compare(
                start_idx, "<=", pos
            ) and self.text_widget.compare(pos, "<", end_idx):
                self.remove_selection(i)
                return

    def on_release(self, event):
        try:
            if self.text_widget.tag_ranges(tk.SEL):
                start_idx = self.text_widget.index(tk.SEL_FIRST)
                end_idx = self.text_widget.index(tk.SEL_LAST)
                start_pos = self.index_to_char(start_idx)
                end_pos = self.index_to_char(end_idx)

                # Check for overlaps
                for sel in self.selections:
                    if start_pos < sel["end"] and end_pos > sel["start"]:
                        self.status.config(text="Overlapping selection rejected")
                        return

                # Add selection
                selected_text = self.text_widget.get(start_idx, end_idx)
                self.selections.append(
                    {"start": start_pos, "end": end_pos, "text": selected_text}
                )

                # Highlight
                tag = f"sel_{len(self.selections)}"
                color = self.colors[(len(self.selections) - 1) % len(self.colors)]
                self.text_widget.tag_configure(tag, background=color)
                self.text_widget.tag_add(tag, start_idx, end_idx)

                self.update_preview()
                self.status.config(text=f"Selected: '{selected_text}'")
        except tk.TclError:
            pass

    def remove_selection(self, index):
        sel = self.selections.pop(index)
        self.text_widget.tag_delete(f"sel_{index + 1}")

        # Recolor remaining selections
        for i, selection in enumerate(self.selections):
            if i >= index:
                old_tag = f"sel_{i + 2}"
                new_tag = f"sel_{i + 1}"
                ranges = self.text_widget.tag_ranges(old_tag)
                if ranges:
                    color = self.colors[i % len(self.colors)]
                    self.text_widget.tag_configure(new_tag, background=color)
                    for j in range(0, len(ranges), 2):
                        self.text_widget.tag_add(new_tag, ranges[j], ranges[j + 1])
                    self.text_widget.tag_delete(old_tag)

        self.update_preview()
        self.status.config(text=f"Removed: '{sel['text']}'")

    def char_to_index(self, char_pos):
        pos, line = 0, 1
        while True:
            line_text = self.text_widget.get(f"{line}.0", f"{line}.end")
            if pos + len(line_text) >= char_pos:
                return f"{line}.{char_pos - pos}"
            pos += len(line_text) + 1
            line += 1
            if line > 1000:
                return "1.0"

    def index_to_char(self, tk_index):
        line, col = map(int, tk_index.split("."))
        pos = sum(
            len(self.text_widget.get(f"{i}.0", f"{i}.end")) + 1 for i in range(1, line)
        )
        return pos + col

    def new_text(self):
        if self.text:
            self.entries.append({"text": self.text, "selections": self.selections})
        self.entry.delete(0, tk.END)
        self.text = ""
        self.selections = []
        self.text_widget.config(state=tk.NORMAL)
        self.text_widget.delete(1.0, tk.END)
        self.text_widget.config(state=tk.DISABLED)
        self.update_preview()
        self.status.config(text="New text ready")

    def export_jsonl(self):
        entries = self.entries.copy()
        if self.text:
            entries.append({"text": self.text, "selections": self.selections})

        if not entries:
            messagebox.showwarning("Warning", "No data to export")
            return

        file_path = filedialog.asksaveasfilename(
            defaultextension=".jsonl", filetypes=[("JSONL files", "*.jsonl")]
        )
        if file_path:
            with open(file_path, "w") as f:
                for entry in entries:
                    f.write(json.dumps(entry) + "\n")
            messagebox.showinfo("Success", f"Exported {len(entries)} entries")

    def update_preview(self):
        entries = self.entries.copy()
        if self.text:
            entries.append({"text": self.text, "selections": self.selections})

        self.preview.config(state=tk.NORMAL)
        self.preview.delete(1.0, tk.END)
        for entry in entries:
            self.preview.insert(tk.END, json.dumps(entry) + "\n")
        self.preview.config(state=tk.DISABLED)

    def run(self):
        self.root.mainloop()


if __name__ == "__main__":
    TextAnnotator().run()