import argparse
import json
import re
from pathlib import Path

import cv2
import numpy as np


IMG_EXTS = {".png", ".jpg", ".jpeg", ".bmp", ".webp"}


# ============================================================
# RAW10 unpack
# ============================================================

def unpack_raw10_packed(raw: bytes, width: int, height: int) -> np.ndarray:
    """
    Desempacota RAW10 packed padrão:
    a cada 5 bytes = 4 pixels de 10 bits.

    Layout comum:
      b0 = p0[9:2]
      b1 = p1[9:2]
      b2 = p2[9:2]
      b3 = p3[9:2]
      b4 = p0[1:0] | p1[1:0]<<2 | p2[1:0]<<4 | p3[1:0]<<6

    Retorna uint16 HxW com valores 0..1023.
    """
    arr = np.frombuffer(raw, dtype=np.uint8)

    expected_groups = (width * height) // 4
    expected_bytes = expected_groups * 5

    if arr.size < expected_bytes:
        raise RuntimeError(
            f"RAW10 menor que esperado. bytes={arr.size}, esperado={expected_bytes}, "
            f"width={width}, height={height}"
        )

    arr = arr[:expected_bytes]
    groups = arr.reshape(-1, 5).astype(np.uint16)

    p0 = (groups[:, 0] << 2) | ((groups[:, 4] >> 0) & 0x03)
    p1 = (groups[:, 1] << 2) | ((groups[:, 4] >> 2) & 0x03)
    p2 = (groups[:, 2] << 2) | ((groups[:, 4] >> 4) & 0x03)
    p3 = (groups[:, 3] << 2) | ((groups[:, 4] >> 6) & 0x03)

    out = np.empty(groups.shape[0] * 4, dtype=np.uint16)
    out[0::4] = p0
    out[1::4] = p1
    out[2::4] = p2
    out[3::4] = p3

    out = out.reshape(height, width)
    return out


def normalize_to_u8(img: np.ndarray, p_low: float = 1.0, p_high: float = 99.0) -> np.ndarray:
    arr = img.astype(np.float32)
    valid = np.isfinite(arr)

    if np.count_nonzero(valid) < 20:
        return np.zeros(arr.shape[:2], dtype=np.uint8)

    vals = arr[valid]
    lo = np.percentile(vals, p_low)
    hi = np.percentile(vals, p_high)

    out = (arr - lo) / (hi - lo + 1e-6)
    out = np.clip(out, 0.0, 1.0)
    return (out * 255).astype(np.uint8)


def raw10_bin_to_gray_u8(path: Path, width: int, height: int, clahe: bool = True) -> tuple[np.ndarray, np.ndarray]:
    raw = path.read_bytes()
    mono10 = unpack_raw10_packed(raw, width=width, height=height)

    gray_u8 = normalize_to_u8(mono10)

    if clahe:
        eq = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
        gray_u8 = eq.apply(gray_u8)

    bgr = cv2.cvtColor(gray_u8, cv2.COLOR_GRAY2BGR)
    return gray_u8, bgr


# ============================================================
# Metadata
# ============================================================

def find_meta_for_bin(bin_path: Path) -> Path | None:
    """
    Procura meta.json na pasta do .bin ou nas pastas acima próximas.
    """
    candidates = [
        bin_path.parent / "meta.json",
        bin_path.parent / "metadata.json",
        bin_path.parent.parent / "meta.json",
        bin_path.parent.parent / "metadata.json",
    ]

    for c in candidates:
        if c.exists():
            return c

    return None


