import tkinter as tk
from tkinter import messagebox
from PIL import Image, ImageTk, ImageDraw
import math
import numpy as np
from simplify_spline import generate_spline_from_rdp, generate_cubicspline, find_nearest_cubicsegment

class SketchCanvas(tk.Canvas):
    def __init__(self, app,parent, size, res, cell_size):
        super().__init__(parent, width=size, height=size, bg="white")
        self.app = app
        self.size = size
        self.res = res
        self.cell_size = cell_size
        
        self.stroke_width = 3
        self.stroke_color = "#214FA0"
        self.skeleton_color = "#00B050"
        self.layout_color = "pink"
        self.highlight_color = "green"
        self.control_point_color = "blue" 
        self.control_point_radius = 3 
        self.selection_box_color = "#C00000" 
        
        self.is_drawing = False
        self.points = []
        
        self.last_image = None
        
        self.original_image = None
        
        self.highlighted_id = None
        self.highlighted_layout_id = None

        self.stroke_item_ids = {}
        self.skeleton_item_ids = {}
        self.layout_item_ids = {}
        
        self.edit_mode = False 
        self.control_points = [] 
        self.active_stroke_index = -1 
        self.active_skeleton_index = -1 
        self.active_layout_index = -1
        self.selected_control_point = None 
        self.dragging = False 
        
        self.selection_mode = False 
        self.selection_start = None 
        self.selection_current = None  
        self.selection_box = None  

        self.edit_layout = False 
        self.layout_handle_ids = []
        self.layout_drag_handle = None
        
        self.stroke_callback = None
        self.edit_callback = None  
        self.selection_callback = None
        self.status_callback = None
        
        self.bind("<Button-1>", self.on_lmouse_down)
        self.bind("<Button-3>", self.on_rmouse_down)
        self.bind("<B1-Motion>", self.on_lmouse_drag)
        self.bind("<ButtonRelease-1>", self.on_lmouse_up)
        
        self.init_canvas()
    
    def init_canvas(self):
        self.img = Image.new('RGB', (self.size, self.size), 'white')
        self.original_image = self.img.copy()
        
        self.photo = ImageTk.PhotoImage(self.img)
        self.create_image(0, 0, image=self.photo, anchor=tk.NW)
    
    def on_lmouse_down(self, event):
        if self.edit_layout:
            clicked = self.find_closest(event.x, event.y)
            if clicked:
                tags = self.gettags(clicked[0])
                for t in tags:
                    if t.startswith("handle_"):
                        self.layout_drag_handle = int(t.split("_")[1])
                        self.dragging = True
                        rect_id = self.layout_item_ids[self.active_layout_index][0]
                        x1, y1, x2, y2 = self.coords(rect_id)

                        self.origin_handle_pos = [
                            (x1, y1), ((x1+x2)/2, y1), (x2, y1),
                            (x2, (y1+y2)/2), (x2, y2), ((x1+x2)/2, y2),
                            (x1, y2), (x1, (y1+y2)/2)
                        ]
                        anchor_pairs = {
                            0: 4, 1: 5, 2: 6, 3: 7, 4: 0, 5: 1, 6: 2, 7: 3
                        }
                        anchor_idx = anchor_pairs[self.layout_drag_handle]
                        if self.layout_drag_handle % 2 == 0:
                            self.anchor_x, self.anchor_y = self.origin_handle_pos[anchor_idx]
                        elif self.layout_drag_handle in [1,5]:
                            self.anchor_x = [self.origin_handle_pos[0][0],self.origin_handle_pos[2][0]]
                            self.anchor_y = self.origin_handle_pos[anchor_idx][1]
                        elif self.layout_drag_handle in [3,7]:
                            self.anchor_x = self.origin_handle_pos[anchor_idx][0]
                            self.anchor_y = [self.origin_handle_pos[0][1],self.origin_handle_pos[6][1]]
                        return
        elif self.selection_mode:
            self.selection_start = (event.x, event.y)
            self.selection_current = (event.x, event.y)
            return

        elif self.edit_mode:
            for i, (x, y) in enumerate(self.control_points):
                if abs(event.x - x) <= self.control_point_radius * 2 and abs(event.y - y) <= self.control_point_radius * 2:
                    self.selected_control_point = i
                    self.dragging = True
                    return

        else:
            self.is_drawing = True
            self.points = [(event.x, event.y)]

    def on_rmouse_down(self,event):
        if self.edit_mode:
            clicked_point = (event.x, event.y)
            deleted = False

            for i, (x, y) in enumerate(self.control_points):
                if abs(event.x - x) <= self.control_point_radius * 2 and abs(event.y - y) <= self.control_point_radius * 2:
                    del self.control_points[i]
                    deleted = True
                    break

            if not deleted:
                insert_index = find_nearest_cubicsegment(self.control_points, clicked_point)
                self.control_points.insert(insert_index, clicked_point)

            self.delete("preview")
            spline = generate_cubicspline(self.control_points, isPeriodic=self.is_periodic)
            draw_points = [(x, y) for x, y in spline]
            self.draw_current_stroke(draw_points, tag="preview")
            self.draw_control_points(self.control_points, tag="preview")
            return
    
    def on_lmouse_drag(self, event):
        if self.edit_layout and self.dragging and self.layout_drag_handle is not None:
            rect_id = self.layout_item_ids[self.active_layout_index][0]

            if self.layout_drag_handle % 2 == 0:
                new_x1, new_x2 = sorted([self.anchor_x, event.x])
                new_y1, new_y2 = sorted([self.anchor_y, event.y])
            elif self.layout_drag_handle in [1,5]:
                new_x1, new_x2 = sorted(self.anchor_x)
                new_y1, new_y2 = sorted([self.anchor_y, event.y])
            elif self.layout_drag_handle in [3,7]:
                new_x1, new_x2 = sorted([self.anchor_x, event.x])
                new_y1, new_y2 = sorted(self.anchor_y)

            self.coords(rect_id, new_x1, new_y1, new_x2, new_y2)

            x1, y1, x2, y2 = new_x1, new_y1, new_x2, new_y2
            new_handles = [
                (x1, y1), ((x1+x2)/2, y1), (x2, y1),
                (x2, (y1+y2)/2), (x2, y2), ((x1+x2)/2, y2),
                (x1, y2), (x1, (y1+y2)/2),
            ]
            r = self.control_point_radius
            for i, (hx, hy) in enumerate(new_handles):
                self.coords(
                    self.layout_handle_ids[i],
                    hx - r, hy - r, hx + r, hy + r
                )
            text_id = self.layout_item_ids[self.active_layout_index][1]
            self.coords(text_id, int(x1)+3, int(y1)+3)
            
        elif self.selection_mode and self.selection_start:
            self.selection_current = (event.x, event.y)
            
            self.draw_selection_box()
            return
        elif self.edit_mode and self.dragging and self.selected_control_point is not None:
            self.control_points[self.selected_control_point] = (event.x, event.y)

            self.status_callback(f"控制点 {self.selected_control_point} 坐标: ({event.x}, {event.y})")

            self.delete("preview")

            spline = generate_cubicspline(self.control_points, isPeriodic=self.is_periodic)
            draw_points = [(x, y) for x, y in spline]
            if self.active_skeleton_index != -1:
                self.create_line(draw_points, fill=self.skeleton_color, width=self.stroke_width,tags="preview")
            elif self.active_stroke_index != -1:
                self.create_line(draw_points, fill=self.stroke_color, width=self.stroke_width,tags="preview")
            self.draw_control_points(self.control_points,tag="preview")
            
        elif not self.edit_mode and self.is_drawing: 
            self.points.append((event.x, event.y))
            
            self.draw_current_stroke(self.points, tag="preview")
    
    def on_lmouse_up(self, event):
        if self.edit_layout and self.dragging:
            self.dragging = False
            self.layout_drag_handle = None
            self.anchor_x = None
            self.anchor_y = None
            self.layout_drag_handle = None
        if self.selection_mode and self.selection_start:
            self.selection_current = (event.x, event.y)
            
            x1 = min(self.selection_start[0], self.selection_current[0])
            y1 = max(self.selection_start[1], self.selection_current[1])
            x2 = max(self.selection_start[0], self.selection_current[0])
            y2 = min(self.selection_start[1], self.selection_current[1])
            
            self.selection_box = (x1, y1, x2, y2)
            
            self.draw_selection_box()
            
            grid_x1 = max(1, min(self.res, x1))
            grid_y1 = max(1, min(self.res, y1))
            grid_x2 = max(1, min(self.res, x2))
            grid_y2 = max(1, min(self.res, y2))
            
            grid_selection = (
                min(grid_x1, grid_x2), 
                max(grid_y1, grid_y2), 
                max(grid_x1, grid_x2), 
                min(grid_y1, grid_y2)
            )

            if self.selection_callback:
                self.selection_callback(self.selection_box, grid_selection)
            
            self.selection_mode = False
            self.selection_start = None
            self.selection_current = None
            
            return
        
        elif self.edit_mode and self.dragging:
            self.dragging = False
            self.delete("preview")  
            if self.edit_callback:
                updated = list(self.control_points)
                if self.active_skeleton_index != -1:
                    self.edit_callback(self.active_skeleton_index, updated)
                elif self.active_stroke_index != -1:
                    self.edit_callback(self.active_stroke_index, updated)

            self.selected_control_point = None

        elif not self.edit_mode and self.is_drawing:
            if self.points[-1] != (event.x, event.y):
                self.points.append((event.x, event.y))
 
            self.delete("preview")
            curve, simplified = generate_spline_from_rdp(self.points, initial_tol=4.0, is_periodic=False, resolution_scale=1.0)

            self.points = [tuple(pt) for pt in curve.tolist()]

            self.draw_current_stroke(self.points, stroke_index=len(self.stroke_item_ids))

            if self.stroke_callback and len(simplified) > 1:
                self.stroke_callback(simplified.tolist())

            self.is_drawing = False
    
    def draw_current_stroke(self,points, stroke_index=None,skeleton_index=None,tag=None):
        if len(points) < 2:
            return
        
        if skeleton_index is not None:
            color = self.skeleton_color
        else:
            color = self.stroke_color
        item_id = self.create_line(points, fill=color, width=self.stroke_width, tags=tag)

        if tag != "preview":
            if skeleton_index is not None:
                self.skeleton_item_ids[skeleton_index] = item_id
            if stroke_index is not None:
                self.stroke_item_ids[stroke_index] = item_id
        
    def draw_layout_rec(self, rec_xyxy,layout_index=None,tag=None):
        xmin,ymax,xmax,ymin = rec_xyxy

        box_item = self.create_rectangle(
            xmin,ymax,xmax,ymin,
            outline=self.selection_box_color,
            width=2,
            dash=(4, 2),tags=tag
        )

        if tag != "preview":
            if layout_index is not None:
                text_item = self.create_text(
                    int(xmin) + 3, int(ymax) + 3, 
                    anchor='nw',       
                    text=self.app.layouts[layout_index]["meaning"],
                    fill=self.selection_box_color,
                    font=("Arial", 10, "bold"),
                    tags=tag
                )
                self.layout_item_ids[layout_index] = [box_item,text_item]

    
    def draw_cubic_bezier(self, p0, p1, p2, p3, steps=20):
        bezier_points = []
        for t in np.linspace(0, 1, steps):
            point = (1-t)**3 * p0 + 3*(1-t)**2*t * p1 + 3*(1-t)*t**2 * p2 + t**3 * p3
            bezier_points.append((int(point[0]), int(point[1])))
        
        if len(bezier_points) > 1:
            print(f"draw_cubic_bezier {bezier_points}")
            self.draw.line(bezier_points, fill=self.stroke_color, width=self.stroke_width)
    
    def draw_cubic_spline(self, points, highlight=False, show_control=False):
        if not points:
            return
        
        original_color = self.stroke_color
        
        if highlight:
            self.stroke_color = self.highlight_color

        full_points = generate_cubicspline(points)
        int_points = [(int(x), int(y)) for x, y in full_points]
        self.draw.line(int_points, fill=self.stroke_color, width=self.stroke_width)

        if show_control and full_points:
            self.draw_control_points(points)
            self.control_points = points
        
        self.stroke_color = original_color
        
        self.photo = ImageTk.PhotoImage(self.img)
        self.create_image(0, 0, image=self.photo, anchor=tk.NW)

    def draw_control_points(self, control_points,tag=None):
        for x, y in control_points:
            radius = self.control_point_radius
            self.create_oval(x - radius, y - radius, x + radius, y + radius, fill='red',tags=tag)

    def show_skeletons(self,is_show=True):
        if is_show:
            state = "normal"
        elif not is_show:
            state = "hidden"
        for i in self.skeleton_item_ids.keys():
            sk_item_id = self.skeleton_item_ids.get(i,None)
            if sk_item_id:
                self.itemconfigure(sk_item_id, fill=self.skeleton_color, state=state)
    
    def show_layouts(self,is_show=True):
        if is_show:
            state = "normal"
        elif not is_show:
            state = "hidden"
        for i in self.layout_item_ids.keys():
            lay_item_id = self.layout_item_ids.get(i,None)
            if lay_item_id:
                self.itemconfigure(lay_item_id[0], state=state)
                self.itemconfigure(lay_item_id[1], state=state)

                

    def enter_edit_mode(self, stroke_index, points, is_skeleton=False, is_periodic=False):
        self.edit_mode = True

        self.control_points = list(points)
        spline = generate_cubicspline(self.control_points)
        self.points = [(x, y) for x, y in spline]
        self.is_periodic = is_periodic 

        if is_skeleton:
            self.active_skeleton_index = stroke_index
            pt_item_id = self.stroke_item_ids.get(stroke_index)
            sk_item_id = self.skeleton_item_ids.get(stroke_index,None)
            if sk_item_id:
                self.itemconfigure(pt_item_id, fill=self.highlight_color, state='normal')
                self.itemconfigure(sk_item_id, fill=self.skeleton_color, state='hidden')
                self.draw_current_stroke(self.points, skeleton_index=stroke_index, tag="preview")
                self.draw_control_points(self.control_points, tag="preview")
            else:
                messagebox.showinfo("tip", "There is no valid point in this skeleton")
        else:
            self.active_stroke_index = stroke_index
            item_id = self.stroke_item_ids.get(stroke_index)
            if item_id:
                self.itemconfigure(item_id, fill=self.stroke_color, state='hidden')
                self.draw_current_stroke(self.points, tag="preview")
            # 只添加控制点预览即可
            self.draw_control_points(self.control_points, tag="preview")


    def exit_edit_mode(self,register=False,stroke_id=-1,skeleton_id=-1,layout_id=-1):
        self.edit_mode = False
        if register:
            if skeleton_id >= 0:
                if self.active_skeleton_index in self.skeleton_item_ids:
                    self.delete("preview")
                    spline = generate_cubicspline(self.control_points)
                    draw_points = [(x, y) for x, y in spline]
                    self.draw_current_stroke(draw_points, skeleton_index=self.active_skeleton_index)
                    self.itemconfigure(self.skeleton_item_ids[self.active_skeleton_index], state='normal')
                    self.itemconfigure(self.stroke_item_ids[self.active_skeleton_index],fill=self.stroke_color, state='normal')

                    if self.stroke_callback and len(self.control_points) > 1:
                        self.stroke_callback(self.control_points,skeleton_id,is_skeleton=True)

            if stroke_id >= 0:
                if self.active_stroke_index in self.stroke_item_ids:
                    self.delete("preview")
                    spline = generate_cubicspline(self.control_points, isPeriodic=self.is_periodic)
                    draw_points = [(x, y) for x, y in spline]
                    self.draw_current_stroke(draw_points, stroke_index=self.active_stroke_index)
                    self.itemconfigure(self.stroke_item_ids[self.active_stroke_index], state='normal')

                    if self.stroke_callback and len(self.control_points) > 1:
                        self.stroke_callback(self.control_points,stroke_id)
            if layout_id >= 0:
                if self.active_layout_index == layout_id:
                    self.edit_layout = False
                    self.delete("layout_handle")
                    self.delete("preview")
                    rect_id, text_id = self.layout_item_ids.get(self.active_layout_index, [None, None])
                    if rect_id:
                        self.itemconfigure(rect_id, state='normal', outline=self.selection_box_color)
                    if text_id:
                        self.itemconfigure(text_id, state='normal', fill=self.selection_box_color)
                    rect_coords = self.coords(rect_id)
                    self.app.on_layout_edited(rect_coords,layout_id)
        self.active_stroke_index = -1
        self.active_skeleton_index = -1
        self.active_layout_index = -1
        self.control_points = []
        self.active_layout_rect_id = None
        self.layout_drag_handle = None
        self.selected_control_point = None
    
    def set_edit_callback(self, callback):
        self.edit_callback = callback
    
    def clear(self):
        self.img = Image.new('RGB', (self.size, self.size), 'white')
        self.draw = ImageDraw.Draw(self.img)
        self.original_image = self.img.copy()
        self.highlighted_id = None
        self.photo = ImageTk.PhotoImage(self.img)
        self.create_image(0, 0, image=self.photo, anchor=tk.NW)
    
    def set_stroke_callback(self, callback):
        self.stroke_callback = callback
    
    def set_status_callback(self, callback):
        self.status_callback = callback
    
    def save_as_png(self, file_path):
        if file_path.endswith('.svg'):
            file_path = file_path[:-4] + '.png'
        self.img.save(file_path)
    
    def highlight_stroke(self, stroke_id=None):
        if stroke_id is None or stroke_id == self.highlighted_id:
            for idx, item_id in self.stroke_item_ids.items():
                try:
                    self.itemconfigure(item_id, state='normal', fill=self.stroke_color)
                except:
                    pass
            self.highlighted_id = None
            return False

        for idx, item_id in self.stroke_item_ids.items():
            if idx == stroke_id:
                self.itemconfigure(item_id, state='normal', fill=self.highlight_color)
            else:
                self.itemconfigure(item_id, state='normal', fill=self.stroke_color)

        self.highlighted_id = stroke_id
        return True
    
    def highlight_layout(self, layout_id=None):
        if layout_id is None or layout_id == self.highlighted_layout_id:
            for idx, item_id in self.layout_item_ids.items():
                self.itemconfigure(item_id[0], state='normal', outline=self.selection_box_color)
                self.itemconfigure(item_id[1], state='normal', fill=self.selection_box_color)
            self.highlighted_id = None
            return False

        for idx, item_id in self.layout_item_ids.items():
            if idx == layout_id:
                self.itemconfigure(item_id[0], state='normal', outline=self.highlight_color)
                self.itemconfigure(item_id[1], state='normal', fill=self.highlight_color)
            else:
                self.itemconfigure(item_id[0], state='normal', outline=self.selection_box_color)
                self.itemconfigure(item_id[1], state='normal', fill=self.selection_box_color)

        self.highlighted_layout_id = layout_id
        return True
    
    def draw_selection_box(self):
        """在 Canvas 上绘制选择框，不破坏原有笔画"""
        if not self.selection_start or not self.selection_current:
            return
        
        x1 = min(self.selection_start[0], self.selection_current[0])
        y1 = max(self.selection_start[1], self.selection_current[1])
        x2 = max(self.selection_start[0], self.selection_current[0])
        y2 = min(self.selection_start[1], self.selection_current[1])
        
        if self.selection_box_item:
            self.delete(self.selection_box_item)
        if self.selection_fill_item:
            self.delete(self.selection_fill_item)

        self.selection_fill_item = self.create_rectangle(
            x1, y1, x2, y2,
            fill="#f0f0f0",
            stipple="gray12",
            outline=""
        )

        self.selection_box_item = self.create_rectangle(
            x1, y1, x2, y2,
            outline=self.selection_box_color,
            width=2,
            dash=(4, 2)
        )
    
    def enter_selection_mode(self):
        if self.edit_mode:
            self.exit_edit_mode()
            
        self.selection_mode = True
        self.selection_start = None
        self.selection_current = None
        
        self.selection_box_item = None
        self.selection_fill_item = None
    
    def exit_selection_mode(self):
        self.selection_mode = False
        self.selection_start = None
        self.selection_current = None
        
        if self.selection_box_item:
            self.delete(self.selection_box_item)
            self.selection_box_item = None
        if self.selection_fill_item:
            self.delete(self.selection_fill_item)
            self.selection_fill_item = None
    
    def set_selection_callback(self, callback):
        self.selection_callback = callback 