import os
import json
import argparse
from pathlib import Path

import cv2
import numpy as np

from cam_3.pi.raw_processor_core import RawProcessorCore
from cam_3.pi.raw_processor_preview import RawProcessorPreview


def load_json(path: Path) -> dict:
    with open(path, "r", encoding="utf-8") as f:
        return json.load(f)


def normalize_float01_to_bgr(img_float: np.ndarray) -> np.ndarray:
    """
    Recebe RGB float32 [0..1] em HWC e devolve BGR uint8.
    """
    rgb_u8 = np.clip(img_float * 255.0, 0, 255).astype(np.uint8)
    return cv2.cvtColor(rgb_u8, cv2.COLOR_RGB2BGR)


def chw_to_hwc(arr: np.ndarray) -> np.ndarray:
    if arr.ndim != 3:
        raise ValueError(f"Esperado CHW 3D, recebido shape={arr.shape}")
    return np.transpose(arr, (1, 2, 0))


def build_visual_from_saved_payload(payload_path: Path, meta: dict, cam_id: str | None = None) -> tuple[np.ndarray, str]:
    """
    Retorna:
      preview_bgr_reconstructed
      texto_descritivo
    """
    saved_type = meta.get("saved_payload_type")

    # =========================================================
    # Caso MULTI payload por câmera
    # =========================================================
    if saved_type == "raw_native_multi":
        if cam_id is None:
            raise RuntimeError("cam_id é obrigatório para saved_payload_type='raw_native_multi'")

        saved_dtypes = meta.get("saved_payload_dtypes", {}) or {}
        saved_shapes = meta.get("saved_payload_shapes", {}) or {}

        saved_dtype = saved_dtypes.get(cam_id)
        saved_shape = saved_shapes.get(cam_id)

        if saved_dtype is None or saved_shape is None:
            raise RuntimeError(
                f"JSON não contém saved_payload_dtypes/saved_payload_shapes para {cam_id}"
            )

        np_dtype = np.dtype(saved_dtype)
        raw = np.fromfile(str(payload_path), dtype=np_dtype)
        arr = raw.reshape(tuple(saved_shape))

        # Busca metadados da câmera no stream_meta
        stream_meta = meta.get("stream_meta", {}) or {}
        cam_frames = stream_meta.get("camera_frames", {}) or {}
        cam_meta = cam_frames.get(cam_id, {}) or {}

        role = cam_meta.get("role", cam_id)
        interface = cam_meta.get("interface", "")
        bit_depth = int(cam_meta.get("bit_depth", 10))
        bayer = cam_meta.get("bayer_pattern", meta.get("bayer_pattern", "GBRG"))

        # USB RGB nativo
        if interface.upper() == "USB" or (arr.ndim == 3 and arr.shape[2] == 3 and arr.dtype == np.uint8):
            preview_bgr = arr.copy()
            desc = f"{cam_id} | role={role} | USB/RGB nativo | dtype={arr.dtype} | shape={arr.shape}"
            return preview_bgr, desc

        # CSI RAW packed mono
        sensor_width, sensor_height = resolve_sensor_dims_for_raw10_packed(arr, cam_meta, meta)

        core = RawProcessorCore(
            sensor_width=sensor_width,
            sensor_height=sensor_height,
            bayer_pattern=bayer,
        )
        preview = RawProcessorPreview(
            sensor_width=sensor_width,
            sensor_height=sensor_height,
            bayer_pattern=bayer,
        )

        packed = arr
        if packed.ndim == 3 and packed.shape[2] == 1:
            packed = packed[:, :, 0]

        raw16 = core.unpack_raw10_packed(packed)
        preview_bgr = preview.raw16_to_preview_bgr(raw16, bit_depth=bit_depth)

        desc = (
            f"{cam_id} | role={role} | RAW packed mono | "
            f"dtype={arr.dtype} | shape={arr.shape} | "
            f"sensor={sensor_width}x{sensor_height} | "
            f"bayer={bayer} | bit_depth={bit_depth}"
        )
        return preview_bgr, desc

    # =========================================================
    # Caso payload único
    # =========================================================
    saved_dtype = meta.get("saved_payload_dtype")
    saved_shape = meta.get("saved_payload_shape")

    if saved_type is None or saved_dtype is None or saved_shape is None:
        raise RuntimeError(
            "JSON não contém saved_payload_type / saved_payload_dtype / saved_payload_shape"
        )

    np_dtype = np.dtype(saved_dtype)
    raw = np.fromfile(str(payload_path), dtype=np_dtype)
    arr = raw.reshape(tuple(saved_shape))

    if saved_type == "rgb":
        if arr.ndim != 3 or arr.shape[0] != 3:
            raise RuntimeError(f"Payload RGB inválido, shape={arr.shape}")

        rgb_hwc = chw_to_hwc(arr.astype(np.float32))
        preview_bgr = normalize_float01_to_bgr(rgb_hwc)
        desc = f"Reconstruido de RGB salvo | dtype={arr.dtype} | shape={arr.shape}"
        return preview_bgr, desc

    if saved_type == "multispec":
        if arr.ndim != 3 or arr.shape[0] < 3:
            raise RuntimeError(f"Payload MULTISPEC inválido, shape={arr.shape}")

        rgb_hwc = chw_to_hwc(arr[:3].astype(np.float32))
        preview_bgr = normalize_float01_to_bgr(rgb_hwc)
        desc = f"Reconstruido de MULTISPEC salvo | dtype={arr.dtype} | shape={arr.shape}"
        return preview_bgr, desc

    if saved_type == "raw_native_single":
        if arr.ndim == 3 and arr.shape[2] == 3 and arr.dtype == np.uint8:
            preview_bgr = arr.copy()
            desc = f"Reconstruido de RAW nativo USB | dtype={arr.dtype} | shape={arr.shape}"
            return preview_bgr, desc

        stream_meta = meta.get("stream_meta", {})
        source_camera = stream_meta.get("source_camera", {}) or {}

        bayer = source_camera.get("bayer_pattern", meta.get("bayer_pattern", "GBRG"))
        bit_depth = int(source_camera.get("bit_depth", 10))

        sensor_height = int(meta.get("sensor_height"))
        sensor_width = int(meta.get("sensor_width"))

        core = RawProcessorCore(
            sensor_width=sensor_width,
            sensor_height=sensor_height,
            bayer_pattern=bayer,
        )
        preview = RawProcessorPreview(
            sensor_width=sensor_width,
            sensor_height=sensor_height,
            bayer_pattern=bayer,
        )

        packed = arr
        if packed.ndim == 3 and packed.shape[2] == 1:
            packed = packed[:, :, 0]

        raw16 = core.unpack_raw10_packed(packed)
        preview_bgr = preview.raw16_to_preview_bgr(raw16, bit_depth=bit_depth)

        desc = f"Reconstruido de RAW packed mono | dtype={arr.dtype} | shape={arr.shape} | bayer={bayer} | bit_depth={bit_depth}"
        return preview_bgr, desc

    if saved_type == "raw10_packed":
        stream_meta = meta.get("stream_meta", {}) or {}
        source_camera = stream_meta.get("source_camera", {}) or {}

        bayer = source_camera.get("bayer_pattern", meta.get("bayer_pattern", "GBRG"))
        bit_depth = int(source_camera.get("bit_depth", 10))

        sensor_height = int(meta.get("sensor_height"))
        sensor_width = int(meta.get("sensor_width"))

        core = RawProcessorCore(
            sensor_width=sensor_width,
            sensor_height=sensor_height,
            bayer_pattern=bayer,
        )
        preview = RawProcessorPreview(
            sensor_width=sensor_width,
            sensor_height=sensor_height,
            bayer_pattern=bayer,
        )

        packed = arr
        if packed.ndim == 3 and packed.shape[2] == 1:
            packed = packed[:, :, 0]

        raw16 = core.unpack_raw10_packed(packed)
        preview_bgr = preview.raw16_to_preview_bgr(raw16, bit_depth=bit_depth)

        desc = (
            f"Reconstruido de RAW10 packed | "
            f"dtype={arr.dtype} | shape={arr.shape} | "
            f"sensor={sensor_width}x{sensor_height} | "
            f"bayer={bayer} | bit_depth={bit_depth}"
        )
        return preview_bgr, desc

    raise RuntimeError(f"saved_payload_type não suportado neste script: {saved_type}")


