import os
import time
import json
import argparse
from datetime import datetime

import numpy as np
import cv2

from raw_segformer_service import make_bgr_preview_from_raw
from gal5000.gal_service import Gal5000Camera  # ajuste o nome do módulo se estiver diferente


# =========================
# Helpers gerais
# =========================

def clamp(v, lo, hi):
    return lo if v < lo else hi if v > hi else v


def ts_name() -> str:
    """Timestamp legível e único para nome de arquivo."""
    return datetime.now().strftime("%Y%m%d_%H%M%S_%f")[:-3]


def norm8(x: np.ndarray, p_lo=2, p_hi=98) -> np.ndarray:
    """
    Normaliza um canal (float32 0..1 ou uint8) em 0..255 com cortes por percentil.
    Pensado pra deixar o preview bonitinho sem estourar tudo.
    """
    x = np.asarray(x)
    if x.dtype != np.float32 and x.dtype != np.float64:
        x = x.astype(np.float32)

    # Se o canal já está em 0..1, escala pra 0..255 antes de cortar
    if x.max() <= 1.5:
        x = x * 255.0

    lo = np.percentile(x, p_lo)
    hi = np.percentile(x, p_hi)

    if hi <= lo + 1e-3:
        y = x
    else:
        y = (x - lo) * (255.0 / (hi - lo))

    return np.clip(y, 0, 255).astype(np.uint8)


def overlay_hud(
    img_bgr: np.ndarray,
    lines: list[str],
    base_h: int = 720,
    base_font_scale: float = 0.75,
    base_line_step: int = 28,
):
    """
    Escreve textos empilhados no canto superior esquerdo,
    ajustando o tamanho do texto de acordo com a altura da imagem.

    base_h: altura de referência (ex: 720 ou a RAW_H original).
    """
    h, w = img_bgr.shape[:2]

    # Fator de escala com base na altura atual
    scale = h / float(base_h)

    # Evita ficar microscópico em resoluções muito baixas
    scale = max(scale, 0.4)

    font_scale = base_font_scale * scale
    line_step  = int(base_line_step * scale)

    # Espessuras proporcionais
    thick_outline = max(1, int(3 * scale))
    thick_text    = max(1, int(2 * scale))

    # Margem superior / esquerda também escaladas
    y = int(24 * scale)
    x = int(12 * scale)

    for s in lines:
        # contorno preto
        cv2.putText(img_bgr, s, (x, y), cv2.FONT_HERSHEY_SIMPLEX, font_scale, (0, 0, 0), thick_outline, cv2.LINE_AA)
        # texto branco
        cv2.putText(img_bgr, s, (x, y), cv2.FONT_HERSHEY_SIMPLEX, font_scale, (255, 255, 255), thick_text, cv2.LINE_AA)
        y += line_step


def save_sample_raw4(
    base_dir: str,
    raw4: np.ndarray,
    preview_bgr: np.ndarray,
    meta: dict,
):
    """
    Salva:
      - RAW4 como .npy float32 (4,H,W)
      - preview RGB como .png
      - metadados como .json
    dentro de base_dir.
    """
    os.makedirs(base_dir, exist_ok=True)
    name = ts_name()

    raw_path = os.path.join(base_dir, f"{name}.raw")
    png_path = os.path.join(base_dir, f"{name}.png")
    json_path = os.path.join(base_dir, f"{name}.json")

    # RAW4
    #np.save(raw_path, raw4.astype(np.float32))
    raw4.astype(np.float32).tofile(raw_path)

    # Preview
    cv2.imwrite(png_path, preview_bgr)

    # Metadados
    with open(json_path, "w", encoding="utf-8") as f:
        json.dump(meta, f, ensure_ascii=False, indent=2)

    return raw_path, png_path, json_path


# =========================
# MAIN
# =========================

