# target_simulator/gui/trajectory_editor_window.py

import tkinter as tk
from tkinter import ttk, messagebox
from typing import List, Optional, cast
import copy
from queue import Queue

# Use absolute imports
from target_simulator.core.models import (
    Target,
    Waypoint,
    ManeuverType,
    TurnDirection,
    MIN_TARGET_ID,
    MAX_TARGET_ID,
    FPS_TO_KNOTS,
    Scenario,
)
from target_simulator.gui.ppi_display import PPIDisplay
from target_simulator.gui.waypoint_editor_window import WaypointEditorWindow
from target_simulator.core.simulation_engine import SimulationEngine

GUI_QUEUE_POLL_INTERVAL_MS = 50


class TrajectoryEditorWindow(tk.Toplevel):
    """A dialog for managing a target's trajectory with a visual preview."""

    def __init__(
        self,
        master,
        existing_ids: List[int],
        max_range_nm: int,
        target_to_edit: Optional[Target] = None,
        scenario_name: Optional[str] = None,
    ):
        super().__init__(master)
        self.transient(master)
        self.grab_set()

        self.existing_ids = existing_ids
        self.result_target: Optional[Target] = None
        self.initial_max_range = max_range_nm
        self.time_multiplier = 1.0
        self.total_sim_time = 0.0

        self.waypoints: List[Waypoint] = []
        is_editing = target_to_edit is not None
        if is_editing:
            target = cast(Target, target_to_edit)
            self.target_id = target.target_id
            self.waypoints = copy.deepcopy(target.trajectory)
            self.use_spline_var = tk.BooleanVar(value=target.use_spline)
        else:
            self.target_id = next(
                (
                    i
                    for i in range(MIN_TARGET_ID, MAX_TARGET_ID + 1)
                    if i not in existing_ids
                ),
                -1,
            )
            self.use_spline_var = tk.BooleanVar(value=False)

        scenario_name_str = scenario_name or f"Target_{self.target_id}"
        self.title(f"Trajectory Editor - {scenario_name_str}")
        self.geometry("1100x700")

        self.gui_update_queue: Queue = Queue()
        self.preview_engine: Optional[SimulationEngine] = None
        self.is_preview_running = tk.BooleanVar(value=False)
        self.is_paused = tk.BooleanVar(value=False)

        self._create_widgets()

        if not self.waypoints:
            self._create_initial_waypoint()

        self._populate_waypoint_list()
        self.after(10, self._update_static_preview)

        self.protocol("WM_DELETE_WINDOW", self._on_cancel)
        self.wait_window(self)

    def _create_widgets(self):
        main_pane = ttk.PanedWindow(self, orient=tk.HORIZONTAL)
        main_pane.pack(fill=tk.BOTH, expand=True, padx=10, pady=10)

        left_frame = ttk.Frame(main_pane)
        main_pane.add(left_frame, weight=1)

        list_frame = ttk.LabelFrame(left_frame, text="Trajectory Waypoints")
        list_frame.pack(fill=tk.BOTH, expand=True, pady=(0, 5))
        self._create_waypoint_list_widgets(list_frame)

        button_frame = ttk.Frame(left_frame)
        button_frame.pack(fill=tk.X, side=tk.BOTTOM, pady=(5, 0))
        ttk.Button(button_frame, text="OK", command=self._on_ok).pack(
            side=tk.RIGHT, padx=5
        )
        ttk.Button(button_frame, text="Cancel", command=self._on_cancel).pack(
            side=tk.RIGHT
        )

        preview_frame = ttk.LabelFrame(main_pane, text="Trajectory Preview")
        main_pane.add(preview_frame, weight=2)
        self._create_preview_widgets(preview_frame)

    def _create_waypoint_list_widgets(self, parent):
        self.wp_tree = ttk.Treeview(
            parent, columns=("num", "type", "details"), show="headings", height=10
        )
        self.wp_tree.heading("num", text="#")
        self.wp_tree.heading("type", text="Maneuver")
        self.wp_tree.heading("details", text="Parameters")
        self.wp_tree.column("num", width=30, anchor=tk.CENTER, stretch=False)
        self.wp_tree.column("type", width=120, stretch=False)
        self.wp_tree.column("details", width=250)

        scrollbar = ttk.Scrollbar(
            parent, orient=tk.VERTICAL, command=self.wp_tree.yview
        )
        self.wp_tree.configure(yscrollcommand=scrollbar.set)

        self.wp_tree.grid(row=0, column=0, sticky="nsew")
        scrollbar.grid(row=0, column=1, sticky="ns")
        parent.grid_rowconfigure(0, weight=1)
        parent.grid_columnconfigure(0, weight=1)

        self.wp_tree.bind("<Double-1>", self._on_edit_waypoint)

        btn_frame = ttk.Frame(parent)
        btn_frame.grid(row=1, column=0, columnspan=2, sticky="ew", pady=(5, 0))

        ttk.Button(btn_frame, text="Add", command=self._on_add_waypoint).pack(
            side=tk.LEFT, padx=(0, 2)
        )
        ttk.Button(btn_frame, text="Edit", command=self._on_edit_waypoint).pack(
            side=tk.LEFT, padx=2
        )
        ttk.Button(btn_frame, text="Remove", command=self._on_remove_waypoint).pack(
            side=tk.LEFT, padx=2
        )

        self.spline_checkbox = ttk.Checkbutton(
            btn_frame,
            text="Use Spline",
            variable=self.use_spline_var,
            command=self._on_spline_toggle,
        )
        self.spline_checkbox.pack(side=tk.RIGHT, padx=5)

    def _create_preview_widgets(self, parent):
        self.ppi_preview = PPIDisplay(parent, max_range_nm=self.initial_max_range)
        self.ppi_preview.pack(fill=tk.BOTH, expand=True)

        controls_frame = ttk.Frame(parent)
        controls_frame.pack(fill=tk.X, pady=5)

        self.play_button = ttk.Button(
            controls_frame, text="▶ Play", command=self._on_preview_play
        )
        self.play_button.pack(side=tk.LEFT, padx=5)
        self.pause_button = ttk.Button(
            controls_frame,
            text="⏸ Pause",
            command=self._on_preview_pause,
            state=tk.DISABLED,
        )
        self.pause_button.pack(side=tk.LEFT, padx=5)
        self.stop_button = ttk.Button(
            controls_frame,
            text="■ Stop",
            command=self._on_preview_stop,
            state=tk.DISABLED,
        )
        self.stop_button.pack(side=tk.LEFT, padx=5)

        ttk.Label(controls_frame, text="Speed:").pack(side=tk.LEFT, padx=(10, 2))
        self.time_multiplier_var = tk.StringVar(value="1x")
        self.multiplier_combo = ttk.Combobox(
            controls_frame,
            textvariable=self.time_multiplier_var,
            values=["1x", "2x", "4x", "10x", "20x"],
            state="readonly",
            width=4,
        )
        self.multiplier_combo.pack(side=tk.LEFT)
        self.multiplier_combo.bind(
            "<<ComboboxSelected>>", self._on_time_multiplier_changed
        )

        progress_frame = ttk.Frame(parent)
        progress_frame.pack(fill=tk.X, padx=5, pady=(0, 5))
        self.sim_progress_var = tk.DoubleVar(value=0.0)
        self.sim_progress = ttk.Scale(
            progress_frame,
            variable=self.sim_progress_var,
            from_=0.0,
            to=1.0,
            orient=tk.HORIZONTAL,
        )
        self.sim_progress.pack(side=tk.LEFT, fill=tk.X, expand=True)
        self.sim_progress.bind("<ButtonRelease-1>", self._on_progress_seek)
        self.sim_time_label = ttk.Label(progress_frame, text="0.0s / 0.0s", width=15)
        self.sim_time_label.pack(side=tk.LEFT, padx=(5, 0))

    def _populate_waypoint_list(self):
        self.wp_tree.delete(*self.wp_tree.get_children())
        for i, wp in enumerate(self.waypoints):
            details_parts = []
            if wp.maneuver_type == ManeuverType.FLY_TO_POINT:
                details_parts.append(f"R:{wp.target_range_nm:.1f} NM")
                details_parts.append(f"Az:{wp.target_azimuth_deg:.1f}°")
                details_parts.append(f"Alt:{wp.target_altitude_ft:.0f} ft")

            elif wp.maneuver_type == ManeuverType.FLY_FOR_DURATION:
                vel_kn = (wp.target_velocity_fps or 0.0) * FPS_TO_KNOTS
                details_parts.append(f"Vel:{vel_kn:.1f} kn")
                details_parts.append(f"Hdg:{wp.target_heading_deg:.1f}°")
                details_parts.append(f"T:{wp.duration_s:.1f}s")

            elif wp.maneuver_type == ManeuverType.DYNAMIC_MANEUVER:
                if wp.longitudinal_acceleration_g != 0:
                    details_parts.append(f"Lon G:{wp.longitudinal_acceleration_g:+.1f}")
                if wp.lateral_acceleration_g != 0:
                    dir = "R" if wp.turn_direction == TurnDirection.RIGHT else "L"
                    details_parts.append(f"Lat G:{wp.lateral_acceleration_g:.1f} {dir}")
                if wp.vertical_acceleration_g != 0:
                    details_parts.append(f"Ver G:{wp.vertical_acceleration_g:+.1f}")
                details_parts.append(f"T:{wp.duration_s:.1f}s")

            details = ", ".join(details_parts)
            self.wp_tree.insert(
                "", tk.END, iid=str(i), values=(i + 1, wp.maneuver_type.value, details)
            )

    def _get_total_simulation_time(self) -> float:
        if not self.waypoints:
            return 0.0
        _, total_duration = Target.generate_path_from_waypoints(
            self.waypoints, self.use_spline_var.get()
        )
        return total_duration

    def _update_static_preview(self):
        if self.is_preview_running.get():
            self._on_preview_stop()

        self.ppi_preview.draw_trajectory_preview(
            waypoints=copy.deepcopy(self.waypoints),
            use_spline=self.use_spline_var.get(),
        )
        self.total_sim_time = self._get_total_simulation_time()
        self.sim_progress_var.set(0.0)
        self.sim_time_label.config(text=f"0.0s / {self.total_sim_time:.1f}s")

    def _on_preview_play(self):
        if (
            self.is_preview_running.get()
            and self.is_paused.get()
            and self.preview_engine
        ):
            self.is_paused.set(False)
            self.preview_engine.resume()
            self._update_preview_button_states()
            return

        if self.is_preview_running.get():
            return
        if (
            not self.waypoints
            or self.waypoints[0].maneuver_type != ManeuverType.FLY_TO_POINT
        ):
            messagebox.showinfo(
                "Incomplete Trajectory",
                "First waypoint must be 'Fly to Point'.",
                parent=self,
            )
            return

        self.is_preview_running.set(True)
        self._update_preview_button_states()

        preview_scenario = Scenario(name=f"Preview_{self.target_id}")
        preview_scenario.add_target(
            Target(
                target_id=self.target_id,
                trajectory=copy.deepcopy(self.waypoints),
                use_spline=self.use_spline_var.get(),
            )
        )
        self.preview_engine = SimulationEngine(
            communicator=None, update_queue=self.gui_update_queue
        )
        self.preview_engine.set_time_multiplier(self.time_multiplier)
        self.preview_engine.load_scenario(preview_scenario)
        self.preview_engine.start()
        self.after(GUI_QUEUE_POLL_INTERVAL_MS, self._process_preview_queue)

    def _on_preview_pause(self):
        if (
            self.preview_engine
            and self.is_preview_running.get()
            and not self.is_paused.get()
        ):
            self.is_paused.set(True)
            self.preview_engine.pause()
            self._update_preview_button_states()

    def _on_preview_stop(self):
        if not self.is_preview_running.get() or not self.preview_engine:
            return

        self.preview_engine.stop()
        self.preview_engine = None
        self.is_preview_running.set(False)
        self.is_paused.set(False)
        self._update_preview_button_states()
        self._update_static_preview()

    def _process_preview_queue(self):
        try:
            while not self.gui_update_queue.empty():
                update = self.gui_update_queue.get_nowait()
                if update == "SIMULATION_FINISHED":
                    self._on_preview_stop()
                    self.sim_progress_var.set(1.0)
                    self.sim_time_label.config(
                        text=f"{self.total_sim_time:.1f}s / {self.total_sim_time:.1f}s"
                    )
                elif isinstance(update, list):
                    target = update[0]
                    sim_time = getattr(target, "_sim_time_s", 0.0)
                    if self.total_sim_time > 0:
                        progress = min(sim_time / self.total_sim_time, 1.0)
                        self.sim_progress_var.set(progress)
                    self.sim_time_label.config(
                        text=f"{sim_time:.1f}s / {self.total_sim_time:.1f}s"
                    )
                    self.ppi_preview.update_targets(update)
        finally:
            if self.is_preview_running.get():
                self.after(GUI_QUEUE_POLL_INTERVAL_MS, self._process_preview_queue)

    def _on_progress_seek(self, event):
        if not self.preview_engine:
            return
        seek_value = self.sim_progress_var.get()
        seek_time = seek_value * self.total_sim_time
        self.sim_time_label.config(
            text=f"{seek_time:.1f}s / {self.total_sim_time:.1f}s"
        )
        self.preview_engine.set_simulation_time(seek_time)

    def _update_preview_button_states(self):
        running = self.is_preview_running.get()
        paused = self.is_paused.get()

        if not running:
            self.play_button.config(state=tk.NORMAL, text="▶ Play")
        elif paused:
            self.play_button.config(state=tk.NORMAL, text="▶ Resume")
        else:
            self.play_button.config(state=tk.DISABLED, text="▶ Play")

        self.pause_button.config(
            state=tk.NORMAL if (running and not paused) else tk.DISABLED
        )
        self.stop_button.config(state=tk.NORMAL if running else tk.DISABLED)
        self.multiplier_combo.config(state="readonly" if not running else tk.DISABLED)

        is_editable = not running
        self.spline_checkbox.config(state=tk.NORMAL if is_editable else tk.DISABLED)
        btn_frame = self.wp_tree.master.children.get("!frame2")
        if btn_frame:
            for child in btn_frame.winfo_children():
                if isinstance(child, (ttk.Button, ttk.Checkbutton)):
                    child.config(state=tk.NORMAL if is_editable else tk.DISABLED)
        self.wp_tree.config(selectmode="browse" if is_editable else "none")

    def _on_spline_toggle(self):
        if self.use_spline_var.get() and len(self.waypoints) < 4:
            self.use_spline_var.set(False)
            messagebox.showinfo(
                "Spline Not Available",
                "Spline mode requires at least 4 waypoints.",
                parent=self,
            )
            return
        self._update_static_preview()

    def _create_initial_waypoint(self):
        editor = WaypointEditorWindow(self, is_first_waypoint=True)
        if editor.result_waypoint:
            self.waypoints.append(editor.result_waypoint)
        else:
            self.after(10, self._on_cancel)

    def _on_add_waypoint(self):
        editor = WaypointEditorWindow(self, is_first_waypoint=False)
        if editor.result_waypoint:
            self.waypoints.append(editor.result_waypoint)
            self._populate_waypoint_list()
            self._update_static_preview()
            new_item_id = str(len(self.waypoints) - 1)
            self.wp_tree.see(new_item_id)
            self.wp_tree.focus(new_item_id)
            self.wp_tree.selection_set(new_item_id)

    def _on_edit_waypoint(self, event=None):
        selected_item = self.wp_tree.focus()
        if not selected_item:
            if not event:
                messagebox.showwarning(
                    "No Selection", "Please select a waypoint to edit.", parent=self
                )
            return
        wp_index = int(selected_item)
        editor = WaypointEditorWindow(
            self,
            is_first_waypoint=(wp_index == 0),
            waypoint_to_edit=self.waypoints[wp_index],
        )
        if editor.result_waypoint:
            self.waypoints[wp_index] = editor.result_waypoint
            self._populate_waypoint_list()
            self._update_static_preview()
            self.wp_tree.focus(selected_item)
            self.wp_tree.selection_set(selected_item)

    def _on_remove_waypoint(self):
        selected_item = self.wp_tree.focus()
        if not selected_item:
            messagebox.showwarning(
                "No Selection", "Please select a waypoint to remove.", parent=self
            )
            return
        if len(self.waypoints) == 1:
            messagebox.showerror(
                "Action Denied",
                "Cannot remove the initial position waypoint.",
                parent=self,
            )
            return
        wp_index = int(selected_item)
        del self.waypoints[wp_index]
        self._populate_waypoint_list()
        self._update_static_preview()

    def _on_time_multiplier_changed(self, event=None):
        try:
            self.time_multiplier = float(
                self.time_multiplier_var.get().replace("x", "")
            )
            if self.preview_engine and self.preview_engine.is_running():
                self.preview_engine.set_time_multiplier(self.time_multiplier)
        except ValueError:
            self.time_multiplier = 1.0

    def _on_ok(self):
        if (
            not self.waypoints
            or self.waypoints[0].maneuver_type != ManeuverType.FLY_TO_POINT
        ):
            messagebox.showerror(
                "Invalid Trajectory",
                "The first waypoint must be 'Fly to Point'.",
                parent=self,
            )
            return
        try:
            self.result_target = Target(
                target_id=self.target_id,
                trajectory=self.waypoints,
                use_spline=self.use_spline_var.get(),
            )
            self.destroy()
        except (ValueError, tk.TclError) as e:
            messagebox.showerror("Validation Error", str(e), parent=self)

    def _on_cancel(self):
        if self.is_preview_running.get():
            self._on_preview_stop()
        self.result_target = None
        self.destroy()