def build_panels_from_group(group):
    panels = []

    meta = load_json(group["json"])

    preview_saved = cv2.imread(str(group["png"]), cv2.IMREAD_COLOR)
    if preview_saved is None:
        raise RuntimeError(f"Falha ao ler preview PNG: {group['png']}")
    panels.append(("Preview salvo", preview_saved, f"{preview_saved.shape[1]}x{preview_saved.shape[0]}"))

    if group["final_raw"] is not None:
        img, desc = build_visual_from_saved_payload(group["final_raw"], meta)
        panels.append(("Reconstruido (final)", img, desc))

    for cam_id, path in group["cameras"].items():
        img, desc = build_visual_from_saved_payload(path, meta, cam_id=cam_id)
        panels.append((f"{cam_id} reconstruido", img, desc))

    return panels


def compose_panels(panels, max_width=1600):
    imgs = []

    # aplica label
    for title, img, subtitle in panels:
        img_labeled = put_label(img, title, subtitle)
        imgs.append(img_labeled)

    # normaliza tamanho base
    max_h = max(img.shape[0] for img in imgs)

    resized = []
    for img in imgs:
        scale = max_h / img.shape[0]
        w = int(img.shape[1] * scale)
        resized.append(cv2.resize(img, (w, max_h), interpolation=cv2.INTER_NEAREST))

    # =========================
    # Montagem em grid 2x2
    # =========================
    rows = []
    gap = np.full((max_h, 20, 3), 30, dtype=np.uint8)

    for i in range(0, len(resized), 2):
        row_imgs = resized[i:i+2]

        # se só tiver 1 imagem na linha, duplica espaço vazio
        if len(row_imgs) == 1:
            blank = np.zeros_like(row_imgs[0])
            row_imgs.append(blank)

        row = np.hstack([row_imgs[0], gap, row_imgs[1]])
        rows.append(row)

    # junta linhas
    gap_h = np.full((20, rows[0].shape[1], 3), 30, dtype=np.uint8)

    canvas = rows[0]
    for r in rows[1:]:
        canvas = np.vstack([canvas, gap_h, r])

    # =========================
    # Resize final
    # =========================
    if canvas.shape[1] > max_width:
        scale = max_width / canvas.shape[1]
        canvas = cv2.resize(
            canvas,
            (int(canvas.shape[1] * scale), int(canvas.shape[0] * scale)),
            interpolation=cv2.INTER_AREA
        )

    return canvas


