#!/usr/bin/env python3

import tkinter as tk
from tkinter import ttk
from plotter import Plotter
from plotterStepperMakerBackend import PlotterBackend
import argparse
import time
import os
import json
import sys


class PlotterGUI:
    def __init__(self, root, save_file):
        self.root = root
        self.root.title("Plotter Stepper Maker")
        self.root.resizable(False, False)

        self.save_file = save_file

        # Backend connection
        self.backend = PlotterBackend(Plotter())
        pos = self.backend.get_position()

        # Variables
        self.step_size = tk.DoubleVar(value=1.0)
        self.z_value = tk.DoubleVar(value=pos.get("z", 0.0))

        self.x_step = tk.DoubleVar(value=1.0)
        self.y_step = tk.DoubleVar(value=1.0)
        self.z_step = tk.DoubleVar(value=1.0)

        self.point_a = {"x": 0.0, "y": 0.0, "z": self.z_value.get()}
        self.point_b = {"x": 0.0, "y": 0.0, "z": self.z_value.get()}

        # --- GUI Layout ---
        self._build_gui(pos)
        self._center_window()

        # Load saved state if available
        self._load_state()

        # Periodic updates
        self._schedule_update()

        

    def _build_gui(self, pos):
        top_frame = ttk.Frame(self.root, padding=8)
        top_frame.grid(row=0, column=0, columnspan=3, sticky="ew")
        self.coord_label = ttk.Label(
            top_frame,
            text=self._fmt_pos(pos),
            font=("Arial", 12, "bold")
        )
        self.coord_label.pack(side="left", padx=(0, 15))

        self._build_move_controls()
        self._build_z_controls()
        self._build_step_controls()
        self._build_ab_controls()
        self._build_save_controls()

    # --- Build Frames ---
    def _build_move_controls(self):
        main_frame = ttk.Frame(self.root, padding=10)
        main_frame.grid(row=1, column=0, sticky="nw")
        move_frame = ttk.Frame(main_frame, padding=5, relief="groove")
        move_frame.grid(row=0, column=0, padx=5, pady=5)

        ttk.Button(move_frame, text="↑", width=5,
                   command=lambda: self.backend.move(0, self.step_size.get())).grid(row=0, column=1, pady=3)
        ttk.Button(move_frame, text="←", width=5,
                   command=lambda: self.backend.move(-self.step_size.get(), 0)).grid(row=1, column=0, padx=3)
        ttk.Button(move_frame, text="Home", width=7,
                   command=self.backend.set_home).grid(row=1, column=1, padx=3)
        ttk.Button(move_frame, text="→", width=5,
                   command=lambda: self.backend.move(self.step_size.get(), 0)).grid(row=1, column=2, padx=3)
        ttk.Button(move_frame, text="↓", width=5,
                   command=lambda: self.backend.move(0, -self.step_size.get())).grid(row=2, column=1, pady=3)

        tk.Scale(main_frame, from_=0.1, to=10, resolution=0.1,
                 orient="horizontal", variable=self.step_size,
                 length=260, label="Step size (mm)", font=("Arial", 10)).grid(row=1, column=0, pady=(10, 0))

    def _build_z_controls(self):
        z_frame = ttk.Frame(self.root, padding=8, relief="groove")
        z_frame.grid(row=0, rowspan=3, column=1, sticky="n", padx=5, pady=5)

        ttk.Label(z_frame, text="", font=("Arial", 10, "bold")).pack(side="top", pady=(0, 2))
        self.z_slider = tk.Scale(
            z_frame, from_=0, to=40, orient="vertical",
            variable=self.z_value, length=180, font=("Arial", 9)
        )
        self.z_slider.pack(side="top", fill="y", padx=2, pady=2)

        ttk.Button(
            z_frame, text="Set Z",
            command=lambda: self.backend.set_z(self.z_value.get())
        ).pack(side="top", pady=(5, 2))

    def _build_step_controls(self):
        step_frame = ttk.Frame(self.root, padding=8, relief="groove")
        step_frame.grid(row=0, rowspan=3, column=2, sticky="n", padx=5, pady=5)

        ttk.Label(step_frame, text="X", font=("Arial", 10, "bold")).grid(row=0, column=0)
        ttk.Label(step_frame, text="Y", font=("Arial", 10, "bold")).grid(row=0, column=1)
        ttk.Label(step_frame, text="Z", font=("Arial", 10, "bold")).grid(row=0, column=2)

        self.x_step_slider = tk.Scale(
            step_frame, from_=0.1, to=10, resolution=0.1,
            orient="vertical", variable=self.x_step,
            length=180, font=("Arial", 9)
        )
        self.x_step_slider.grid(row=1, column=0, padx=2, pady=2)

        self.y_step_slider = tk.Scale(
            step_frame, from_=0.1, to=10, resolution=0.1,
            orient="vertical", variable=self.y_step,
            length=180, font=("Arial", 9)
        )
        self.y_step_slider.grid(row=1, column=1, padx=2, pady=2)

        self.z_step_slider = tk.Scale(
            step_frame, from_=0, to=40, resolution=0.1,
            orient="vertical", variable=self.z_step,
            length=180, font=("Arial", 9)
        )
        self.z_step_slider.grid(row=1, column=2, padx=2, pady=2)

        ttk.Button(step_frame, text="Save Step Values", command=self.save_steps).grid(
            row=2, column=0, columnspan=3, pady=(5, 0)
        )

    def _build_ab_controls(self):
        ab_frame = ttk.Frame(self.root, padding=10, relief="groove")
        ab_frame.grid(row=3, column=0, columnspan=3, sticky="ew", pady=10)

        ttk.Button(ab_frame, text="Set A", width=8, command=self.set_a).grid(row=0, column=0, padx=5)
        ttk.Button(ab_frame, text="Go A", width=8, command=self.backend.go_a).grid(row=0, column=1, padx=5)
        ttk.Button(ab_frame, text="Set B", width=8, command=self.set_b).grid(row=0, column=2, padx=5)
        ttk.Button(ab_frame, text="Go B", width=8, command=self.backend.go_b).grid(row=0, column=4, padx=5)
        ttk.Button(ab_frame, text="Go 1", width=8, command=self.backend.go_one).grid(row=0, column=5, padx=5)

        tk.Label(ab_frame, text="", font=("Arial", 10, "bold")).grid(row=2, column=0, sticky="w")
        self.point_a_value = tk.Label(ab_frame, text="X=+000.00, Y=+000.00, Z=+000.00", font=("Arial", 10))

        tk.Label(ab_frame, text="A:", font=("Arial", 10, "bold")).grid(row=3, column=0, sticky="w")
        self.point_a_value = tk.Label(ab_frame, text="X=+000.00, Y=+000.00, Z=+000.00", font=("Arial", 10))
        self.point_a_value.grid(row=3, column=1, columnspan=3, sticky="w")

        tk.Label(ab_frame, text="B:", font=("Arial", 10, "bold")).grid(row=4, column=0, sticky="w")
        self.point_b_value = tk.Label(ab_frame, text="X=+000.00, Y=+000.00, Z=+000.00", font=("Arial", 10))
        self.point_b_value.grid(row=4, column=1, columnspan=3, sticky="w")

        tk.Label(ab_frame, text="Steps:", font=("Arial", 10, "bold")).grid(row=5, column=0, sticky="w")
        self.step_values_label = tk.Label(ab_frame, text="X=0.0, Y=0.0, Z=0.0", font=("Arial", 10))
        self.step_values_label.grid(row=5, column=1, columnspan=3, sticky="w")

        tk.Label(ab_frame, text="Step count:", font=("Arial", 10, "bold")).grid(row=6, column=0, sticky="w")
        self.step_count_label = tk.Label(ab_frame, text="0", font=("Arial", 10))
        self.step_count_label.grid(row=6, column=1, columnspan=3, sticky="w")

        tk.Label(ab_frame, text="Current step:", font=("Arial", 10, "bold")).grid(row=7, column=0, sticky="w")
        self.current_step_label = tk.Label(ab_frame, text="0", font=("Arial", 10))
        self.current_step_label.grid(row=7, column=1, columnspan=3, sticky="w")

        tk.Label(ab_frame, text="Next step:", font=("Arial", 10, "bold")).grid(row=8, column=0, sticky="w")
        self.next_step_label = tk.Label(ab_frame, text="0", font=("Arial", 10))
        self.next_step_label.grid(row=8, column=1, columnspan=3, sticky="w")

    def _build_save_controls(self):
        save_frame = ttk.Frame(self.root, padding=8)
        save_frame.grid(row=4, column=0, columnspan=3, sticky="ew")

        tk.Label(save_frame, text="Save file:", font=("Arial", 10, "bold")).pack(side="left", padx=(0, 5))
        self.file_entry = ttk.Entry(save_frame, width=50)
        self.file_entry.insert(0, self.save_file)
        self.file_entry.pack(side="left", fill="x", expand=True)
        ttk.Button(save_frame, text="Save", command=self._save_state).pack(side="left", padx=(5, 0))

    # --- State management ---
    def _save_state(self):
        # prefer GUI-stored values, fall back to backend values, then sensible defaults
        point_a = getattr(self, "point_a", None) or getattr(self.backend, "point_a", None) or {"x": 0.0, "y": 0.0, "z": self.z_value.get()}
        point_b = getattr(self, "point_b", None) or getattr(self.backend, "point_b", None) or {"x": 0.0, "y": 0.0, "z": self.z_value.get()}

        state = {
            "point_a": point_a,
            "point_b": point_b,
            "steps": {
                "x": float(self.x_step.get()),
                "y": float(self.y_step.get()),
                "z": float(self.z_step.get())
            }
        }

        with open(self.file_entry.get(), "w") as f:
            json.dump(state, f, indent=2)

    def _load_state(self, filename=None):
        if filename is None:
            filename = self.save_file

        try:
            with open(filename, "r") as f:
                state = json.load(f)
        except (OSError, json.JSONDecodeError):
            return

        # Restore A
        if "point_a" in state:
            self.point_a = state["point_a"]
            self.point_a_value.config(
                text=f"X={self.point_a['x']:+07.2f}, Y={self.point_a['y']:+07.2f}, Z={self.point_a['z']:+07.2f}"
            )
            # Restore in backend
            self.backend.point_a = dict(self.point_a)

        # Restore B
        if "point_b" in state:
            self.point_b = state["point_b"]
            self.point_b_value.config(
                text=f"X={self.point_b['x']:+07.2f}, Y={self.point_b['y']:+07.2f}, Z={self.point_b['z']:+07.2f}"
            )
            # Restore in backend
            self.backend.point_b = dict(self.point_b)

        # Restore step sliders
        if "steps" in state:
            self.x_step.set(state["steps"].get("x", 1.0))
            self.y_step.set(state["steps"].get("y", 1.0))
            self.z_step.set(state["steps"].get("z", 1.0))
            # Restore in backend
            self.backend.set_step_sizes(self.x_step.get(), self.y_step.get(), self.z_step.get())

        # Update frontend step display
        self.step_values_label.config(
            text=f"X={self.x_step.get():.1f}, Y={self.y_step.get():.1f}, Z={self.z_step.get():.1f}"
        )
        # Update step info (current/next step)
        self.update_step_info()




    # --- Other GUI logic (unchanged from earlier split) ---
    def save_steps(self):
        self.backend.set_step_sizes(self.x_step.get(), self.y_step.get(), self.z_step.get())
        self.step_values_label.config(
            text=f"X={self.x_step.get():.1f}, Y={self.y_step.get():.1f}, Z={self.z_step.get():.1f}"
        )
        self.update_step_info()

    def set_a(self):
        pos = self.backend.set_a()
        if pos:
            # store the coordinates in the GUI instance for saving later
            self.point_a = {"x": float(pos["x"]), "y": float(pos["y"]), "z": float(pos["z"])}
            self.point_a_value.config(text=f"X={self.point_a['x']:+07.2f}, Y={self.point_a['y']:+07.2f}, Z={self.point_a['z']:+07.2f}")
            self.update_step_info()
        return pos

    def set_b(self):
        pos = self.backend.set_b()
        if pos:
            self.point_b = {"x": float(pos["x"]), "y": float(pos["y"]), "z": float(pos["z"])}
            self.point_b_value.config(text=f"X={self.point_b['x']:+07.2f}, Y={self.point_b['y']:+07.2f}, Z={self.point_b['z']:+07.2f}")
            self.update_step_info()
        return pos

    def go_one(self):
        first_coords = self.backend.go_one()
        if first_coords:
            self.update_step_info()

    def update_step_info(self):
        current_step, total_steps, next_step_number, next_pos = self.backend.next_raster_step()
        if current_step is None:
            self.step_count_label.config(text="0")
            self.current_step_label.config(text="N/A")
            self.next_step_label.config(text="N/A")
            return
        self.step_count_label.config(text=f"{total_steps}")
        self.current_step_label.config(text=f"{current_step}")
        if current_step >= total_steps:
            self.next_step_label.config(text="End reached")
        else:
            self.next_step_label.config(
                text=f"{next_step_number}, X={next_pos['x']:+07.2f}, Y={next_pos['y']:+07.2f}, Z={next_pos['z']:+07.2f}"
            )

    def _schedule_update(self):
        try:
            pos = self.backend.get_position()
            self.coord_label.config(text=self._fmt_pos(pos))
            self.update_step_info()
        except Exception:
            pass
        self.root.after(500, self._schedule_update)

    def _fmt_pos(self, pos):
        try:
            return f"X: {pos['x']:+07.2f}  Y: {pos['y']:+07.2f}  Z: {pos['z']:+07.2f}"
        except Exception:
            return "X: +000.00  Y: +000.00  Z: +000.00"

    def _center_window(self):
        self.root.update_idletasks()
        w, h = self.root.winfo_width(), self.root.winfo_height()
        sw, sh = self.root.winfo_screenwidth(), self.root.winfo_screenheight()
        x, y = (sw // 2) - (w // 2), (sh // 2) - (h // 2)
        self.root.geometry(f"+{x}+{y}")


class DummyPlotter:
    def move_to(self, x, y, z=None):
        return {"x": x, "y": y, "z": z}
    def set_z(self, z):
        return {"z": z}
    def move_absolute(self, x=None, y=None, feed=1000):
        # Simulate absolute move by updating state only
        if x is not None:
            self.state["x"] = float(x)
        if y is not None:
            self.state["y"] = float(y)
        # Feed rate ignored in dummy
        print(f"[DummyPlotter] move_absolute -> {self.state}")

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("-f", "--file", help="Save file path")
    parser.add_argument("-xyz", help="JSON position to check step number")
    parser.add_argument("-n", type=int, help="Step number to get coordinates for")
    parser.add_argument("-count", action="store_true", help="Return total number of steps in grid")
    args = parser.parse_args()

    # Save file
    if args.file:
        save_file = args.file
    else:
        timestamp = int(time.time())
        save_file = f"/tmp/plotter_{timestamp}.json"

    # Use dummy if not starting GUI
    plotter = DummyPlotter()
    backend = PlotterBackend(plotter)

    # Load state if available
    if os.path.exists(save_file):
        with open(save_file, "r") as f:
            state = json.load(f)
            backend.load_state(state)

    if args.count and args.file:
        with open(args.file, "r") as f:
            state = json.load(f)
        backend.load_state(state)
        print(backend.get_total_steps())
        sys.exit(0)

    if args.xyz:
        pos = json.loads(args.xyz)
        step_no = backend.get_step_number(pos)
        print(step_no)
        sys.exit(0)

    if args.n is not None:
        coords = backend.get_step_coords(args.n)
        print(json.dumps(coords))
        sys.exit(0)

    # GUI mode
    root = tk.Tk()
    from plotterStepperMaker import PlotterGUI
    app = PlotterGUI(root, save_file)
    root.mainloop()

if __name__ == "__main__":
    main()