def extract_camera_info_from_meta(meta: dict, cam_key: str):
    """
    Tenta extrair width/height/raw_format para CAM_B ou CAM_C em vários formatos de meta.
    """
    # Caso meta["camera_info"]["CAM_B"]
    for root_key in ["camera_info", "cameras", "camera_meta", "payload_sources_info"]:
        root = meta.get(root_key)
        if isinstance(root, dict):
            info = root.get(cam_key)
            if isinstance(info, dict):
                return info

    # Caso meta tenha lista de câmeras
    for root_key in ["camera_info", "cameras", "sources"]:
        root = meta.get(root_key)
        if isinstance(root, list):
            for item in root:
                if not isinstance(item, dict):
                    continue
                name = item.get("camera") or item.get("name") or item.get("id") or item.get("socket")
                if name == cam_key:
                    return item

    # Fallback: procura recursivamente dicionário que tenha CAM_B/C
    stack = [meta]
    while stack:
        obj = stack.pop()
        if isinstance(obj, dict):
            if cam_key in obj and isinstance(obj[cam_key], dict):
                return obj[cam_key]
            for v in obj.values():
                if isinstance(v, (dict, list)):
                    stack.append(v)
        elif isinstance(obj, list):
            for v in obj:
                if isinstance(v, (dict, list)):
                    stack.append(v)

    return None


def parse_width_height_from_meta(meta_path: Path, cam_key: str):
    try:
        meta = json.loads(meta_path.read_text(encoding="utf-8"))
    except Exception:
        meta = json.loads(meta_path.read_text(encoding="latin-1"))

    info = extract_camera_info_from_meta(meta, cam_key)

    if not info:
        return None, None

    width = (
        info.get("width")
        or info.get("w")
        or info.get("sensor_width")
        or info.get("frame_width")
    )

    height = (
        info.get("height")
        or info.get("h")
        or info.get("sensor_height")
        or info.get("frame_height")
    )

    if width is None or height is None:
        return None, None

    return int(width), int(height)


def resolve_width_height(bin_path: Path, cam_key: str, args):
    if args.width > 0 and args.height > 0:
        return args.width, args.height

    meta_path = find_meta_for_bin(bin_path)
    if meta_path:
        w, h = parse_width_height_from_meta(meta_path, cam_key)
        if w and h:
            return w, h

    raise RuntimeError(
        f"Não consegui descobrir width/height para {bin_path}. "
        f"Passe manualmente: --width 1280 --height 800"
    )


# ============================================================
# Pairing CAM_B/C
# ============================================================

def clean_stem_for_pair(path: Path, cam_key: str):
    """
    Remove CAM_B/C do nome para tentar parear frames.
    """
    stem = path.stem
    s = stem

    patterns = [
        cam_key,
        cam_key.lower(),
        cam_key.replace("_", ""),
        cam_key.replace("_", "").lower(),
    ]

    for p in patterns:
        s = s.replace(p, "")

    s = re.sub(r"[_\-\s]+", "_", s).strip("_").lower()
    return s


def find_cam_bins(root_dir: Path, cam_key: str):
    bins = []
    for p in root_dir.rglob("*.bin"):
        name = p.name.lower()
        if cam_key.lower() in name:
            bins.append(p)
    return sorted(bins)


def pair_cam_bins(root_dir: Path, left_cam: str, right_cam: str):
    left_bins = find_cam_bins(root_dir, left_cam)
    right_bins = find_cam_bins(root_dir, right_cam)

    right_by_folder_and_key = {}
    right_by_key = {}

    for rp in right_bins:
        key = clean_stem_for_pair(rp, right_cam)
        right_by_folder_and_key[(rp.parent, key)] = rp
        right_by_key.setdefault(key, rp)

    pairs = []

    for lp in left_bins:
        key = clean_stem_for_pair(lp, left_cam)

        rp = right_by_folder_and_key.get((lp.parent, key))
        if rp is None:
            rp = right_by_key.get(key)

        # Fallback comum: CAM_B/C dentro da mesma pasta, mas nomes não batem
        if rp is None:
            candidates_same_folder = [r for r in right_bins if r.parent == lp.parent]
            if len(candidates_same_folder) == 1:
                rp = candidates_same_folder[0]

        if rp is not None:
            pairs.append((lp, rp))

    if pairs:
        return pairs, left_bins, right_bins, "cam-key"

    # Fallback por ordem
    n = min(len(left_bins), len(right_bins))
    pairs = list(zip(left_bins[:n], right_bins[:n]))
    return pairs, left_bins, right_bins, "order"


# ============================================================
# Stereo SGBM
# ============================================================

