import json
import os
import cv2
from utils import carregar_labelmap_completo, converter_mask_rgb_para_ids

# ⚙️ Configurações
with open("config.json", "r") as f:
    config = json.load(f)
MODELO = config["camera"]
RESOLUCAO = config["resolucao"]
pasta_base = os.path.join(MODELO, "dataset")
fonte_dados = ["original", "augmented"]
labelmap_path = os.path.join(pasta_base, "labelmap.txt")

RESOLUCOES = {
    f"{RESOLUCAO[0]}x{RESOLUCAO[1]}": (RESOLUCAO[0], RESOLUCAO[1]),
}

# === Início do processamento ===
cor_para_id, _, _, ignore_rgb = carregar_labelmap_completo(labelmap_path)
ignore_id = ignore_rgb[0]

# Cria pastas de saída
for nome_res, dim in RESOLUCOES.items():
    os.makedirs(os.path.join(pasta_base, nome_res, "images"), exist_ok=True)
    os.makedirs(os.path.join(pasta_base, nome_res, "masks"), exist_ok=True)

# Processa cada fonte de dados (original + augmented)
for fonte in fonte_dados:
    if (not os.path.exists(os.path.join(pasta_base, fonte))):
        continue
    pasta_rgb = os.path.join(pasta_base, fonte, "images")
    pasta_masks = os.path.join(pasta_base, fonte, "masks")

    nomes_arquivos = sorted([f for f in os.listdir(pasta_rgb) if f.endswith(".jpg") or f.endswith(".jpeg")])
    total = len(nomes_arquivos)

    for i, nome in enumerate(nomes_arquivos, 1):
        caminho_rgb = os.path.join(pasta_rgb, nome)
        caminho_mask = os.path.join(pasta_masks, nome.replace(".jpg", ".png").replace(".jpeg", ".png"))

        img_rgb = cv2.imread(caminho_rgb)
        if img_rgb is None:
            print(f"[!] Erro ao ler imagem {nome}")
            continue

        # Tenta carregar a máscara RGB (se existir)
        if os.path.exists(caminho_mask):
            img_mask_rgb = cv2.cvtColor(cv2.imread(caminho_mask), cv2.COLOR_BGR2RGB)
            #img_mask_rgb = cv2.cvtColor(img_mask_rgb, cv2.COLOR_BGR2RGB)  # ← CORRIGE isso!
            if img_mask_rgb is not None:
                mask_ids = converter_mask_rgb_para_ids(img_mask_rgb, cor_para_id, ignore_id)
            else:
                print(f"[!] Erro ao ler máscara {caminho_mask}, ignorando.")
                mask_ids = None
        else:
            mask_ids = None

        for nome_res, dim in RESOLUCOES.items():
            # Cria nomes únicos baseados na fonte
            nome_saida_img = f"{fonte}_{nome}"
            nome_saida_mask = nome_saida_img.replace(".jpg", ".png").replace(".jpeg", ".png")

            # Redimensiona e salva imagem
            img_resized = cv2.resize(img_rgb, dim, interpolation=cv2.INTER_AREA)
            path_img_saida = os.path.join(pasta_base, nome_res, "images", nome_saida_img)
            cv2.imwrite(path_img_saida, img_resized)

            # Redimensiona e salva máscara (se existir)
            if mask_ids is not None:
                mask_resized = cv2.resize(mask_ids, dim, interpolation=cv2.INTER_NEAREST)
                path_mask_saida = os.path.join(pasta_base, nome_res, "masks", nome_saida_mask)
                cv2.imwrite(path_mask_saida, mask_resized)

        print(f"[{fonte}] [{i}/{total}] Redimensionado: {nome}")

print("\n✅ Concluído com sucesso! Todas as fontes foram processadas.")