import json
import os
import argparse
from pathlib import Path

import numpy as np
import cv2
from PIL import Image, ImageDraw
import matplotlib.pyplot as plt
from utils import carregar_labelmap_completo


# =========================
# CONFIGURAÇÃO DE CLASSES
# =========================
# Ajuste aqui conforme suas máscaras:
# - Se sua máscara for "indexada" (modo P) ou grayscale com IDs por pixel:
#   class_ids = {0: 1, 1: 2}  # exemplo: erva=1, cana=2
# - Se sua máscara for RGB com cores fixas:
#   class_colors = {0: (0,255,0), 1: (0,0,255)}  # exemplo
#
# Por padrão abaixo: tenta RGB primeiro; se a máscara vier indexada, usa IDs.
def build_maps_from_labelmap(alpha: int = 90, ignore_names=None):
    """
    Lê labelmap e constrói maps dinâmicos:
      - class_names: {new_id: name}
      - class_colors_rgb: {new_id: (r,g,b)}
      - class_ids: {new_id: new_id}  (para máscaras indexed alinhadas com o new_id)
      - overlay_rgba: {new_id: (r,g,b,alpha)}
      - ignore_rgb: cor da classe "ignore" (se existir no labelmap original)
      - id_old_to_new: {old_id: new_id}  (útil se sua máscara indexed usa ids antigos)
      - id_new_to_old: {new_id: old_id}
    """
    if ignore_names is None:
        ignore_names = []

    ignore_set = {n.strip().lower() for n in ignore_names if n and n.strip()}

    with open("config.json", "r", encoding="utf-8") as f:
        config = json.load(f)

    MODELO = config["camera"]
    labelmap_path = os.path.join(MODELO, "dataset", "labelmap.txt")

    # Lê tudo do labelmap (mantém seus ids originais)
    cor_para_id, cores_rgb, id_para_nome, ignore_rgb = carregar_labelmap_completo(labelmap_path)

    # --- Filtra classes por nome ---
    kept_old_ids = []
    for old_id, name in id_para_nome.items():
        if name.strip().lower() in ignore_set:
            continue
        kept_old_ids.append(old_id)

    # Reindexa para ficar 0..N-1
    kept_old_ids = sorted(kept_old_ids)
    id_old_to_new = {old_id: new_id for new_id, old_id in enumerate(kept_old_ids)}
    id_new_to_old = {new_id: old_id for old_id, new_id in id_old_to_new.items()}

    # Constrói maps novos (compactos)
    class_names = {}
    class_colors_rgb = {}
    overlay_rgba = {}

    # cor_para_id: { (r,g,b): old_id }
    # id_para_nome: { old_id: name }
    for cor_rgb, old_id in cor_para_id.items():
        if old_id not in id_old_to_new:
            continue
        new_id = id_old_to_new[old_id]
        class_names[new_id] = id_para_nome[old_id]
        class_colors_rgb[new_id] = cor_rgb
        overlay_rgba[new_id] = (cor_rgb[0], cor_rgb[1], cor_rgb[2], alpha)

    # Para máscara indexed:
    # - Se sua máscara indexed já usa os IDs NOVOS (compactos), isso aqui está ok.
    # - Se ela usa IDs ANTIGOS, você precisa mapear (old -> new) antes de extrair polígonos.
    class_ids = {new_id: new_id for new_id in class_names.keys()}

    return {
        "class_names": class_names,
        "class_colors_rgb": class_colors_rgb,
        "class_ids": class_ids,
        "overlay_rgba": overlay_rgba,
        "ignore_rgb": ignore_rgb,
        "id_old_to_new": id_old_to_new,
        "id_new_to_old": id_new_to_old,
        "labelmap_path": labelmap_path,
    }

maps = build_maps_from_labelmap(ignore_names=["chao", "ignore"])

CLASS_NAMES = maps["class_names"]
CLASS_IDS = maps["class_ids"]
CLASS_COLORS_RGB = maps["class_colors_rgb"]
OVERLAY_RGBA = maps["overlay_rgba"]
IGNORE_RGB = maps["ignore_rgb"]


# =========================
# UTILITÁRIOS
# =========================
def imread_unicode(path: Path) -> np.ndarray:
    """Lê imagem com caminho unicode no Windows."""
    data = np.fromfile(str(path), dtype=np.uint8)
    img = cv2.imdecode(data, cv2.IMREAD_UNCHANGED)
    return img


