import base64
import math
import os
from pathlib import Path
import tkinter as tk
from tkinter import messagebox

import cv2
import numpy as np


INPUT_FOLDER = "pages"
OUTPUT_FOLDER = "chars_gui"

CX0, CY0 = 1027, 1257
CX1, CY1 = 4029, 5673

CHAR_W = 700
CHAR_H = 819
ROWS = 5
COLS = 5

OFFSET_MIN = -300
OFFSET_MAX = 300

SCALE_MIN = 0.70
SCALE_MAX = 1.30
SCALE_RES = 0.001

ROT_MIN = -30
ROT_MAX = 30
ROT_RES = 0.1

DISPLAY_MAX_W = 1100
DISPLAY_MAX_H = 900


def imread_unicode(path: str):
    return cv2.imdecode(np.fromfile(path, dtype=np.uint8), cv2.IMREAD_COLOR)


def imwrite_unicode(path: str, img: np.ndarray):
    cv2.imencode(".png", img)[1].tofile(path)


def clamp_int(v: int, lo: int, hi: int) -> int:
    return int(max(lo, min(hi, int(v))))

def clamp_float(v: float, lo: float, hi: float) -> float:
    return float(max(lo, min(hi, float(v))))


def crop_fixed(img: np.ndarray, x1: int, y1: int, size: int) -> np.ndarray:
    h, w = img.shape[:2]
    x2 = x1 + size
    y2 = y1 + size

    out = np.full((size, size, 3), 255, dtype=np.uint8)

    sx1 = max(0, x1)
    sy1 = max(0, y1)
    sx2 = min(w, x2)
    sy2 = min(h, y2)

    src = img[sy1:sy2, sx1:sx2]
    if src.size == 0:
        return out

    dx = max(0, -x1)
    dy = max(0, -y1)
    out[dy : dy + src.shape[0], dx : dx + src.shape[1]] = src
    return out