def ensure_same_size(left, right):
    h = min(left.shape[0], right.shape[0])
    w = min(left.shape[1], right.shape[1])

    left2 = cv2.resize(left, (w, h), interpolation=cv2.INTER_AREA)
    right2 = cv2.resize(right, (w, h), interpolation=cv2.INTER_AREA)

    return left2, right2


def make_sgbm(num_disp, block_size, min_disp=0, uniqueness=8, speckle_window=80, speckle_range=2):
    num_disp = max(16, int(round(num_disp / 16)) * 16)

    block_size = max(3, int(block_size))
    if block_size % 2 == 0:
        block_size += 1

    matcher = cv2.StereoSGBM_create(
        minDisparity=min_disp,
        numDisparities=num_disp,
        blockSize=block_size,
        P1=8 * block_size * block_size,
        P2=32 * block_size * block_size,
        disp12MaxDiff=1,
        uniquenessRatio=uniqueness,
        speckleWindowSize=speckle_window,
        speckleRange=speckle_range,
        preFilterCap=63,
        mode=cv2.STEREO_SGBM_MODE_SGBM_3WAY,
    )

    return matcher, num_disp, block_size


def compute_disparity(left_gray, right_gray, args):
    matcher, num_disp, block_size = make_sgbm(
        num_disp=args.num_disp,
        block_size=args.block_size,
        min_disp=args.min_disp,
        uniqueness=args.uniqueness,
        speckle_window=args.speckle_window,
        speckle_range=args.speckle_range,
    )

    disp_raw = matcher.compute(left_gray, right_gray).astype(np.float32) / 16.0

    valid = disp_raw > args.min_valid_disp

    disp_vis = disp_raw.copy()
    disp_vis[~valid] = 0.0

    if np.count_nonzero(valid) > 20:
        vals = disp_vis[valid]
        p2 = np.percentile(vals, 2)
        p98 = np.percentile(vals, 98)
        disp_norm = (disp_vis - p2) / (p98 - p2 + 1e-6)
        disp_norm = np.clip(disp_norm, 0.0, 1.0)
    else:
        disp_norm = np.zeros_like(disp_vis, dtype=np.float32)

    disp_color = cv2.applyColorMap((disp_norm * 255).astype(np.uint8), cv2.COLORMAP_TURBO)

    valid_mask = np.zeros_like(disp_color)
    valid_mask[valid] = (255, 255, 255)

    stats = {
        "num_disp": num_disp,
        "block_size": block_size,
        "valid_pct": float(np.mean(valid) * 100.0),
        "disp_p05": float(np.percentile(disp_vis[valid], 5)) if np.count_nonzero(valid) > 20 else 0.0,
        "disp_p50": float(np.percentile(disp_vis[valid], 50)) if np.count_nonzero(valid) > 20 else 0.0,
        "disp_p95": float(np.percentile(disp_vis[valid], 95)) if np.count_nonzero(valid) > 20 else 0.0,
    }

    return disp_raw, disp_color, valid_mask, valid, stats


# ============================================================
# View
# ============================================================

def resize_to_height(img, target_h):
    h, w = img.shape[:2]
    if h == target_h:
        return img

    scale = target_h / h
    new_w = max(1, int(w * scale))
    return cv2.resize(img, (new_w, target_h), interpolation=cv2.INTER_AREA)


def draw_header(canvas, lines):
    header_h = 24 + 24 * len(lines)
    cv2.rectangle(canvas, (0, 0), (canvas.shape[1], header_h), (0, 0, 0), -1)

    y = 24
    for line in lines:
        cv2.putText(
            canvas,
            line,
            (12, y),
            cv2.FONT_HERSHEY_SIMPLEX,
            0.55,
            (255, 255, 255),
            1,
            cv2.LINE_AA,
        )
        y += 24

    return canvas


def make_overlay(left_bgr, disp_color, alpha=0.45):
    disp_resized = cv2.resize(
        disp_color,
        (left_bgr.shape[1], left_bgr.shape[0]),
        interpolation=cv2.INTER_AREA,
    )
    return cv2.addWeighted(left_bgr, 1.0 - alpha, disp_resized, alpha, 0)