def load_mask(mask_path: Path):
    """
    Retorna:
      mask_type: 'indexed' ou 'rgb'
      mask_data:
        - indexed: np.ndarray (H,W) int
        - rgb:     np.ndarray (H,W,3) uint8 em RGB
    """
    pil = Image.open(mask_path)
    if pil.mode == "P":
        arr = np.array(pil, dtype=np.int32)
        return "indexed", arr
    if pil.mode in ("L", "I;16"):
        arr = np.array(pil, dtype=np.int32)
        return "indexed", arr

    # RGB/RGBA
    pil = pil.convert("RGBA")
    rgba = np.array(pil, dtype=np.uint8)
    rgb = rgba[:, :, :3]
    return "rgb", rgb


def class_binary_from_mask(mask_type, mask_data, cls, class_ids, class_colors, rgb_tol=10):
    """Gera máscara binária (uint8 0/255) para uma classe."""
    if mask_type == "indexed":
        target_id = class_ids[cls]
        bin_mask = (mask_data == target_id).astype(np.uint8) * 255
        return bin_mask

    # rgb
    target = np.array(class_colors[cls], dtype=np.int16)
    img = mask_data.astype(np.int16)
    diff = np.abs(img - target[None, None, :])
    ok = (diff[:, :, 0] <= rgb_tol) & (diff[:, :, 1] <= rgb_tol) & (diff[:, :, 2] <= rgb_tol)
    return ok.astype(np.uint8) * 255


def simplify_contour(cnt, epsilon_px=1.0, epsilon_rel=0.001):
    peri = cv2.arcLength(cnt, True)
    eps = max(epsilon_px, epsilon_rel * peri)
    return cv2.approxPolyDP(cnt, eps, True)


