#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
corridor_state_labeler.py
-------------------------------------------------
Rotulador rápido de estado global do corredor para dataset de segmentação.

Objetivo
- Percorrer uma raiz de dataset no formato:

    dataset_root/
    ├── grupo_1/
    │   ├── images/
    │   ├── masks/
    │   └── labels/        # criado automaticamente
    ├── grupo_2/
    │   ├── images/
    │   ├── masks/
    │   └── labels/
    └── grupo_3/
        ├── images/
        ├── masks/
        └── labels/

- Mostrar imagem original e máscara lado a lado.
- Permitir classificar cada frame com teclas numéricas.
- Salvar um JSON por imagem dentro de labels/.
- Suportar resume, skip, voltar, avançar e undo.

Exemplo de labels globais:
    Direcionando,EntrandoRua,CaminhandoRua,SaindoRua

Dependências:
    pip install pillow

Uso com GUI:
    python corridor_state_labeler.py

Uso por linha de comando:
    python corridor_state_labeler.py --dataset-root "C:/.../oak-d/dataset/original/group" --states Direcionando EntrandoRua CaminhandoRua SaindoRua --resume
"""

import argparse
import json
import os
import sys
from dataclasses import dataclass
from datetime import datetime
from pathlib import Path
from typing import Dict, List, Optional, Tuple

import tkinter as tk
from tkinter import filedialog, messagebox

try:
    from PIL import Image, ImageTk, ImageOps
except Exception:
    print("ERRO: Pillow não encontrado. Instale com: pip install pillow", file=sys.stderr)
    raise


IMG_EXTS = [".png", ".jpg", ".jpeg", ".bmp", ".webp"]
MASK_EXTS = [".png", ".jpg", ".jpeg", ".bmp", ".webp"]

IMAGE_DIR_NAMES = ["images", "image", "imgs", "rgb"]
MASK_DIR_NAMES = ["masks", "mask", "segmentacao", "seg", "segmentations"]
LABEL_DIR_NAME = "labels"


@dataclass
class SampleItem:
    group: str
    image_path: Path
    mask_path: Path
    label_path: Path
    base: str


def normalize_stem(stem: str) -> str:
    """
    Remove sufixos comuns para tentar casar image e mask mesmo quando nomes diferem.
    Exemplos:
        001_rgb -> 001
        001_image -> 001
        001_mask -> 001
        001_segmentacao -> 001
    """
    suffixes = [
        "_rgb", "_RGB", "_Rgb",
        "_image", "_img", "_frame",
        "_mask", "_masks",
        "_seg", "_SEG", "_segment", "_segmentacao", "_Segmentacao",
    ]
    out = stem
    changed = True
    while changed:
        changed = False
        for sfx in suffixes:
            if out.endswith(sfx):
                out = out[: -len(sfx)]
                changed = True
                break
    return out


def natural_key(text: str):
    """Ordenação natural simples: 2 antes de 10."""
    import re
    parts = re.split(r"(\d+)", text)
    return [int(p) if p.isdigit() else p.lower() for p in parts]


def find_first_existing_dir(group_dir: Path, candidates: List[str]) -> Optional[Path]:
    for name in candidates:
        p = group_dir / name
        if p.is_dir():
            return p
    return None


def list_image_files(folder: Path, exts: List[str]) -> List[Path]:
    files = []
    for p in folder.iterdir():
        if p.is_file() and p.suffix.lower() in exts:
            files.append(p)
    files.sort(key=lambda x: natural_key(x.name))
    return files


def collect_samples(dataset_root: Path) -> List[SampleItem]:
    """
    Procura grupos dentro de dataset_root.
    Cada grupo precisa ter uma pasta images/ e masks/.
    Cria labels/ automaticamente.
    """
    if not dataset_root.exists() or not dataset_root.is_dir():
        raise FileNotFoundError(f"Raiz inválida: {dataset_root}")

    samples: List[SampleItem] = []

    group_dirs = [p for p in dataset_root.iterdir() if p.is_dir()]
    group_dirs.sort(key=lambda x: natural_key(x.name))

    for group_dir in group_dirs:
        image_dir = find_first_existing_dir(group_dir, IMAGE_DIR_NAMES)
        mask_dir = find_first_existing_dir(group_dir, MASK_DIR_NAMES)

        if image_dir is None or mask_dir is None:
            continue

        label_dir = group_dir / LABEL_DIR_NAME
        label_dir.mkdir(parents=True, exist_ok=True)

        image_files = list_image_files(image_dir, IMG_EXTS)
        mask_files = list_image_files(mask_dir, MASK_EXTS)

        mask_by_base: Dict[str, Path] = {}
        for m in mask_files:
            mask_by_base[normalize_stem(m.stem)] = m

        for img in image_files:
            base = normalize_stem(img.stem)
            mask = mask_by_base.get(base)
            if mask is None:
                # fallback: tenta mesmo stem exato
                exact_candidates = [m for m in mask_files if m.stem == img.stem]
                mask = exact_candidates[0] if exact_candidates else None

            if mask is None:
                continue

            label_path = label_dir / f"{base}.json"
            samples.append(
                SampleItem(
                    group=group_dir.name,
                    image_path=img,
                    mask_path=mask,
                    label_path=label_path,
                    base=base,
                )
            )

    samples.sort(key=lambda s: (natural_key(s.group), natural_key(s.base)))
    return samples


def read_existing_label(label_path: Path) -> Optional[Dict]:
    if not label_path.exists():
        return None
    try:
        with label_path.open("r", encoding="utf-8") as f:
            return json.load(f)
    except Exception:
        return None


def safe_relative(path: Path, root: Path) -> str:
    try:
        return str(path.resolve().relative_to(root.resolve())).replace("\\", "/")
    except Exception:
        return str(path).replace("\\", "/")


class CorridorStateLabelerApp:
    def __init__(
        self,
        dataset_root: Path,
        samples: List[SampleItem],
        states: List[str],
        resume: bool = True,
        display_height: int = 520,
        show_existing: bool = False,
    ):
        self.dataset_root = dataset_root
        self.states = states
        self.display_height = int(display_height)
        self.show_existing = show_existing

        if resume:
            self.samples = [s for s in samples if not s.label_path.exists()]
        else:
            self.samples = samples

        self.idx = 0
        self.history: List[Dict] = []

        self.root = tk.Tk()
        self.root.title("Agrobot - Corridor State Labeler")
        self.root.geometry("1600x880")
        self.root.bind("<Key>", self.on_key)

        self.top_frame = tk.Frame(self.root)
        self.top_frame.pack(side=tk.TOP, fill=tk.X)

        self.info_label = tk.Label(self.top_frame, text="", font=("Segoe UI", 11), anchor="w")
        self.info_label.pack(side=tk.LEFT, padx=10, pady=6, fill=tk.X, expand=True)

        self.legend_label = tk.Label(self.top_frame, text=self.build_legend_text(), font=("Segoe UI", 10), anchor="e")
        self.legend_label.pack(side=tk.RIGHT, padx=10, pady=6)

        self.image_frame = tk.Frame(self.root)
        self.image_frame.pack(side=tk.TOP, fill=tk.BOTH, expand=True)

        self.left_frame = tk.LabelFrame(self.image_frame, text="Imagem")
        self.left_frame.pack(side=tk.LEFT, fill=tk.BOTH, expand=True, padx=8, pady=8)

        self.right_frame = tk.LabelFrame(self.image_frame, text="Máscara")
        self.right_frame.pack(side=tk.LEFT, fill=tk.BOTH, expand=True, padx=8, pady=8)

        self.image_label = tk.Label(self.left_frame)
        self.image_label.pack(fill=tk.BOTH, expand=True)

        self.mask_label = tk.Label(self.right_frame)
        self.mask_label.pack(fill=tk.BOTH, expand=True)

        self.current_label_var = tk.StringVar(value="")
        self.current_label = tk.Label(self.root, textvariable=self.current_label_var, font=("Segoe UI", 11), anchor="w")
        self.current_label.pack(side=tk.BOTTOM, fill=tk.X, padx=8, pady=2)

        self.status_var = tk.StringVar(value="Pronto.")
        self.status_label = tk.Label(self.root, textvariable=self.status_var, font=("Segoe UI", 10), anchor="w")
        self.status_label.pack(side=tk.BOTTOM, fill=tk.X, padx=8, pady=4)

        self.footer = tk.Label(
            self.root,
            text="1..9/0=classificar | Espaço=saltar | →/n=próxima | ←/p=anterior | b=desfazer | r=recarregar | q/Esc=sair",
            font=("Segoe UI", 10),
        )
        self.footer.pack(side=tk.BOTTOM, fill=tk.X, pady=2)

        self.render()

    def build_legend_text(self) -> str:
        parts = []
        for i, state in enumerate(self.states, start=1):
            key = i if i <= 9 else 0
            parts.append(f"[{key}] {state}")
        return " | ".join(parts)

    def load_and_resize(self, path: Path) -> ImageTk.PhotoImage:
        img = Image.open(path).convert("RGB")
        img = ImageOps.exif_transpose(img)

        w, h = img.size
        if h <= 0:
            raise ValueError(f"Imagem inválida: {path}")

        new_h = self.display_height
        new_w = max(1, int(w * (new_h / h)))
        img = img.resize((new_w, new_h), Image.BILINEAR)
        return ImageTk.PhotoImage(img)

    def render(self):
        if not self.samples:
            messagebox.showinfo("Fim", "Nenhum item para rotular. Talvez tudo já esteja rotulado com resume ligado.")
            self.root.destroy()
            return

        self.idx = max(0, min(self.idx, len(self.samples) - 1))
        item = self.samples[self.idx]

        try:
            self.tk_img = self.load_and_resize(item.image_path)
            self.tk_mask = self.load_and_resize(item.mask_path)
            self.image_label.configure(image=self.tk_img)
            self.mask_label.configure(image=self.tk_mask)
        except Exception as e:
            self.status_var.set(f"ERRO ao abrir imagem/máscara: {e}")
            return

        existing = read_existing_label(item.label_path)
        existing_txt = ""
        if existing:
            existing_txt = f" | Já rotulado: {existing.get('estado_corredor')}"

        self.info_label.configure(
            text=(
                f"{self.idx + 1}/{len(self.samples)} | Grupo: {item.group} | Base: {item.base} "
                f"| Imagem: {item.image_path.name} | Máscara: {item.mask_path.name}{existing_txt}"
            )
        )

        self.current_label_var.set(
            f"Arquivo de label: {safe_relative(item.label_path, self.dataset_root)}"
        )

    def save_label(self, state_index: int):
        if state_index < 0 or state_index >= len(self.states):
            return

        item = self.samples[self.idx]
        state = self.states[state_index]

        previous_content = None
        previous_existed = item.label_path.exists()
        if previous_existed:
            try:
                previous_content = item.label_path.read_text(encoding="utf-8")
            except Exception:
                previous_content = None

        rec = {
            "estado_corredor": state,
            "label_id": state_index,
            "states": self.states,
            "group": item.group,
            "base": item.base,
            "image": safe_relative(item.image_path, self.dataset_root),
            "mask": safe_relative(item.mask_path, self.dataset_root),
            "label": safe_relative(item.label_path, self.dataset_root),
            "timestamp": datetime.now().isoformat(timespec="seconds"),
            "source": "corridor_state_labeler",
        }

        try:
            item.label_path.parent.mkdir(parents=True, exist_ok=True)
            with item.label_path.open("w", encoding="utf-8") as f:
                json.dump(rec, f, ensure_ascii=False, indent=2)

            self.history.append(
                {
                    "action": "label",
                    "index": self.idx,
                    "label_path": item.label_path,
                    "previous_existed": previous_existed,
                    "previous_content": previous_content,
                    "new_state": state,
                }
            )

            self.status_var.set(f"Salvo: {item.group}/{item.base} → {state}")
            self.idx += 1
            if self.idx >= len(self.samples):
                messagebox.showinfo("Concluído", "Você chegou ao final da fila.")
                self.root.destroy()
                return
            self.render()

        except Exception as e:
            messagebox.showerror("Erro", f"Falha ao salvar label: {e}")
            self.status_var.set(f"ERRO ao salvar label: {e}")

    def skip(self):
        item = self.samples[self.idx]
        self.history.append({"action": "skip", "index": self.idx})
        self.status_var.set(f"Pulou: {item.group}/{item.base}")
        self.idx += 1
        if self.idx >= len(self.samples):
            messagebox.showinfo("Concluído", "Você chegou ao final da fila.")
            self.root.destroy()
            return
        self.render()

    def undo(self):
        if not self.history:
            self.status_var.set("Nada para desfazer.")
            return

        last = self.history.pop()
        action = last.get("action")

        if action == "label":
            label_path: Path = last["label_path"]
            try:
                if last.get("previous_existed"):
                    previous_content = last.get("previous_content")
                    if previous_content is not None:
                        label_path.write_text(previous_content, encoding="utf-8")
                else:
                    if label_path.exists():
                        label_path.unlink()

                self.idx = max(0, min(int(last.get("index", self.idx)), len(self.samples) - 1))
                self.status_var.set(f"Desfeito label: {label_path.name}")
                self.render()
            except Exception as e:
                messagebox.showerror("Erro", f"Falha ao desfazer: {e}")
                self.status_var.set(f"ERRO no undo: {e}")

        elif action == "skip":
            self.idx = max(0, min(int(last.get("index", self.idx)), len(self.samples) - 1))
            self.status_var.set("Desfeito pulo.")
            self.render()

    def next_item(self):
        self.idx = min(len(self.samples) - 1, self.idx + 1)
        self.status_var.set("Avançou uma imagem.")
        self.render()

    def prev_item(self):
        self.idx = max(0, self.idx - 1)
        self.status_var.set("Voltou uma imagem.")
        self.render()

    def on_key(self, event):
        ch = event.keysym.lower()

        if ch in [str(i) for i in range(1, 10)] or ch == "0":
            idx = 9 if ch == "0" else int(ch) - 1
            self.save_label(idx)
            return

        if ch in ("space", "s"):
            self.skip()
            return

        if ch in ("right", "n"):
            self.next_item()
            return

        if ch in ("left", "p"):
            self.prev_item()
            return

        if ch == "b":
            self.undo()
            return

        if ch == "r":
            self.status_var.set("Recarregado.")
            self.render()
            return

        if ch in ("q", "escape"):
            self.root.destroy()
            return

    def run(self):
        self.root.mainloop()


class SetupWindow:
    def __init__(self):
        self.root = tk.Tk()
        self.root.title("Configurar - Corridor State Labeler")
        self.root.geometry("760x430")

        frm_root = tk.LabelFrame(self.root, text="Raiz do dataset")
        frm_root.pack(fill=tk.X, padx=10, pady=8)

        self.dataset_root_var = tk.StringVar(value="")
        tk.Entry(frm_root, textvariable=self.dataset_root_var).pack(side=tk.LEFT, fill=tk.X, expand=True, padx=6, pady=6)
        tk.Button(frm_root, text="Escolher...", command=self.choose_dataset_root).pack(side=tk.RIGHT, padx=6, pady=6)

        frm_states = tk.LabelFrame(self.root, text="Estados do corredor separados por vírgula")
        frm_states.pack(fill=tk.X, padx=10, pady=8)

        self.states_var = tk.StringVar(value="Parado,EntrandoRua,CaminhandoRua,SaindoRua,Manobrando,Direcionando,RetornandoBase,Indefinido")
        tk.Entry(frm_states, textvariable=self.states_var).pack(fill=tk.X, padx=6, pady=6)

        frm_opts = tk.LabelFrame(self.root, text="Opções")
        frm_opts.pack(fill=tk.X, padx=10, pady=8)

        self.resume_var = tk.BooleanVar(value=True)
        self.height_var = tk.IntVar(value=520)

        tk.Checkbutton(frm_opts, text="Retomar: pular imagens que já possuem labels/*.json", variable=self.resume_var).pack(anchor="w", padx=6, pady=4)

        frm_height = tk.Frame(frm_opts)
        frm_height.pack(fill=tk.X, padx=6, pady=4)
        tk.Label(frm_height, text="Altura de exibição em px:").pack(side=tk.LEFT)
        tk.Entry(frm_height, textvariable=self.height_var, width=8).pack(side=tk.LEFT, padx=6)

        self.preview_var = tk.StringVar(value="")
        tk.Label(self.root, textvariable=self.preview_var, fg="#555", anchor="w", justify="left").pack(fill=tk.X, padx=12, pady=4)

        btn_frame = tk.Frame(self.root)
        btn_frame.pack(fill=tk.X, padx=10, pady=10)
        tk.Button(btn_frame, text="Verificar dataset", command=self.preview_dataset).pack(side=tk.LEFT, padx=4)
        tk.Button(btn_frame, text="Iniciar classificação", command=self.start).pack(side=tk.RIGHT, padx=4)

        tk.Label(
            self.root,
            text="Teclas: 1..9/0=classificar | Espaço=saltar | b=desfazer | ←/→ navegar | q/Esc=sair",
            fg="#555",
        ).pack(pady=4)

        self.result = None

    def choose_dataset_root(self):
        p = filedialog.askdirectory(title="Selecione a raiz do dataset")
        if p:
            self.dataset_root_var.set(p)
            self.preview_dataset()

    def parse_states(self) -> List[str]:
        return [s.strip() for s in self.states_var.get().split(",") if s.strip()]

    def preview_dataset(self):
        root_raw = self.dataset_root_var.get().strip()
        if not root_raw:
            self.preview_var.set("Selecione uma raiz de dataset.")
            return
        root = Path(root_raw)
        try:
            samples = collect_samples(root)
            groups = sorted(set(s.group for s in samples), key=natural_key)
            labeled = sum(1 for s in samples if s.label_path.exists())
            self.preview_var.set(
                f"Amostras encontradas: {len(samples)} | Já rotuladas: {labeled} | Grupos: {', '.join(groups) if groups else '-'}"
            )
        except Exception as e:
            self.preview_var.set(f"Erro ao verificar dataset: {e}")

    def start(self):
        root_raw = self.dataset_root_var.get().strip()
        states = self.parse_states()

        if not root_raw:
            messagebox.showwarning("Faltando raiz", "Selecione a raiz do dataset.")
            return
        if not states:
            messagebox.showwarning("Faltando estados", "Informe pelo menos um estado.")
            return
        if len(states) > 10:
            messagebox.showwarning("Muitos estados", "Este rotulador suporta até 10 estados nas teclas 1..9 e 0.")
            return

        root = Path(root_raw)
        try:
            samples = collect_samples(root)
        except Exception as e:
            messagebox.showerror("Erro", f"Falha ao ler dataset: {e}")
            return

        if not samples:
            messagebox.showinfo(
                "Sem amostras",
                "Nenhum par image/mask encontrado. Verifique se cada grupo possui images/ e masks/.",
            )
            return

        self.result = {
            "dataset_root": root,
            "states": states,
            "resume": self.resume_var.get(),
            "display_height": self.height_var.get(),
        }
        self.root.destroy()

    def run(self):
        self.root.mainloop()
        return self.result


def run_gui_setup():
    setup = SetupWindow()
    res = setup.run()
    if not res:
        return

    dataset_root = res["dataset_root"]
    samples = collect_samples(dataset_root)

    app = CorridorStateLabelerApp(
        dataset_root=dataset_root,
        samples=samples,
        states=res["states"],
        resume=res["resume"],
        display_height=res["display_height"],
    )
    app.run()


def main():
    parser = argparse.ArgumentParser(description="Rotulador de estado global do corredor para dataset de segmentação.")
    parser.add_argument("--dataset-root", help="Raiz do dataset contendo grupos com images/ e masks/.")
    parser.add_argument("--states", nargs="+", help="Estados do corredor. Ex: Direcionando EntrandoRua CaminhandoRua SaindoRua")
    parser.add_argument("--resume", action="store_true", help="Pular itens que já possuem labels/*.json.")
    parser.add_argument("--no-resume", action="store_true", help="Não pular itens já rotulados.")
    parser.add_argument("--display-height", type=int, default=520, help="Altura de exibição das imagens em px.")
    parser.add_argument("--no-gui-setup", action="store_true", help="Não abrir janela de configuração.")

    args = parser.parse_args()

    need_gui = not args.no_gui_setup and (not args.dataset_root or not args.states)
    if need_gui:
        run_gui_setup()
        return

    if not args.dataset_root or not args.states:
        print("ERRO: informe --dataset-root e --states, ou rode sem argumentos para abrir a GUI.", file=sys.stderr)
        sys.exit(1)

    dataset_root = Path(args.dataset_root)
    states = args.states
    if len(states) > 10:
        print("ERRO: máximo de 10 estados: teclas 1..9 e 0.", file=sys.stderr)
        sys.exit(1)

    samples = collect_samples(dataset_root)
    if not samples:
        print("ERRO: nenhum par image/mask encontrado.", file=sys.stderr)
        sys.exit(1)

    resume = True
    if args.no_resume:
        resume = False
    elif args.resume:
        resume = True

    app = CorridorStateLabelerApp(
        dataset_root=dataset_root,
        samples=samples,
        states=states,
        resume=resume,
        display_height=args.display_height,
    )
    app.run()


if __name__ == "__main__":
    main()