def compose_view(left_bgr, right_bgr, disp_color, valid_mask, pair, idx, total, mode, stats, args):
    if mode == "disp":
        third = disp_color
        mode_name = "disparity"
    elif mode == "mask":
        third = valid_mask
        mode_name = "valid mask"
    else:
        third = make_overlay(left_bgr, disp_color)
        mode_name = "overlay"

    left = resize_to_height(left_bgr, args.view_h)
    right = resize_to_height(right_bgr, args.view_h)
    third = resize_to_height(third, args.view_h)

    h = min(left.shape[0], right.shape[0], third.shape[0])
    left = left[:h]
    right = right[:h]
    third = third[:h]

    canvas = np.hstack([left, right, third])

    lp, rp = pair

    lines = [
        f"{idx + 1}/{total} | L={lp.name} | R={rp.name}",
        f"mode={mode_name} | valid={stats['valid_pct']:.1f}% | disp p05={stats['disp_p05']:.2f} p50={stats['disp_p50']:.2f} p95={stats['disp_p95']:.2f}",
        f"numDisp={stats['num_disp']} | block={stats['block_size']} | uniqueness={args.uniqueness}",
        "N/SPACE prox | A ant | M modo | [ ] numDisp | - + block | S salvar | Q sair",
    ]

    return draw_header(canvas, lines)


# ============================================================
# Main
# ============================================================