def contours_to_polygons(bin_mask, min_area_px=50, epsilon_px=2.0):
    """
    bin_mask: uint8 0/255
    Retorna lista de polígonos, cada um como array (N,2) em pixels (float).
    """
    # limpa ruído e fecha pequenos buracos
    kernel = np.ones((3, 3), np.uint8)
    m = cv2.morphologyEx(bin_mask, cv2.MORPH_OPEN, kernel, iterations=1)
    m = cv2.morphologyEx(m, cv2.MORPH_CLOSE, kernel, iterations=1)

    contours, _hier = cv2.findContours(m, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
    polys = []
    for cnt in contours:
        area = cv2.contourArea(cnt)
        if area < min_area_px:
            continue
        approx = simplify_contour(cnt, epsilon_px=epsilon_px)
        if len(approx) < 3:
            continue
        pts = approx.reshape(-1, 2).astype(np.float32)
        polys.append(pts)
    return polys


def polygon_px_to_yolo(poly_px, w, h):
    """(N,2) px -> lista [x1,y1,x2,y2,...] normalizada 0..1"""
    xs = np.clip(poly_px[:, 0] / float(w), 0.0, 1.0)
    ys = np.clip(poly_px[:, 1] / float(h), 0.0, 1.0)
    coords = []
    for x, y in zip(xs, ys):
        coords.append(float(x))
        coords.append(float(y))
    return coords


def draw_polygons_on_preview(preview_path: Path, polygons_by_class, out_path: Path):
    """Cria overlay (PIL) com polígonos extraídos por cima do preview."""
    img = Image.open(preview_path).convert("RGB")
    w, h = img.size
    draw = ImageDraw.Draw(img, "RGBA")

    for cls, polys in polygons_by_class.items():
        color = OVERLAY_RGBA.get(cls, (255, 255, 255, 90))
        outline = color[:3] + (255,)
        for poly in polys:
            pts = [(float(x), float(y)) for x, y in poly]
            if len(pts) >= 3:
                draw.polygon(pts, fill=color, outline=outline)

    img.save(out_path)


def make_triview(preview_path: Path, mask_path: Path, overlay_path: Path, out_path: Path, title: str = ""):
    """Salva uma imagem com 3 colunas: preview | mask | overlay."""
    prev = Image.open(preview_path).convert("RGB")
    msk = Image.open(mask_path).convert("RGB")
    ovl = Image.open(overlay_path).convert("RGB")

    fig = plt.figure(figsize=(16, 6))
    fig.suptitle(title, fontsize=12)

    ax1 = fig.add_subplot(1, 3, 1)
    ax1.imshow(prev)
    ax1.set_title("Preview")
    ax1.axis("off")

    ax2 = fig.add_subplot(1, 3, 2)
    ax2.imshow(msk)
    ax2.set_title("Mask (manual)")
    ax2.axis("off")

    ax3 = fig.add_subplot(1, 3, 3)
    ax3.imshow(ovl)
    ax3.set_title("Overlay (polígonos extraídos)")
    ax3.axis("off")

    plt.tight_layout()
    fig.savefig(out_path, dpi=140)
    plt.close(fig)


def find_matching_preview(previews_dir: Path, stem: str):
    """Procura preview com mesmo stem em extensões comuns."""
    for ext in [".jpg", ".jpeg", ".png", ".bmp", ".tif", ".tiff"]:
        p = previews_dir / f"{stem}{ext}"
        if p.exists():
            return p
    return None


# =========================
# PIPELINE PRINCIPAL
# =========================
def process_dataset(root_dir: Path, out_labels_dir: Path, out_vis_dir: Path,
                    class_ids, class_colors, rgb_tol=10,
                    min_area_px=50, epsilon_px=2.0):
    previews_dir = root_dir / "previews"
    masks_dir = root_dir / "masks"

    if not previews_dir.exists() or not masks_dir.exists():
        raise FileNotFoundError(f"Esperado encontrar previews/ e masks/ dentro de {root_dir}")

    out_labels_dir.mkdir(parents=True, exist_ok=True)
    out_vis_dir.mkdir(parents=True, exist_ok=True)

    mask_files = sorted(list(masks_dir.glob("*.png")) + list(masks_dir.glob("*.jpg")) + list(masks_dir.glob("*.jpeg")))
    if not mask_files:
        print(f"[WARN] Nenhuma máscara encontrada em: {masks_dir}")
        return

    total = 0
    for mask_path in mask_files:
        stem = mask_path.stem
        preview_path = find_matching_preview(previews_dir, stem)
        if preview_path is None:
            print(f"[WARN] Sem preview para máscara: {mask_path.name}")
            continue

        # tamanhos
        prev_img = Image.open(preview_path)
        w, h = prev_img.size

        mask_type, mask_data = load_mask(mask_path)

        polygons_by_class = {}
        yolo_lines = []

        for cls in sorted(CLASS_NAMES.keys()):
            if cls not in class_ids or cls not in class_colors:
                continue

            bin_mask = class_binary_from_mask(mask_type, mask_data, cls, class_ids, class_colors, rgb_tol=rgb_tol)
            polys = contours_to_polygons(bin_mask, min_area_px=min_area_px, epsilon_px=epsilon_px)
            if not polys:
                continue

            polygons_by_class[cls] = polys

            for poly_px in polys:
                coords = polygon_px_to_yolo(poly_px, w, h)
                # YOLOv8-seg exige pelo menos 3 pontos (6 nums)
                if len(coords) >= 6:
                    line = str(cls) + " " + " ".join(f"{v:.6f}" for v in coords)
                    yolo_lines.append(line)

        # salva label
        label_path = out_labels_dir / f"{stem}.txt"
        label_path.write_text("\n".join(yolo_lines) + ("\n" if yolo_lines else ""), encoding="utf-8")

        # gera overlay e triview
        overlay_path = out_vis_dir / f"{stem}_overlay.png"
        triview_path = out_vis_dir / f"{stem}_triview.png"
        draw_polygons_on_preview(preview_path, polygons_by_class, overlay_path)
        make_triview(preview_path, mask_path, overlay_path, triview_path, title=stem)

        total += 1
        print(f"[OK] {stem}: polys={sum(len(v) for v in polygons_by_class.values())} -> {label_path.name}")

    print(f"\nFeito ✅  Processados: {total} arquivos")
    print(f"Labels: {out_labels_dir}")
    print(f"Vis:    {out_vis_dir}")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--root", type=str, required=True, help="Pasta raiz no formato antigo (contendo previews/ masks/ raws/ metas/)")

    ap.add_argument("--rgb_tol", type=int, default=10, help="Tolerância p/ match de cor RGB na máscara")
    ap.add_argument("--min_area", type=int, default=50, help="Área mínima (px) pra descartar sujeira")
    ap.add_argument("--eps", type=float, default=2.0, help="Epsilon (px) pra simplificar polígonos")

    args = ap.parse_args()

    root_dir = Path(args.root)
    out_labels_dir = Path(f"{args.root}/labels")
    out_vis_dir = Path(f"{args.root}/vis")

    process_dataset(
        root_dir=root_dir,
        out_labels_dir=out_labels_dir,
        out_vis_dir=out_vis_dir,
        class_ids=CLASS_IDS,
        class_colors=CLASS_COLORS_RGB,
        rgb_tol=args.rgb_tol,
        min_area_px=args.min_area,
        epsilon_px=args.eps
    )


if __name__ == "__main__":
    main()