def main():
    parser = argparse.ArgumentParser(
        description="Captura de dataset RAW4 (SegFormer B0) usando Gal5000 + AutoExposure.",
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    )

    parser.add_argument("--cana", required=True, choices=["baixa", "media", "alta"], help="Estado da cana no momento da coleta.")
    parser.add_argument("--horario", required=True, choices=["cedo", "meio_dia", "entardecer", "nublado"], help="Janela de iluminação / horário da coleta.")
    parser.add_argument("--out_root", default="dataset", help="Pasta raiz do dataset.")
    parser.add_argument("--dll_dir", default=r"C:\ZendionInc\agrobot_base\Python\gal5000\dlls", help="Pasta onde está a VT_SDK64.dll (usada pelo Gal5000Camera).")
    parser.add_argument("--dll_name", default="VT_SDK64.dll", help="Nome da DLL da câmera.")
    parser.add_argument("--interval", type=float, default=1.0, help="Intervalo em segundos para auto-save quando ligado.")
    parser.add_argument("--no_ae", action="store_true", help="Desliga o AutoExposure do service (por padrão ele vem ligado).")
    parser.add_argument("--upscale", type=int, default=2, help="Fator de upscale visual do preview.")

    args = parser.parse_args()

    # ===== config =====
    with open("config.json", "r", encoding="utf-8") as f:
        config = json.load(f)
    MODELO = config["camera"]
    RAW_W = config["raw_size"][0]
    RAW_H = config["raw_size"][1]
    
    # Define diretório de sessão:
    # dataset/cana_<estado>/<horario>/<YYYYMMDD>/
    session_dir = os.path.join(MODELO, args.out_root, "brutas", f"cana_{args.cana}", args.horario, datetime.now().strftime("%Y%m%d"))
    os.makedirs(session_dir, exist_ok=True)

    print("============================================")
    print("Coleta de dataset RAW4 - SegFormer B0")
    print(f"Cana    : {args.cana}")
    print(f"Horário : {args.horario}")
    print(f"Saída   : {session_dir}")
    print("============================================")

    window_name = "Dataset Capture - RAW4 (C/SPACE=save | A=auto-save | E=AE | Q=quit)"
    cv2.namedWindow(window_name, cv2.WINDOW_NORMAL)

    auto_save = False
    last_auto_t = 0.0
    upscale = args.upscale

    # Estatísticas simples de FPS
    t_fps = time.time()
    frames = 0
    fps = 0.0

    last_msg = ""
    last_msg_t = 0.0

    try:
        cam = Gal5000Camera(dll_dir=args.dll_dir, dll_name=args.dll_name, raw_w=RAW_W, raw_h=RAW_H, use_auto_exposure=(not args.no_ae))
        with cam:
            print("[CAM] Status inicial:", cam.get_status())

            # opcional: você pode ligar streaming se quiser, mas grab_raw4 já usa single-frame
            cam.configure_fps(20)
            cam.start_streaming()

            while True:
                t0 = time.time()
                raw4_base, dbg = cam.grab_raw4(out_h=RAW_H, out_w=RAW_W, timeout_ms=2000, do_ae=True)
                t1 = time.time()

                ae_dbg = dbg.get("ae", {}) or {}
                exp_raw = dbg.get("exp_raw", None)
                gain_a = dbg.get("gain_a", None)
                gain_d = dbg.get("gain_d", None)

                bgr = make_bgr_preview_from_raw(raw4_base, rgirb=True, preview_fast=upscale > 0, preview_scale=upscale)

                # FPS
                frames += 1
                dt_fps = time.time() - t_fps
                if dt_fps >= 1.0:
                    fps = frames / dt_fps
                    frames = 0
                    t_fps = time.time()

                # AE info
                p95_disp = ae_dbg.get("p95_ema", ae_dbg.get("p95", 0.0))
                sat_disp = ae_dbg.get("sat", 0.0)
                hold = ae_dbg.get("hold", False)

                ae_on = cam.is_auto_exposure_enabled()

                # HUD principal
                lines = [
                    f"CANA: {args.cana} | HORA: {args.horario} | Pasta: {os.path.basename(session_dir)}",
                    f"AE: {'ON' if ae_on else 'OFF'} | AutoSave: {'ON' if auto_save else 'OFF'} | Intervalo: {args.interval:.1f}s",
                    f"exp_raw={exp_raw}  gain_a={gain_a}  gain_d={gain_d}  | FPS={fps:.1f}",
                    f"AEdbg: p95={p95_disp:.1f}  sat={sat_disp:.3f}  hold={hold}",
                    "Keys: C/SPACE=save | A=auto-save | E=AE toggle | M=preview scale | Q/Esc=quit",
                ]
                overlay_hud(bgr, lines, base_h=RAW_H)

                # Mensagem rápida (ex: arquivo salvo)
                if last_msg and (time.time() - last_msg_t) < 2.0:
                    cv2.putText(bgr, last_msg, (12, bgr.shape[0] - 18), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 0), 2, cv2.LINE_AA)

                cv2.imshow(window_name, bgr)

                # Auto-save
                now = time.time()
                if auto_save and (now - last_auto_t) >= args.interval:
                    meta = {
                        "ts": datetime.now().isoformat(timespec="milliseconds"),
                        "cana": args.cana,
                        "horario": args.horario,
                        "raw4_shape": list(raw4_base.shape),
                        "out_h": RAW_H,
                        "out_w": RAW_W,
                        "ae_enabled": bool(ae_on),
                        "exp_raw": int(exp_raw) if exp_raw is not None else None,
                        "gain_a": int(gain_a) if gain_a is not None else None,
                        "gain_d": int(gain_d) if gain_d is not None else None,
                        "ae_dbg": {
                            k: (float(v) if isinstance(v, (int, float, np.floating)) else v)
                            for k, v in ae_dbg.items()
                        },
                        "note": "autosave",
                    }
                    rgb_clean_bgr_save = make_bgr_preview_from_raw(raw4_base, rgirb=True, preview_fast=False)
                    raw_path, _, _ = save_sample_raw4(session_dir, raw4_base, rgb_clean_bgr_save, meta)
                    last_msg = f"SALVO (auto): {os.path.basename(raw_path)}"
                    last_msg_t = now
                    last_auto_t = now

                # Teclado
                k = cv2.waitKey(1) & 0xFF
                if k in (ord("q"), ord("Q"), 27):  # Q ou ESC
                    break

                elif k in (ord("a"), ord("A")):
                    auto_save = not auto_save
                    last_msg = f"AutoSave -> {'ON' if auto_save else 'OFF'}"
                    last_msg_t = time.time()

                elif k in (ord("e"), ord("E")):
                    cam.enable_auto_exposure(not ae_on)
                    last_msg = f"AE -> {'ON' if cam.is_auto_exposure_enabled() else 'OFF'}"
                    last_msg_t = time.time()

                elif k in (ord("m"), ord("M")):
                    upscale = 0 if upscale else args.upscale
                    last_msg = f"Preview UPSCALE -> {upscale}"
                    last_msg_t = time.time()

                elif k in (ord("c"), ord("C"), 32):  # C ou SPACE
                    meta = {
                        "ts": datetime.now().isoformat(timespec="milliseconds"),
                        "cana": args.cana,
                        "horario": args.horario,
                        "raw4_shape": list(raw4_base.shape),
                        "out_h": RAW_H,
                        "out_w": RAW_W,
                        "ae_enabled": bool(ae_on),
                        "exp_raw": int(exp_raw) if exp_raw is not None else None,
                        "gain_a": int(gain_a) if gain_a is not None else None,
                        "gain_d": int(gain_d) if gain_d is not None else None,
                        "ae_dbg": {
                            k2: (float(v2) if isinstance(v2, (int, float, np.floating)) else v2)
                            for k2, v2 in ae_dbg.items()
                        },
                        "note": "manual",
                    }
                    rgb_clean_bgr_save = make_bgr_preview_from_raw(raw4_base, rgirb=True, preview_fast=False)
                    raw_path, _, _ = save_sample_raw4(session_dir, raw4_base, rgb_clean_bgr_save, meta)
                    last_msg = f"SALVO (manual): {os.path.basename(raw_path)}"
                    last_msg_t = time.time()

                # você pode adicionar mais atalhos depois (ex: mudar intervalo, etc.)

                # Só pra não ficar rodando a 1000 FPS na UI
                # mas sem travar muito a captura
                dt_loop = time.time() - t0
                if dt_loop < 0.001:
                    time.sleep(0.001)

    finally:
        cv2.destroyAllWindows()
        print("Fim da captura.")


if __name__ == "__main__":
    main()