class PageCutterApp:
    def __init__(self, root: tk.Tk):
        self.root = root
        self.root.title("篆文 5×5 分块裁切（整体偏移 / 缩放 / 旋转）")

        self.pages = sorted(Path(INPUT_FOLDER).glob("*.png"))
        if not self.pages:
            messagebox.showerror("错误", f"未找到图片：{Path(INPUT_FOLDER).resolve()}")
            raise SystemExit(1)

        self.page_index = 0
        self.src_bgr = None
        self.src_h = 0
        self.src_w = 0
        self.scale = 1.0

        self.offset_x = tk.IntVar(value=0)
        self.offset_y = tk.IntVar(value=0)
        self.factor = tk.DoubleVar(value=1.0)
        self.angle = tk.DoubleVar(value=0.0)

        self.mode = tk.StringVar(value="all")
        self.selected_row = tk.IntVar(value=1)
        self.params_all = {"dx": 0, "dy": 0, "factor": 1.0, "angle": 0.0}
        self.params_row = {r: {"dx": 0, "dy": 0, "factor": 1.0, "angle": 0.0} for r in range(1, ROWS + 1)}
        self._suspend = False
        self._notice_after = None

        self._dragging = False
        self._drag_start = None

        self._build_ui()
        self._bind_keys()
        self.mode.trace_add("write", lambda *_: self._on_mode_or_row_change())
        self.selected_row.trace_add("write", lambda *_: self._on_mode_or_row_change())

        self._load_page(0)

    def _save_current_params(self):
        mode = str(self.mode.get())
        if mode == "all":
            store = self.params_all
        else:
            r = clamp_int(int(self.selected_row.get()), 1, ROWS)
            store = self.params_row[r]
        store["dx"] = int(self.offset_x.get())
        store["dy"] = int(self.offset_y.get())
        store["factor"] = float(self.factor.get())
        store["angle"] = float(self.angle.get())

    def _load_params(self):
        mode = str(self.mode.get())
        if mode == "all":
            store = self.params_all
        else:
            r = clamp_int(int(self.selected_row.get()), 1, ROWS)
            store = self.params_row[r]
        self.offset_x.set(int(store["dx"]))
        self.offset_y.set(int(store["dy"]))
        self.factor.set(float(store["factor"]))
        self.angle.set(float(store["angle"]))

    def _on_mode_or_row_change(self):
        if self._suspend:
            return
        self._save_current_params()
        self._suspend = True
        self._load_params()
        self._suspend = False
        self._update_boxes()

    def _build_ui(self):
        self.root.geometry("1400x980")

        main = tk.Frame(self.root)
        main.pack(fill=tk.BOTH, expand=True)

        left = tk.Frame(main)
        left.pack(side=tk.LEFT, fill=tk.BOTH, expand=True)

        right = tk.Frame(main, width=260)
        right.pack(side=tk.RIGHT, fill=tk.Y)
        right.pack_propagate(False)

        self.canvas = tk.Canvas(left, bg="#202020", highlightthickness=0)
        self.canvas.pack(fill=tk.BOTH, expand=True)

        self.status = tk.Label(right, text="", anchor="w", justify="left")
        self.status.pack(fill=tk.X, padx=10, pady=(10, 6))

        mode_frame = tk.Frame(right)
        mode_frame.pack(fill=tk.X, padx=10, pady=(0, 10))
        tk.Label(mode_frame, text="模式").pack(anchor="w")
        tk.Radiobutton(mode_frame, text="25 框", variable=self.mode, value="all").pack(anchor="w")
        tk.Radiobutton(mode_frame, text="单行 5 框", variable=self.mode, value="row").pack(anchor="w")

        row_frame = tk.Frame(right)
        row_frame.pack(fill=tk.X, padx=10, pady=(0, 12))
        tk.Label(row_frame, text="选择行").pack(anchor="w")
        self.row_menu = tk.OptionMenu(row_frame, self.selected_row, *list(range(1, ROWS + 1)))
        self.row_menu.pack(fill=tk.X)

        tk.Label(right, text="整体偏移 X").pack(anchor="w", padx=10)
        self.scale_x = tk.Scale(
            right,
            from_=OFFSET_MIN,
            to=OFFSET_MAX,
            orient=tk.HORIZONTAL,
            variable=self.offset_x,
            command=lambda _=None: self._update_boxes(),
        )
        self.scale_x.pack(fill=tk.X, padx=10, pady=(0, 10))

        tk.Label(right, text="整体偏移 Y").pack(anchor="w", padx=10)
        self.scale_y = tk.Scale(
            right,
            from_=OFFSET_MIN,
            to=OFFSET_MAX,
            orient=tk.HORIZONTAL,
            variable=self.offset_y,
            command=lambda _=None: self._update_boxes(),
        )
        self.scale_y.pack(fill=tk.X, padx=10, pady=(0, 12))

        tk.Label(right, text="整体缩放").pack(anchor="w", padx=10)
        self.scale_factor = tk.Scale(
            right,
            from_=SCALE_MIN,
            to=SCALE_MAX,
            resolution=SCALE_RES,
            orient=tk.HORIZONTAL,
            variable=self.factor,
            command=lambda _=None: self._update_boxes(),
        )
        self.scale_factor.pack(fill=tk.X, padx=10, pady=(0, 12))

        tk.Label(right, text="整体旋转（°）").pack(anchor="w", padx=10)
        self.scale_angle = tk.Scale(
            right,
            from_=ROT_MIN,
            to=ROT_MAX,
            resolution=ROT_RES,
            orient=tk.HORIZONTAL,
            variable=self.angle,
            command=lambda _=None: self._update_boxes(),
        )
        self.scale_angle.pack(fill=tk.X, padx=10, pady=(0, 12))

        btn_row = tk.Frame(right)
        btn_row.pack(fill=tk.X, padx=10, pady=(0, 10))

        self.btn_prev = tk.Button(btn_row, text="上一页", command=self.prev_page)
        self.btn_prev.pack(side=tk.LEFT, fill=tk.X, expand=True)

        self.btn_next = tk.Button(btn_row, text="下一页", command=self.next_page)
        self.btn_next.pack(side=tk.LEFT, fill=tk.X, expand=True, padx=(8, 0))

        self.btn_reset = tk.Button(right, text="重置参数", command=self.reset_offset)
        self.btn_reset.pack(fill=tk.X, padx=10, pady=(0, 10))

        self.btn_save = tk.Button(right, text="保存本页 25 个字", command=self.save_page)
        self.btn_save.pack(fill=tk.X, padx=10, pady=(0, 10))

        self.help = tk.Label(
            right,
            text="拖拽：整体移动\n滚轮：缩放（Shift 精细）\n方向键：微调（Shift 加速）\nW/S：选上一/下一行（自动切到单行）\nA/D：旋转 -/+0.1°\nQ：保存  R：重置  Z/C：上一/下一页",
            anchor="w",
            justify="left",
        )
        self.help.pack(fill=tk.X, padx=10, pady=(0, 10))

        self.notice = tk.Label(right, text="", anchor="e", justify="right", fg="#00ff7f")
        self.notice.pack(side=tk.BOTTOM, fill=tk.X, padx=10, pady=10)

        self.canvas.bind("<ButtonPress-1>", self._on_mouse_down)
        self.canvas.bind("<B1-Motion>", self._on_mouse_drag)
        self.canvas.bind("<ButtonRelease-1>", self._on_mouse_up)
        self.canvas.bind("<MouseWheel>", self._on_mouse_wheel)

        self._photo = None
        self._img_item = None
        self._boxes = []

    def _bind_keys(self):
        self.root.bind("<Left>", self._on_arrow)
        self.root.bind("<Right>", self._on_arrow)
        self.root.bind("<Up>", self._on_arrow)
        self.root.bind("<Down>", self._on_arrow)
        self.root.bind("<Return>", lambda _=None: self.save_page())
        self.root.bind("q", lambda _=None: self.save_page())
        self.root.bind("Q", lambda _=None: self.save_page())
        self.root.bind("r", lambda _=None: self.reset_offset())
        self.root.bind("R", lambda _=None: self.reset_offset())
        self.root.bind("w", lambda _=None: self._select_row(-1))
        self.root.bind("W", lambda _=None: self._select_row(-1))
        self.root.bind("s", lambda _=None: self._select_row(+1))
        self.root.bind("S", lambda _=None: self._select_row(+1))
        self.root.bind("a", lambda _=None: self._rotate(-ROT_RES))
        self.root.bind("A", lambda _=None: self._rotate(-ROT_RES))
        self.root.bind("d", lambda _=None: self._rotate(+ROT_RES))
        self.root.bind("D", lambda _=None: self._rotate(+ROT_RES))
        self.root.bind("z", lambda _=None: self.prev_page())
        self.root.bind("Z", lambda _=None: self.next_page())
        self.root.bind("C", lambda _=None: self.prev_page())
        self.root.bind("c", lambda _=None: self.next_page())

    def _set_notice(self, text: str, ok: bool = True):
        if self._notice_after is not None:
            try:
                self.root.after_cancel(self._notice_after)
            except Exception:
                pass
            self._notice_after = None
        self.notice.config(text=str(text), fg=("#00ff7f" if ok else "#ff6b6b"))
        self._notice_after = self.root.after(2500, lambda: self.notice.config(text=""))

    def _select_row(self, delta: int):
        if str(self.mode.get()) != "row":
            self.mode.set("row")
        r = clamp_int(int(self.selected_row.get()) + int(delta), 1, ROWS)
        self.selected_row.set(r)

    def _rotate(self, delta_deg: float):
        v = float(self.angle.get()) + float(delta_deg)
        self.angle.set(clamp_float(v, float(ROT_MIN), float(ROT_MAX)))
        self._update_boxes()

    def _on_mouse_wheel(self, e):
        if self._suspend:
            return
        direction = 1.0 if e.delta > 0 else -1.0
        step = 0.001 if (e.state & 0x0001) else 0.005
        v = float(self.factor.get()) + direction * step
        self.factor.set(clamp_float(v, float(SCALE_MIN), float(SCALE_MAX)))
        self._update_boxes()

    def _grid_centers(self):
        step_x = (CX1 - CX0) / (COLS - 1)
        step_y = (CY1 - CY0) / (ROWS - 1)
        centers = []
        for r in range(ROWS):
            for c in range(COLS):
                centers.append((CX0 + c * step_x, CY0 + r * step_y))
        return centers

    def _load_page(self, index: int):
        index = clamp_int(index, 0, len(self.pages) - 1)
        self.page_index = index

        img_path = self.pages[self.page_index]
        bgr = imread_unicode(str(img_path))
        if bgr is None:
            messagebox.showerror("错误", f"无法读取：{img_path}")
            return

        self.src_bgr = bgr
        self.src_h, self.src_w = bgr.shape[:2]

        scale = min(DISPLAY_MAX_W / self.src_w, DISPLAY_MAX_H / self.src_h, 1.0)
        self.scale = float(scale)

        rgb = cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB)
        if self.scale != 1.0:
            rgb = cv2.resize(rgb, (int(self.src_w * self.scale), int(self.src_h * self.scale)), interpolation=cv2.INTER_AREA)

        ok, buf = cv2.imencode(".png", rgb)
        if not ok:
            messagebox.showerror("错误", "图片编码失败")
            return

        b64 = base64.b64encode(buf).decode("ascii")
        self._photo = tk.PhotoImage(data=b64)

        self.canvas.delete("all")
        self._img_item = self.canvas.create_image(0, 0, anchor="nw", image=self._photo)
        self.canvas.config(width=self._photo.width(), height=self._photo.height(), scrollregion=(0, 0, self._photo.width(), self._photo.height()))

        self._boxes = []
        for _ in range(ROWS * COLS):
            rid = self.canvas.create_polygon(0, 0, 0, 0, 0, 0, 0, 0, outline="#00ff7f", fill="", width=2)
            self._boxes.append(rid)

        self._update_boxes()

    def _selected_indices(self):
        mode = str(self.mode.get())
        if mode == "all":
            return list(range(ROWS * COLS))
        r = clamp_int(int(self.selected_row.get()), 1, ROWS) - 1
        return list(range(r * COLS, r * COLS + COLS))

    def _update_boxes(self):
        if self._suspend:
            return
        if self.src_bgr is None:
            return

        dx = float(self.offset_x.get())
        dy = float(self.offset_y.get())
        factor = float(self.factor.get())
        ang_deg = float(self.angle.get())
        ang = math.radians(ang_deg)
        ca = math.cos(ang)
        sa = math.sin(ang)
        img_cx = self.src_w / 2.0
        img_cy = self.src_h / 2.0

        centers = self._grid_centers()
        half_w = CHAR_W / 2.0
        half_h = CHAR_H / 2.0
        half_w_s = half_w * factor
        half_h_s = half_h * factor
        sel = set(self._selected_indices())
        mode = str(self.mode.get())
        row = clamp_int(int(self.selected_row.get()), 1, ROWS)

        def tx(x: float, y: float):
            vx = x - img_cx
            vy = y - img_cy
            rx = vx * ca - vy * sa
            ry = vx * sa + vy * ca
            return img_cx + factor * rx + dx, img_cy + factor * ry + dy

        for i, (cx, cy) in enumerate(centers):
            ccx, ccy = tx(cx, cy)
            p0x = ccx + (-half_w_s) * ca - (-half_h_s) * sa
            p0y = ccy + (-half_w_s) * sa + (-half_h_s) * ca
            p1x = ccx + (+half_w_s) * ca - (-half_h_s) * sa
            p1y = ccy + (+half_w_s) * sa + (-half_h_s) * ca
            p2x = ccx + (+half_w_s) * ca - (+half_h_s) * sa
            p2y = ccy + (+half_w_s) * sa + (+half_h_s) * ca
            p3x = ccx + (-half_w_s) * ca - (+half_h_s) * sa
            p3y = ccy + (-half_w_s) * sa + (+half_h_s) * ca

            pts = [
                p0x * self.scale,
                p0y * self.scale,
                p1x * self.scale,
                p1y * self.scale,
                p2x * self.scale,
                p2y * self.scale,
                p3x * self.scale,
                p3y * self.scale,
            ]
            self.canvas.coords(self._boxes[i], *pts)
            if mode == "row" and i not in sel:
                self.canvas.itemconfigure(self._boxes[i], state="hidden")
            else:
                self.canvas.itemconfigure(self._boxes[i], state="normal")

        stem = self.pages[self.page_index].stem
        if mode == "all":
            mode_hint = "25 框"
            save_n = 25
            self.row_menu.config(state="disabled")
        else:
            mode_hint = f"单行 5 框（第 {row} 行）"
            save_n = 5
            self.row_menu.config(state="normal")
        self.btn_save.config(text=f"保存本页 {save_n} 个字")
        self.status.config(
            text=f"页：{self.page_index + 1}/{len(self.pages)}\n文件：{stem}\n模式：{mode_hint}\n偏移：({self.offset_x.get()}, {self.offset_y.get()})\n缩放：{self.factor.get():.3f}\n旋转：{self.angle.get():.1f}°\n框大小：{CHAR_W}×{CHAR_H}\n输出：{OUTPUT_FOLDER}"
        )

        self.btn_prev.config(state=tk.NORMAL if self.page_index > 0 else tk.DISABLED)
        self.btn_next.config(state=tk.NORMAL if self.page_index < len(self.pages) - 1 else tk.DISABLED)

    def _on_arrow(self, e):
        step = 10 if (e.state & 0x0001) else 1
        if e.keysym == "Left":
            self.offset_x.set(clamp_int(self.offset_x.get() - step, OFFSET_MIN, OFFSET_MAX))
        elif e.keysym == "Right":
            self.offset_x.set(clamp_int(self.offset_x.get() + step, OFFSET_MIN, OFFSET_MAX))
        elif e.keysym == "Up":
            self.offset_y.set(clamp_int(self.offset_y.get() - step, OFFSET_MIN, OFFSET_MAX))
        elif e.keysym == "Down":
            self.offset_y.set(clamp_int(self.offset_y.get() + step, OFFSET_MIN, OFFSET_MAX))
        self._update_boxes()

    def _on_mouse_down(self, e):
        self._dragging = True
        self._drag_start = (e.x, e.y, self.offset_x.get(), self.offset_y.get())

    def _on_mouse_drag(self, e):
        if not self._dragging or self._drag_start is None:
            return
        sx, sy, ox, oy = self._drag_start
        if self.scale <= 0:
            return
        dx = int(round((e.x - sx) / self.scale))
        dy = int(round((e.y - sy) / self.scale))
        self.offset_x.set(clamp_int(ox + dx, OFFSET_MIN, OFFSET_MAX))
        self.offset_y.set(clamp_int(oy + dy, OFFSET_MIN, OFFSET_MAX))
        self._update_boxes()

    def _on_mouse_up(self, _e):
        self._dragging = False
        self._drag_start = None

    def reset_offset(self):
        self.offset_x.set(0)
        self.offset_y.set(0)
        self.factor.set(1.0)
        self.angle.set(0.0)
        self._update_boxes()

    def prev_page(self):
        if self.page_index > 0:
            self._load_page(self.page_index - 1)

    def next_page(self):
        if self.page_index < len(self.pages) - 1:
            self._load_page(self.page_index + 1)

    def save_page(self):
        if self.src_bgr is None:
            return

        out_dir = Path(OUTPUT_FOLDER)
        out_dir.mkdir(parents=True, exist_ok=True)

        sel_indices = self._selected_indices()

        dx = float(self.offset_x.get())
        dy = float(self.offset_y.get())
        factor = float(self.factor.get())
        ang_deg = float(self.angle.get())
        ang = math.radians(ang_deg)
        ca = math.cos(ang)
        sa = math.sin(ang)
        img_cx = self.src_w / 2.0
        img_cy = self.src_h / 2.0

        centers = self._grid_centers()
        half_w = CHAR_W / 2.0
        half_h = CHAR_H / 2.0
        half_w_s = half_w * factor
        half_h_s = half_h * factor

        def tx(x: float, y: float):
            vx = x - img_cx
            vy = y - img_cy
            rx = vx * ca - vy * sa
            ry = vx * sa + vy * ca
            return img_cx + factor * rx + dx, img_cy + factor * ry + dy

        stem = self.pages[self.page_index].stem
        for idx in sel_indices:
            cx, cy = centers[idx]
            r = idx // COLS
            c = idx % COLS
            ccx, ccy = tx(cx, cy)

            p0 = (ccx + (-half_w_s) * ca - (-half_h_s) * sa, ccy + (-half_w_s) * sa + (-half_h_s) * ca)
            p1 = (ccx + (+half_w_s) * ca - (-half_h_s) * sa, ccy + (+half_w_s) * sa + (-half_h_s) * ca)
            p2 = (ccx + (+half_w_s) * ca - (+half_h_s) * sa, ccy + (+half_w_s) * sa + (+half_h_s) * ca)
            p3 = (ccx + (-half_w_s) * ca - (+half_h_s) * sa, ccy + (-half_w_s) * sa + (+half_h_s) * ca)

            src_pts = np.array([p0, p1, p2, p3], dtype=np.float32)
            dst_pts = np.array(
                [[0, 0], [CHAR_W - 1, 0], [CHAR_W - 1, CHAR_H - 1], [0, CHAR_H - 1]],
                dtype=np.float32,

            )
            m = cv2.getPerspectiveTransform(src_pts, dst_pts)
            crop = cv2.warpPerspective(self.src_bgr, m, (int(CHAR_W), int(CHAR_H)), flags=cv2.INTER_LINEAR, borderValue=(255, 255, 255))
            fname = f"{stem}_r{r + 1}_c{c + 1}.png"
            imwrite_unicode(str(out_dir / fname), crop)

        self._set_notice(f"已保存 {stem}（{len(sel_indices)} 个）", ok=True)
        self._update_boxes()


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