def sort_panels(panels):
    order = ["Preview salvo", "cam2", "cam0", "cam1"]

    def key(p):
        title = p[0].lower()
        for i, k in enumerate(order):
            if k in title:
                return i
        return 99

    return sorted(panels, key=key)


def fit_same_height(img_a: np.ndarray, img_b: np.ndarray, target_h: int = None):
    if target_h is None:
        target_h = max(img_a.shape[0], img_b.shape[0])

    def resize_to_h(img, h):
        scale = h / img.shape[0]
        w = int(img.shape[1] * scale)
        return cv2.resize(img, (w, h), interpolation=cv2.INTER_NEAREST)

    return resize_to_h(img_a, target_h), resize_to_h(img_b, target_h)


def put_label(img: np.ndarray, title: str, subtitle: str = "") -> np.ndarray:
    out = img.copy()
    cv2.putText(out, title, (12, 28), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 0, 0), 3, cv2.LINE_AA)
    cv2.putText(out, title, (12, 28), cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 255), 2, cv2.LINE_AA)

    if subtitle:
        cv2.putText(out, subtitle, (12, 56), cv2.FONT_HERSHEY_SIMPLEX, 0.55, (0, 0, 0), 3, cv2.LINE_AA)
        cv2.putText(out, subtitle, (12, 56), cv2.FONT_HERSHEY_SIMPLEX, 0.55, (255, 255, 255), 1, cv2.LINE_AA)

    return out


def resolve_capture_group(input_path: Path):
    """
    Resolve todos os arquivos relacionados a uma captura.

    Retorna:
    {
        "json": Path,
        "png": Path,
        "final_raw": Path | None,
        "cameras": { "cam0": Path, ... }
    }
    """
    input_path = input_path.resolve()
    folder = input_path.parent

    name = input_path.stem

    # remove sufixo _camX se existir
    if "_cam" in name:
        base_name = name.split("_cam")[0]
    else:
        base_name = name

    json_path = folder / f"{base_name}.json"
    png_path = folder / f"{base_name}.png"

    if not json_path.exists():
        raise FileNotFoundError(f"JSON não encontrado: {json_path}")
    if not png_path.exists():
        raise FileNotFoundError(f"PNG não encontrado: {png_path}")

    meta = load_json(json_path)

    group = {
        "json": json_path,
        "png": png_path,
        "final_raw": None,
        "cameras": {}
    }

    # =========================
    # Caso MULTI payload
    # =========================
    if "saved_payload_paths" in meta:
        for cam_id, fname in meta["saved_payload_paths"].items():
            path = folder / fname
            if path.exists():
                group["cameras"][cam_id] = path

    # =========================
    # Caso payload único (.raw)
    # =========================
    elif "saved_payload_path" in meta:
        path = folder / meta["saved_payload_path"]
        if path.exists():
            group["final_raw"] = path

    return group


def resolve_sensor_dims_for_raw10_packed(arr: np.ndarray, cam_meta: dict, meta: dict) -> tuple[int, int]:
    """
    Para CSI RAW10 packed:
      packed_width = ceil(sensor_width * 5 / 4)
    Na prática aqui usamos:
      sensor_width = packed_width * 4 // 5

    Altura permanece a mesma.
    """
    packed_h = int(arr.shape[0])
    packed_w = int(arr.shape[1])

    interface = str(cam_meta.get("interface", "")).upper()
    bit_depth = int(cam_meta.get("bit_depth", 10))

    # USB ou RGB HWC não entra nessa lógica
    if interface == "USB":
        return packed_w, packed_h

    # Caso esperado: CSI RAW10 packed mono
    if bit_depth == 10:
        sensor_w = (packed_w * 4) // 5
        sensor_h = packed_h
        return sensor_w, sensor_h

    # fallback conservador
    return packed_w, packed_h


