import tkinter as tk
from tkinter import ttk, messagebox, filedialog
import yaml
import os

class HiveStudio:
    def __init__(self, root):
        self.root = root
        self.root.title("Hive Studio Lite")

        self.chunks = []

        # Task metadata
        self.task_name_var = tk.StringVar()
        self.capability_var = tk.StringVar(value="display")
        self.execution_mode_var = tk.StringVar(value="concurrent")

        self.setup_ui()

    def setup_ui(self):
        tk.Label(self.root, text="Task Name").grid(row=0, column=0)
        tk.Entry(self.root, textvariable=self.task_name_var).grid(row=0, column=1)

        tk.Label(self.root, text="Capability").grid(row=1, column=0)
        self.capability_box = ttk.Combobox(self.root, textvariable=self.capability_var)
        self.capability_box.grid(row=1, column=1)
        self.capability_box['values'] = ["generic"]

        tk.Button(self.root, text="Manage Capabilities", command=self.manage_capabilities).grid(row=1, column=2, padx=5)

        tk.Label(self.root, text="Execution Mode").grid(row=2, column=0)
        ttk.Combobox(self.root, textvariable=self.execution_mode_var, values=["concurrent", "sequential"]).grid(row=2, column=1)

        # Chunk controls
        self.chunk_frame = tk.LabelFrame(self.root, text="Chunks")
        self.chunk_frame.grid(row=3, column=0, columnspan=2, pady=10, sticky="ew")

        self.chunk_listbox = tk.Listbox(self.chunk_frame, width=50)
        self.chunk_listbox.pack()
        self.chunk_listbox.bind('<Double-1>', self.edit_chunk)

        tk.Button(self.root, text="Add Chunk", command=self.add_chunk).grid(row=4, column=0, sticky="ew")
        tk.Button(self.root, text="Remove Chunk", command=self.remove_chunk).grid(row=4, column=1, sticky="ew")
        tk.Button(self.root, text="Move Up", command=self.move_chunk_up).grid(row=5, column=0, sticky="ew")
        tk.Button(self.root, text="Move Down", command=self.move_chunk_down).grid(row=5, column=1, sticky="ew")
        tk.Button(self.root, text="Load Task", command=self.load_task).grid(row=6, column=0, sticky="ew")
        tk.Button(self.root, text="Save Task", command=self.save_task).grid(row=6, column=1, sticky="ew")
        tk.Button(self.root, text="Quit", command=self.root.quit).grid(row=7, column=0, columnspan=2, sticky="ew")

    def chunk_editor_ui(self, chunk=None, index=None):
        chunk_editor = tk.Toplevel(self.root)
        chunk_editor.title("Edit Chunk" if chunk else "Add Chunk")

        chunk_id = tk.StringVar(value=chunk.get("chunk_id") if chunk else "")
        payload = tk.StringVar(value=chunk.get("payload") if chunk else "")
        ttl = tk.StringVar(value=str(chunk.get("ttl")) if chunk and "ttl" in chunk else "")
        timing = tk.StringVar(value=chunk.get("timing") if chunk else "")
        depends_on = tk.StringVar(value=chunk.get("depends_on") if chunk else "")
        
        # Chunk capability dropdown (only shown if task-level is a list)
        chunk_capability = tk.StringVar(value=chunk["capability_required"] if chunk and "capability_required" in chunk else "")
        
        import re
        raw_caps = self.capability_var.get()
        capability_values = []

        if "[" in raw_caps or "," in raw_caps:
            try:
                import ast
                parsed = ast.literal_eval(raw_caps)
                if isinstance(parsed, list):
                    capability_values = parsed
            except Exception:
                pass
        else:
            # Try space-separated fallback
            capability_values = re.split(r"[ ,]+", raw_caps.strip())

        # Clean up empty entries
        capability_values = [c for c in capability_values if c]

        
        if capability_values:
            tk.Label(chunk_editor, text="Capability").grid(row=6, column=0)
            ttk.Combobox(chunk_editor, textvariable=chunk_capability, values=capability_values).grid(row=6, column=1)


        tk.Label(chunk_editor, text="Chunk ID").grid(row=0, column=0)
        tk.Entry(chunk_editor, textvariable=chunk_id).grid(row=0, column=1)

        tk.Label(chunk_editor, text="Payload").grid(row=1, column=0)
        tk.Entry(chunk_editor, textvariable=payload).grid(row=1, column=1)

        tk.Label(chunk_editor, text="TTL").grid(row=2, column=0)
        tk.Entry(chunk_editor, textvariable=ttl).grid(row=2, column=1)

        tk.Label(chunk_editor, text="Timing (optional)").grid(row=3, column=0)
        ttk.Combobox(chunk_editor, textvariable=timing, values=["", "seconds", "ms"]).grid(row=3, column=1)

        tk.Label(chunk_editor, text="Depends On (optional)").grid(row=4, column=0)
        ttk.Combobox(chunk_editor, textvariable=depends_on, values=[c['chunk_id'] for c in self.chunks if not chunk or c['chunk_id'] != chunk.get('chunk_id')]).grid(row=4, column=1)

        def confirm():
            updated_chunk = {
                "chunk_id": chunk_id.get(),
                "payload": payload.get()
            }
            if ttl.get():
                updated_chunk["ttl"] = ttl.get()
            if timing.get():
                updated_chunk["timing"] = timing.get()
            if depends_on.get():
                updated_chunk["depends_on"] = depends_on.get()
                
            if capability_values:
                if not chunk_capability.get():
                    messagebox.showerror("Missing Capability", "Each chunk must declare a capability.")
                    return
                updated_chunk["capability_required"] = chunk_capability.get()


            if chunk and index is not None:
                self.chunks[index] = updated_chunk
                self.refresh_chunk_list()
            else:
                self.chunks.append(updated_chunk)
                self.refresh_chunk_list()
            chunk_editor.destroy()

        tk.Button(chunk_editor, text="OK", command=confirm).grid(row=7, column=0, columnspan=2)
        
    def manage_capabilities(self):
        import ast

        def parse_caps():
            raw = self.capability_var.get()
            try:
                return ast.literal_eval(raw) if isinstance(ast.literal_eval(raw), list) else [raw]
            except:
                import re
                return list(filter(None, re.split(r"[ ,]+", raw.strip())))

        editor = tk.Toplevel(self.root)
        editor.title("Manage Capabilities")

        cap_list = parse_caps()
        cap_var = tk.StringVar()
        capbox = tk.Listbox(editor, height=6, width=25)
        capbox.grid(row=0, column=0, columnspan=2, padx=10, pady=5)
        for cap in cap_list:
            capbox.insert(tk.END, cap)

        tk.Entry(editor, textvariable=cap_var).grid(row=1, column=0, padx=5, pady=5)

        def add_cap():
            cap = cap_var.get().strip()
            if cap and cap not in cap_list:
                cap_list.append(cap)
                capbox.insert(tk.END, cap)
                cap_var.set("")

        def remove_cap():
            selection = capbox.curselection()
            if selection:
                index = selection[0]
                cap_list.pop(index)
                capbox.delete(index)

        def apply_caps():
            if len(cap_list) == 1:
                self.capability_var.set(cap_list[0])
            else:
                self.capability_var.set(str(cap_list))
            self.capability_box['values'] = cap_list
            editor.destroy()

        tk.Button(editor, text="Add", command=add_cap).grid(row=1, column=1, padx=5)
        tk.Button(editor, text="Remove Selected", command=remove_cap).grid(row=2, column=0, columnspan=2, pady=5)
        tk.Button(editor, text="Apply", command=apply_caps).grid(row=3, column=0, columnspan=2, pady=10)


    def refresh_chunk_list(self):
        self.chunk_listbox.delete(0, tk.END)
        for chunk in self.chunks:
            label = f"{chunk['chunk_id']} -> {chunk['payload']}"
            if 'ttl' in chunk:
                label += f" [TTL: {chunk['ttl']}]"
            if 'depends_on' in chunk:
                label += f" [Depends: {chunk['depends_on']}]"
            if 'capability_required' in chunk:
                label += f" [Cap: {chunk['capability_required']}]"

            self.chunk_listbox.insert(tk.END, label)

    def add_chunk(self):
        self.chunk_editor_ui()

    def edit_chunk(self, event):
        selection = self.chunk_listbox.curselection()
        if selection:
            index = selection[0]
            self.chunk_editor_ui(chunk=self.chunks[index], index=index)

    def remove_chunk(self):
        selected = self.chunk_listbox.curselection()
        if selected:
            index = selected[0]
            self.chunk_listbox.delete(index)
            del self.chunks[index]

    def move_chunk_up(self):
        selection = self.chunk_listbox.curselection()
        if selection:
            index = selection[0]
            if index > 0:
                self.chunks[index - 1], self.chunks[index] = self.chunks[index], self.chunks[index - 1]
                self.refresh_chunk_list()
                self.chunk_listbox.select_set(index - 1)

    def move_chunk_down(self):
        selection = self.chunk_listbox.curselection()
        if selection:
            index = selection[0]
            if index < len(self.chunks) - 1:
                self.chunks[index + 1], self.chunks[index] = self.chunks[index], self.chunks[index + 1]
                self.refresh_chunk_list()
                self.chunk_listbox.select_set(index + 1)

    def load_task(self):
        import ast
        path = filedialog.askopenfilename(title="Load Task YAML", filetypes=[("YAML Files", "*.yaml")])
        if path:
            with open(path, "r") as f:
                task = yaml.safe_load(f)

            self.task_name_var.set(task.get("task_name", ""))
            self.execution_mode_var.set(task.get("execution_mode", "concurrent"))
            self.chunks = task.get("chunks", [])

            # Normalize capabilities
            caps = task.get("capability_required", "display")
            if isinstance(caps, list):
                self.capability_var.set(str(caps))  # This will trigger parsing logic later
            else:
                self.capability_var.set(caps)
                
            # Update dropdown values dynamically
            if isinstance(caps, list):
                self.capability_box['values'] = caps
            else:
                self.capability_box['values'] = [caps]


            self.refresh_chunk_list()


    def save_task(self):
        import re
        import ast

        # Parse task-level capabilities
        raw_caps = self.capability_var.get()
        capability_values = []

        try:
            parsed = ast.literal_eval(raw_caps)
            if isinstance(parsed, list):
                capability_values = parsed
            else:
                raise ValueError
        except Exception:
            capability_values = re.split(r"[ ,]+", raw_caps.strip())

        capability_values = [c for c in capability_values if c]
        cap_obj = capability_values[0] if len(capability_values) == 1 else capability_values

        # Sanitize chunks: ensure 'required_capability' exists and is a string
        for chunk in self.chunks:
            # Migrate field if still using old name
            if "capability_required" in chunk:
                chunk["required_capability"] = chunk.pop("capability_required")

            cap = chunk.get("required_capability")

            if isinstance(cap, list):
                chunk["required_capability"] = cap[0]

            elif cap is None:
                if isinstance(cap_obj, list):
                    messagebox.showerror("Missing Capability", f"Chunk '{chunk.get('chunk_id', '?')}' is missing a capability.")
                    return
                elif isinstance(cap_obj, str):
                    chunk["required_capability"] = cap_obj

        task = {
            "task_name": self.task_name_var.get(),
            "capability_required": cap_obj,
            "execution_mode": self.execution_mode_var.get(),
            "chunks": self.chunks
        }

        script_dir = os.path.dirname(os.path.abspath(__file__))
        save_dir = os.path.join(script_dir, "..")
        os.makedirs(save_dir, exist_ok=True)
        path = os.path.join(save_dir, f"{task['task_name']}.yaml")

        with open(path, "w") as f:
            yaml.dump(task, f)

        messagebox.showinfo("Saved", f"Task saved to {path}")


if __name__ == "__main__":
    root = tk.Tk()
    app = HiveStudio(root)
    root.mainloop()