def main():
    parser = argparse.ArgumentParser()

    parser.add_argument("--root_dir", required=True, help="Pasta raiz onde estão os .bin e meta.json")
    parser.add_argument("--left_cam", default="CAM_B", help="Câmera esquerda. Ex: CAM_B")
    parser.add_argument("--right_cam", default="CAM_C", help="Câmera direita. Ex: CAM_C")

    parser.add_argument("--width", type=int, default=-1, help="Largura manual caso não tenha meta.json")
    parser.add_argument("--height", type=int, default=-1, help="Altura manual caso não tenha meta.json")

    parser.add_argument("--view_h", type=int, default=480)

    parser.add_argument("--num_disp", type=int, default=128)
    parser.add_argument("--block_size", type=int, default=7)
    parser.add_argument("--min_disp", type=int, default=0)
    parser.add_argument("--min_valid_disp", type=float, default=1.0)

    parser.add_argument("--uniqueness", type=int, default=8)
    parser.add_argument("--speckle_window", type=int, default=80)
    parser.add_argument("--speckle_range", type=int, default=2)

    parser.add_argument("--no_clahe", action="store_true")
    parser.add_argument("--swap", action="store_true", help="Inverte left/right depois do pareamento")
    parser.add_argument("--save_dir", default="stereo_raw10_sgbm_saves")

    args = parser.parse_args()

    root_dir = Path(args.root_dir)
    save_dir = Path(args.save_dir)
    save_dir.mkdir(parents=True, exist_ok=True)

    pairs, left_bins, right_bins, pair_mode = pair_cam_bins(
        root_dir=root_dir,
        left_cam=args.left_cam,
        right_cam=args.right_cam,
    )

    if args.swap:
        pairs = [(r, l) for l, r in pairs]
        args.left_cam, args.right_cam = args.right_cam, args.left_cam

    if not left_bins:
        raise RuntimeError(f"Nenhum .bin encontrado com {args.left_cam} em {root_dir}")

    if not right_bins:
        raise RuntimeError(f"Nenhum .bin encontrado com {args.right_cam} em {root_dir}")

    if not pairs:
        raise RuntimeError("Não consegui parear CAM_B/C.")

    print(f"[INFO] root_dir: {root_dir}")
    print(f"[INFO] left_cam: {args.left_cam} | arquivos: {len(left_bins)}")
    print(f"[INFO] right_cam: {args.right_cam} | arquivos: {len(right_bins)}")
    print(f"[INFO] pares: {len(pairs)}")
    print(f"[INFO] pareamento: {pair_mode}")
    print("[INFO] controles:")
    print("       N ou SPACE = próxima")
    print("       A          = anterior")
    print("       M          = modo disparity/mask/overlay")
    print("       [ / ]      = diminui/aumenta numDisparities")
    print("       - / +      = diminui/aumenta blockSize")
    print("       S          = salva visual atual")
    print("       Q ou ESC   = sair")

    idx = 0
    mode = "disp"

    cached_key = None
    cached_data = None

    cv2.namedWindow("Stereo RAW10 SGBM Viewer", cv2.WINDOW_NORMAL)

    while True:
        pair = pairs[idx]
        lp, rp = pair

        key_cache = (
            str(lp),
            str(rp),
            args.num_disp,
            args.block_size,
            args.uniqueness,
            args.speckle_window,
            args.speckle_range,
            args.min_disp,
            args.min_valid_disp,
            args.no_clahe,
            args.width,
            args.height,
        )

        if key_cache != cached_key:
            print(f"[RUN] {idx + 1}/{len(pairs)} - L={lp.name} | R={rp.name}")

            try:
                lw, lh = resolve_width_height(lp, args.left_cam, args)
                rw, rh = resolve_width_height(rp, args.right_cam, args)

                left_gray, left_bgr = raw10_bin_to_gray_u8(
                    lp,
                    width=lw,
                    height=lh,
                    clahe=not args.no_clahe,
                )

                right_gray, right_bgr = raw10_bin_to_gray_u8(
                    rp,
                    width=rw,
                    height=rh,
                    clahe=not args.no_clahe,
                )

                left_gray, right_gray = ensure_same_size(left_gray, right_gray)
                left_bgr, right_bgr = ensure_same_size(left_bgr, right_bgr)

                _, disp_color, valid_mask, valid, stats = compute_disparity(left_gray, right_gray, args)

                cached_data = (left_bgr, right_bgr, disp_color, valid_mask, stats)
                cached_key = key_cache

            except Exception as e:
                print(f"[ERRO] Falha processando par:")
                print(f"       L={lp}")
                print(f"       R={rp}")
                print(f"       erro={e}")

                idx = min(idx + 1, len(pairs) - 1)
                cached_key = None
                cached_data = None
                continue

        else:
            left_bgr, right_bgr, disp_color, valid_mask, stats = cached_data

        view = compose_view(
            left_bgr=left_bgr,
            right_bgr=right_bgr,
            disp_color=disp_color,
            valid_mask=valid_mask,
            pair=pair,
            idx=idx,
            total=len(pairs),
            mode=mode,
            stats=stats,
            args=args,
        )

        cv2.imshow("Stereo RAW10 SGBM Viewer", view)
        key = cv2.waitKey(0) & 0xFF

        if key in [27, ord("q"), ord("Q")]:
            break

        elif key in [ord("n"), ord("N"), 32]:
            idx = min(idx + 1, len(pairs) - 1)
            cached_key = None

        elif key in [ord("a"), ord("A")]:
            idx = max(idx - 1, 0)
            cached_key = None

        elif key in [ord("m"), ord("M")]:
            if mode == "disp":
                mode = "mask"
            elif mode == "mask":
                mode = "overlay"
            else:
                mode = "disp"

        elif key == ord("["):
            args.num_disp = max(16, args.num_disp - 16)
            cached_key = None
            print(f"[PARAM] num_disp={args.num_disp}")

        elif key == ord("]"):
            args.num_disp = min(512, args.num_disp + 16)
            cached_key = None
            print(f"[PARAM] num_disp={args.num_disp}")

        elif key in [ord("-"), ord("_")]:
            args.block_size = max(3, args.block_size - 2)
            if args.block_size % 2 == 0:
                args.block_size -= 1
            cached_key = None
            print(f"[PARAM] block_size={args.block_size}")

        elif key in [ord("+"), ord("=")]:
            args.block_size = min(31, args.block_size + 2)
            if args.block_size % 2 == 0:
                args.block_size += 1
            cached_key = None
            print(f"[PARAM] block_size={args.block_size}")

        elif key in [ord("s"), ord("S")]:
            out_path = save_dir / f"stereo_raw10_{idx:04d}_{mode}.png"
            cv2.imwrite(str(out_path), view)
            print(f"[SAVE] {out_path}")

    cv2.destroyAllWindows()


if __name__ == "__main__":
    main()