def list_capture_groups_from_dir(folder: Path) -> list[Path]:
    """
    Lista todos os JSONs de captura do diretório, ordenados por nome.
    Cada JSON representa uma captura.
    """
    if not folder.exists() or not folder.is_dir():
        raise FileNotFoundError(f"Diretório não encontrado: {folder}")

    items = sorted(folder.glob("*.json"))
    if not items:
        raise RuntimeError(f"Nenhum arquivo .json encontrado em: {folder}")

    return items


def resolve_navigation_inputs(input_path: Path) -> tuple[list[Path], int]:
    """
    Retorna:
      entries: lista de JSONs de captura
      start_index: índice inicial baseado no input fornecido
    """
    input_path = input_path.resolve()

    # Caso 1: usuário passou uma pasta
    if input_path.is_dir():
        entries = list_capture_groups_from_dir(input_path)
        return entries, 0

    # Caso 2: usuário passou arquivo
    if not input_path.exists():
        raise FileNotFoundError(f"Arquivo não encontrado: {input_path}")

    folder = input_path.parent
    entries = list_capture_groups_from_dir(folder)

    # Tenta descobrir qual JSON corresponde ao input
    if input_path.suffix.lower() == ".json":
        target_json = input_path.resolve()
    else:
        group = resolve_capture_group(input_path)
        target_json = group["json"].resolve()

    try:
        idx = entries.index(target_json)
    except ValueError:
        idx = 0

    return entries, idx


def render_group_to_canvas(json_path: Path, max_width: int):
    group = resolve_capture_group(json_path)
    meta = load_json(group["json"])

    panels = build_panels_from_group(group)
    panels = sort_panels(panels)
    canvas = compose_panels(panels, max_width=max_width)

    info = {
        "json": group["json"],
        "png": group["png"],
        "final_raw": group["final_raw"],
        "cameras": group["cameras"],
        "meta": meta,
    }
    return canvas, info


def main():
    parser = argparse.ArgumentParser(
        description="Valida visualmente payload salvo (.bin/.raw/.json/.png) comparando com o preview .png"
    )
    parser.add_argument("--input_path", help="Caminho para .json, .png, .bin, .raw ou diretório")
    parser.add_argument("--max-width", type=int, default=1600, help="Largura máxima da janela final")
    args = parser.parse_args()

    input_path = Path(args.input_path)
    entries, current_idx = resolve_navigation_inputs(input_path)

    window_name = "Validacao do payload salvo | A=anterior | D=proximo | Q/Esc=sair"
    cv2.namedWindow(window_name, cv2.WINDOW_NORMAL)

    while True:
        current_json = entries[current_idx]
        canvas, info = render_group_to_canvas(current_json, max_width=args.max_width)

        # Cabeçalho adicional na imagem
        overlay = canvas.copy()
        text = f"{current_idx + 1}/{len(entries)} | {current_json.name}"
        cv2.putText(overlay, text, (12, overlay.shape[0] - 16),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.65, (0, 0, 0), 3, cv2.LINE_AA)
        cv2.putText(overlay, text, (12, overlay.shape[0] - 16),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.65, (0, 255, 0), 1, cv2.LINE_AA)

        cv2.imshow(window_name, overlay)

        meta = info["meta"]
        print("=" * 60)
        print(f"[{current_idx + 1}/{len(entries)}]")
        print("Entrada JSON   :", info["json"])
        print("PNG            :", info["png"])
        print("Final RAW      :", info["final_raw"])
        print("Câmeras        :", {k: str(v) for k, v in info["cameras"].items()})
        print("saved_payload_type :", meta.get("saved_payload_type"))
        print("saved_payload_dtype:", meta.get("saved_payload_dtype"))
        print("saved_payload_shape:", meta.get("saved_payload_shape"))
        print("saved_payload_dtypes:", meta.get("saved_payload_dtypes"))
        print("saved_payload_shapes:", meta.get("saved_payload_shapes"))
        print("stream frame_type  :", (meta.get("stream_meta") or {}).get("frame_type"))
        print("=" * 60)

        k = cv2.waitKey(0) & 0xFF

        if k in (ord("q"), ord("Q"), 27):
            break
        elif k in (ord("d"), ord("D")):
            current_idx = min(current_idx + 1, len(entries) - 1)
        elif k in (ord("a"), ord("A")):
            current_idx = max(current_idx - 1, 0)

    cv2.destroyAllWindows()


if __name__ == "__main__":
    main()