#!/opt/homebrew/bin/python3.13
"""Sprite frame tagger — navigate frames, assign labels, copy to labeled dirs.

Usage:
    sprite_tagger.py <frames_dir> <label1> [label2 ...] [--out <output_dir>]

Example:
    sprite_tagger.py ./sprites idle walk run jump

Keys:
    ← →         previous / next frame
    1-9         assign label and advance
    N           add a new label
    Space       play/pause all frames
    P           play/pause only the current frame's label (looping)
    [ ]         slower / faster playback
    X           clear label on current frame
    Enter       export labeled frames to output dir
"""

import sys
import shutil
import tkinter as tk
from tkinter import messagebox
from pathlib import Path
from PIL import Image, ImageDraw, ImageTk

import labelstore
import playback

IMAGE_EXTENSIONS = {'.png', '.jpg', '.jpeg', '.gif', '.bmp', '.webp', '.tga'}

BG        = '#1a1a2e'
BG_PANEL  = '#16213e'
ACCENT    = '#4ecca3'
FG        = '#e0e0e0'
FG_DIM    = '#556677'
BTN_BG    = '#0f3460'
BTN_ACT   = '#e94560'
FONT_MONO = ('Courier', 11)
FONT_TINY = ('Courier', 9)

CANVAS_W  = 860
CANVAS_H  = 620
FPS_DEFAULT = 10
FPS_MIN     = 1
FPS_MAX     = 30


def load_frames(directory: Path) -> list:
    return sorted(
        p for p in directory.iterdir()
        if p.suffix.lower() in IMAGE_EXTENSIONS
    )


class SpriteTagger:
    def __init__(self, root: tk.Tk, frames: list, labels: list, output_dir: Path,
                 store_path: Path, assignments: dict = None):
        self.root       = root
        self.frames     = frames
        self.labels     = labels
        self.output_dir = output_dir
        self.store_path = store_path
        self.current    = 0
        self.assignments: dict = assignments or {}   # frame index → label name
        self._photo     = None        # ImageTk ref — must not be GC'd
        self._playing   = False
        self._fps       = FPS_DEFAULT
        self._play_job  = None
        self._preview_label = None    # None = play all; label str = preview that label

        self._build_ui()
        self._bind_keys()
        self._show_frame()

    # ── UI construction ──────────────────────────────────────────────────────

    def _build_ui(self):
        self.root.title('Sprite Tagger')
        self.root.configure(bg=BG)
        self.root.resizable(True, True)

        # canvas
        self.canvas = tk.Canvas(
            self.root, bg=BG_PANEL,
            width=CANVAS_W, height=CANVAS_H,
            highlightthickness=1, highlightbackground='#2a2a4e',
        )
        self.canvas.pack(fill=tk.BOTH, expand=True, padx=10, pady=(10, 0))
        self.canvas.bind('<Configure>', lambda _e: self._show_frame())

        # status row
        status = tk.Frame(self.root, bg=BG)
        status.pack(fill=tk.X, padx=10, pady=4)

        self._var_info  = tk.StringVar()
        self._var_label = tk.StringVar()
        self._var_fps   = tk.StringVar()

        tk.Label(status, textvariable=self._var_info,  bg=BG, fg=FG,     font=FONT_MONO).pack(side=tk.LEFT)
        tk.Label(status, textvariable=self._var_fps,   bg=BG, fg=FG_DIM, font=FONT_TINY).pack(side=tk.LEFT, padx=10)
        tk.Label(status, textvariable=self._var_label, bg=BG, fg=ACCENT,  font=(*FONT_MONO, 'bold')).pack(side=tk.RIGHT)

        # label buttons row
        self._btn_row = tk.Frame(self.root, bg=BG)
        self._btn_row.pack(fill=tk.X, padx=10, pady=(0, 4))

        self._label_btns = {}
        for label in self.labels:
            self._add_label_button(label)

        self._make_utility_buttons()

        # hint row
        tk.Label(
            self.root,
            text='← → navigate   1-9 label+advance   N new label   Space play all   P preview label   [ ] speed   X clear   ↵ export',
            bg=BG, fg=FG_DIM, font=FONT_TINY,
        ).pack(pady=(0, 6))

    def _add_label_button(self, label: str):
        idx = len(self._label_btns)
        key = str(idx + 1) if idx < 9 else '-'
        # tk.Label used instead of tk.Button — macOS ignores fg/bg on native buttons
        lbl = tk.Label(
            self._btn_row, text=f'[{key}] {label}',
            bg=BTN_BG, fg='white',
            padx=8, pady=4, font=FONT_MONO,
            cursor='hand2', relief=tk.FLAT,
        )
        lbl.pack(side=tk.LEFT, padx=2)
        lbl.bind('<Button-1>', lambda _e, l=label: self._assign(l))
        lbl.bind('<Enter>',    lambda _e, w=lbl: w.configure(bg='#1a5a9a'))
        lbl.bind('<Leave>',    lambda _e, w=lbl: self._restore_btn_color(w, label))
        self._label_btns[label] = lbl

    def _make_utility_buttons(self):
        self._make_action_label(self._btn_row, '[X] clear',   '#5a1a3a', '#ff6688', self._clear,              side=tk.LEFT,  padx=(8, 2))
        self._make_action_label(self._btn_row, '[N] + label', '#3a3a1a', '#dddd44', self._prompt_new_label,   side=tk.LEFT,  padx=2)
        self._make_action_label(self._btn_row, '[↵] export',  '#0a3a1a', ACCENT,    self._export,             side=tk.RIGHT, padx=2, bold=True)

    def _make_action_label(self, parent, text, bg, hover_bg, command, side, padx, bold=False):
        font = (*FONT_MONO, 'bold') if bold else FONT_MONO
        w = tk.Label(parent, text=text, bg=bg, fg='white', padx=8, pady=4, font=font, cursor='hand2', relief=tk.FLAT)
        w.pack(side=side, padx=padx)
        w.bind('<Button-1>', lambda _e: command())
        w.bind('<Enter>',    lambda _e: w.configure(bg=hover_bg, fg='black'))
        w.bind('<Leave>',    lambda _e: w.configure(bg=bg,       fg='white'))

    def _restore_btn_color(self, widget: tk.Label, label: str):
        current_label = self.assignments.get(self.current)
        if label == current_label:
            widget.configure(bg=ACCENT, fg='black')
        else:
            widget.configure(bg=BTN_BG, fg='white')

    def _bind_keys(self):
        self.root.bind('<Right>',  lambda _e: self._navigate(1))
        self.root.bind('<Left>',   lambda _e: self._navigate(-1))
        self.root.bind('<Return>', lambda _e: self._export())
        self.root.bind('<space>',  lambda _e: self._toggle_play())
        self.root.bind('p',        lambda _e: self._toggle_preview())
        self.root.bind('P',        lambda _e: self._toggle_preview())
        self.root.bind('[',        lambda _e: self._change_fps(-2))
        self.root.bind(']',        lambda _e: self._change_fps(+2))
        self.root.bind('x',        lambda _e: self._clear())
        self.root.bind('X',        lambda _e: self._clear())
        self.root.bind('n',        lambda _e: self._prompt_new_label())
        self.root.bind('N',        lambda _e: self._prompt_new_label())
        self._rebind_number_keys()

    def _rebind_number_keys(self):
        for i, label in enumerate(self.labels):
            if i < 9:
                self.root.bind(str(i + 1), lambda _e, l=label: self._assign(l))

    # ── Frame display ────────────────────────────────────────────────────────

    def _show_frame(self):
        path  = self.frames[self.current]
        img   = Image.open(path).convert('RGBA')

        cw = self.canvas.winfo_width()  if self.canvas.winfo_width()  > 1 else CANVAS_W
        ch = self.canvas.winfo_height() if self.canvas.winfo_height() > 1 else CANVAS_H

        if img.width <= cw // 2 and img.height <= ch // 2:
            # small sprite — zoom with nearest-neighbor to preserve pixels
            zoom = min(cw // img.width, ch // img.height)
            zoom = max(zoom, 1)
            img = img.resize((img.width * zoom, img.height * zoom), Image.NEAREST)
        elif img.width > cw or img.height > ch:
            # large image — scale down with lanczos
            scale = min(cw / img.width, ch / img.height)
            img = img.resize((int(img.width * scale), int(img.height * scale)), Image.LANCZOS)

        # checkerboard background for transparency
        check_size = max(8, min(img.width, img.height) // 16)
        bg = Image.new('RGBA', img.size, (40, 40, 60, 255))
        draw = ImageDraw.Draw(bg)
        for y in range(0, img.height, check_size):
            for x in range(0, img.width, check_size):
                if (x // check_size + y // check_size) % 2 == 0:
                    draw.rectangle(
                        [x, y, min(x + check_size, img.width) - 1, min(y + check_size, img.height) - 1],
                        fill=(50, 50, 70, 255),
                    )
        bg.paste(img, mask=img.split()[3])

        self._photo = ImageTk.PhotoImage(bg)
        self.canvas.delete('all')
        self.canvas.create_image(cw // 2, ch // 2, image=self._photo, anchor='center')

        self._update_status()

    def _update_status(self):
        path    = self.frames[self.current]
        labeled = len(self.assignments)
        current_label = self.assignments.get(self.current)

        self._var_info.set(f'{path.name}  [{self.current + 1}/{len(self.frames)}]  labeled: {labeled}')
        self._var_label.set(f'▶ {current_label}' if current_label else '—')
        play_sym = '⏸' if self._playing else '▶'
        if self._preview_label is not None:
            n = len(playback.frames_with_label(self.assignments, self._preview_label))
            self._var_fps.set(f'{play_sym} preview: {self._preview_label} ({n}) ⟳  {self._fps}fps')
        else:
            self._var_fps.set(f'{play_sym} {self._fps}fps')

        for label, btn in self._label_btns.items():
            if label == current_label:
                btn.configure(bg=ACCENT, fg='black')
            else:
                btn.configure(bg=BTN_BG, fg='white')

    # ── Navigation ───────────────────────────────────────────────────────────

    def _navigate(self, delta: int):
        self.current = (self.current + delta) % len(self.frames)
        self._show_frame()

    # ── Labeling ─────────────────────────────────────────────────────────────

    def _assign(self, label: str):
        self.assignments[self.current] = label
        self._save()
        self._navigate(1)

    def _clear(self):
        self.assignments.pop(self.current, None)
        self._save()
        self._update_status()

    def _save(self):
        by_name = labelstore.to_name_keyed(self.assignments, self.frames)
        labelstore.save_labels(self.store_path, by_name, self.labels)

    # ── Playback ─────────────────────────────────────────────────────────────

    def _toggle_play(self):
        self._preview_label = None   # Space always plays all frames
        if self._playing:
            self._stop_play()
        else:
            self._start_play()

    def _toggle_preview(self):
        if self._preview_label is not None:
            self._stop_play()
            self._preview_label = None
            self._update_status()
            return
        label = self.assignments.get(self.current)
        if label is None:
            self._flash('tag this frame first to preview its label')
            return
        self._preview_label = label
        indices = playback.frames_with_label(self.assignments, label)
        self.current = indices[0]
        self._show_frame()
        self._start_play()

    def _start_play(self):
        self._playing = True
        delay = max(1, 1000 // self._fps)
        self._play_job = self.root.after(delay, self._tick)
        self._update_status()

    def _stop_play(self):
        self._playing = False
        if self._play_job:
            self.root.after_cancel(self._play_job)
            self._play_job = None
        self._update_status()

    def _tick(self):
        if not self._playing:
            return
        if self._preview_label is not None:
            self._step_preview(1)
        else:
            self._navigate(1)
        delay = max(1, 1000 // self._fps)
        self._play_job = self.root.after(delay, self._tick)

    def _step_preview(self, delta: int):
        indices = playback.frames_with_label(self.assignments, self._preview_label)
        if not indices:
            self._stop_play()
            self._preview_label = None
            return
        self.current = playback.next_in_cycle(indices, self.current, delta)
        self._show_frame()

    def _change_fps(self, delta: int):
        self._fps = max(FPS_MIN, min(FPS_MAX, self._fps + delta))
        self._update_status()

    def _flash(self, msg: str):
        self._var_label.set(msg)
        self.root.after(1500, self._update_status)

    # ── Label management ─────────────────────────────────────────────────────

    def _prompt_new_label(self):
        from tkinter import simpledialog
        name = simpledialog.askstring('New label', 'Label name:', parent=self.root)
        if not name:
            return
        name = name.strip().lower().replace(' ', '_')
        if not name or name in self._label_btns:
            return
        self.labels.append(name)
        self._add_label_button(name)
        self._rebind_number_keys()
        self._save()

    # ── Export ───────────────────────────────────────────────────────────────

    def _export(self):
        if self._playing:
            self._toggle_play()

        if not self.assignments:
            messagebox.showinfo('Nothing to export', 'No frames have been labeled yet.')
            return

        counts: dict = {}
        for idx, label in self.assignments.items():
            dest = self.output_dir / label
            dest.mkdir(parents=True, exist_ok=True)
            shutil.copy2(self.frames[idx], dest / self.frames[idx].name)
            counts[label] = counts.get(label, 0) + 1

        summary = '\n'.join(f'  {label}/: {n} frame{"s" if n != 1 else ""}' for label, n in sorted(counts.items()))
        messagebox.showinfo('Export complete', f'Copied to:\n{self.output_dir}\n\n{summary}')


# ── Entry point ───────────────────────────────────────────────────────────────

def _parse_args(argv: list):
    if len(argv) < 3:
        print(__doc__)
        sys.exit(1)

    args = list(argv[1:])
    output_dir = None

    if '--out' in args:
        i = args.index('--out')
        if i + 1 >= len(args):
            print('Error: --out requires a path argument')
            sys.exit(1)
        output_dir = Path(args[i + 1])
        args = args[:i] + args[i + 2:]

    frames_dir = Path(args[0])
    labels = args[1:]

    if not frames_dir.is_dir():
        print(f'Error: {frames_dir} is not a directory')
        sys.exit(1)
    if not labels:
        print('Error: provide at least one label')
        sys.exit(1)

    if output_dir is None:
        output_dir = frames_dir / 'labeled'

    return frames_dir, labels, output_dir


def main():
    frames_dir, labels, output_dir = _parse_args(sys.argv)

    frames = load_frames(frames_dir)
    if not frames:
        print(f'No image files found in {frames_dir}')
        print(f'Supported formats: {", ".join(sorted(IMAGE_EXTENSIONS))}')
        sys.exit(1)

    store_path = labelstore.sidecar_path(frames_dir)
    saved_assignments, saved_labels = labelstore.load_labels(store_path)
    for lbl in saved_labels:
        if lbl not in labels:
            labels.append(lbl)
    assignments = labelstore.to_index_keyed(saved_assignments, frames)

    print(f'Loaded {len(frames)} frames')
    print(f'Labels : {", ".join(labels)}')
    print(f'Output : {output_dir}')
    if assignments:
        print(f'Resumed {len(assignments)} labels from {store_path.name}')

    root = tk.Tk()
    root.geometry(f'{CANVAS_W + 20}x{CANVAS_H + 130}')
    root.minsize(400, 400)
    SpriteTagger(root, frames, labels, output_dir, store_path, assignments)
    root.mainloop()


if __name__ == '__main__':
    main()