import argparse
import cv2
import depthai as dai
import numpy as np


def socket_from_name(name: str):
    name = name.upper().strip()
    mapping = {
        "CAM_A": dai.CameraBoardSocket.CAM_A,
        "CAM_B": dai.CameraBoardSocket.CAM_B,
        "CAM_C": dai.CameraBoardSocket.CAM_C,
        "LEFT": dai.CameraBoardSocket.CAM_B,
        "RIGHT": dai.CameraBoardSocket.CAM_C,
    }
    if name not in mapping:
        raise ValueError(f"Socket inválido: {name}. Use CAM_B, CAM_C, LEFT ou RIGHT.")
    return mapping[name]


def make_pipeline(args):
    pipeline = dai.Pipeline()

    left_socket = socket_from_name(args.left_socket)
    right_socket = socket_from_name(args.right_socket)

    left = pipeline.create(dai.node.MonoCamera)
    right = pipeline.create(dai.node.MonoCamera)
    stereo = pipeline.create(dai.node.StereoDepth)

    xout_left = pipeline.create(dai.node.XLinkOut)
    xout_right = pipeline.create(dai.node.XLinkOut)
    xout_disp = pipeline.create(dai.node.XLinkOut)
    xout_depth = pipeline.create(dai.node.XLinkOut)

    xout_left.setStreamName("rectified_left")
    xout_right.setStreamName("rectified_right")
    xout_disp.setStreamName("disparity")
    xout_depth.setStreamName("depth")

    left.setBoardSocket(left_socket)
    right.setBoardSocket(right_socket)

    # OV9282 no OAK-FFC-3P geralmente é 1280x800.
    left.setResolution(dai.MonoCameraProperties.SensorResolution.THE_800_P)
    right.setResolution(dai.MonoCameraProperties.SensorResolution.THE_800_P)

    left.setFps(args.fps)
    right.setFps(args.fps)

    stereo.setDefaultProfilePreset(dai.node.StereoDepth.PresetMode.HIGH_DENSITY)

    stereo.setLeftRightCheck(args.lrcheck)
    stereo.setExtendedDisparity(args.extended)
    stereo.setSubpixel(args.subpixel)

    stereo.initialConfig.setConfidenceThreshold(args.confidence)

    # Mantém depth alinhado ao par mono, não ao RGB.
    stereo.setRectifyEdgeFillColor(0)

    # Se quiser ver mais suave.
    try:
        stereo.initialConfig.setMedianFilter(dai.MedianFilter.KERNEL_7x7)
    except Exception:
        pass

    left.out.link(stereo.left)
    right.out.link(stereo.right)

    stereo.rectifiedLeft.link(xout_left.input)
    stereo.rectifiedRight.link(xout_right.input)
    stereo.disparity.link(xout_disp.input)
    stereo.depth.link(xout_depth.input)

    return pipeline


def colorize_disparity(disp_frame):
    disp = disp_frame.astype(np.float32)

    valid = disp > 0
    if np.count_nonzero(valid) > 50:
        lo = np.percentile(disp[valid], 2)
        hi = np.percentile(disp[valid], 98)
        norm = (disp - lo) / (hi - lo + 1e-6)
    else:
        norm = np.zeros_like(disp, dtype=np.float32)

    norm = np.clip(norm, 0, 1)
    vis = (norm * 255).astype(np.uint8)
    return cv2.applyColorMap(vis, cv2.COLORMAP_TURBO)


def colorize_depth(depth_frame, max_depth_m=5.0):
    depth = depth_frame.astype(np.float32)

    # Depth vem em milímetros.
    depth_m = depth / 1000.0
    valid = depth > 0

    norm = np.zeros_like(depth_m, dtype=np.float32)
    norm[valid] = 1.0 - np.clip(depth_m[valid] / max_depth_m, 0.0, 1.0)

    vis = (norm * 255).astype(np.uint8)
    return cv2.applyColorMap(vis, cv2.COLORMAP_TURBO)


def draw_lines(img, step=40):
    out = img.copy()
    if out.ndim == 2:
        out = cv2.cvtColor(out, cv2.COLOR_GRAY2BGR)

    h, w = out.shape[:2]
    for y in range(0, h, step):
        color = (0, 255, 255) if (y // step) % 2 == 0 else (255, 255, 0)
        cv2.line(out, (0, y), (w, y), color, 1, cv2.LINE_AA)

    return out


def resize_h(img, target_h):
    h, w = img.shape[:2]
    if h == target_h:
        return img
    scale = target_h / h
    return cv2.resize(img, (int(w * scale), target_h), interpolation=cv2.INTER_AREA)


def put_label(img, text):
    out = img.copy()
    if out.ndim == 2:
        out = cv2.cvtColor(out, cv2.COLOR_GRAY2BGR)

    cv2.rectangle(out, (0, 0), (out.shape[1], 34), (0, 0, 0), -1)
    cv2.putText(out, text, (10, 24), cv2.FONT_HERSHEY_SIMPLEX, 0.65, (255, 255, 255), 1, cv2.LINE_AA)
    return out


def fit_image_to_window(img, win_name, margin=20):
    """
    Redimensiona o canvas para caber na janela atual do OpenCV,
    mantendo proporção.
    """
    try:
        x, y, win_w, win_h = cv2.getWindowImageRect(win_name)
    except Exception:
        return img

    if win_w <= 50 or win_h <= 50:
        return img

    target_w = max(1, win_w - margin)
    target_h = max(1, win_h - margin)

    h, w = img.shape[:2]
    scale = min(target_w / w, target_h / h)

    if scale <= 0:
        return img

    new_w = max(1, int(w * scale))
    new_h = max(1, int(h * scale))

    return cv2.resize(img, (new_w, new_h), interpolation=cv2.INTER_AREA)


def compose_2x2_responsive(a, b, c, d):
    """
    Monta 2x2 sem depender de altura fixa.
    Primeiro padroniza os tamanhos relativos, depois o fit final
    é feito conforme a janela.
    """
    imgs = [a, b, c, d]

    # Garante BGR
    fixed = []
    for img in imgs:
        if img.ndim == 2:
            img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
        fixed.append(img)

    # Define tamanho base comum pela menor altura
    base_h = min(img.shape[0] for img in fixed)
    resized = [resize_h(img, base_h) for img in fixed]

    max_w = max(img.shape[1] for img in resized)

    padded = []
    for img in resized:
        h, w = img.shape[:2]
        if w < max_w:
            pad = np.zeros((h, max_w - w, 3), dtype=np.uint8)
            img = np.hstack([img, pad])
        padded.append(img)

    top = np.hstack([padded[0], padded[1]])
    bottom = np.hstack([padded[2], padded[3]])

    return np.vstack([top, bottom])


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

    # Pelo seu fluxo de exportação:
    # left  = CAM_C / NIR
    # right = CAM_B / RE
    parser.add_argument("--left_socket", default="CAM_C")
    parser.add_argument("--right_socket", default="CAM_B")

    parser.add_argument("--fps", type=float, default=10)
    parser.add_argument("--confidence", type=int, default=200)

    parser.add_argument("--lrcheck", action="store_true")
    parser.add_argument("--extended", action="store_true")
    parser.add_argument("--subpixel", action="store_true")

    parser.add_argument("--view_h", type=int, default=360)
    parser.add_argument("--max_depth_m", type=float, default=5.0)
    parser.add_argument("--lines", action="store_true")

    args = parser.parse_args()

    print("[INFO] DepthAI version:", dai.__version__)
    print(f"[INFO] left_socket={args.left_socket}")
    print(f"[INFO] right_socket={args.right_socket}")
    print(f"[INFO] lrcheck={args.lrcheck} extended={args.extended} subpixel={args.subpixel}")

    pipeline = make_pipeline(args)

    win_name = "OAK-FFC StereoDepth Test"
    cv2.namedWindow(win_name, cv2.WINDOW_NORMAL)
    cv2.resizeWindow(win_name, 1280, 800)
    try:
        cv2.setWindowProperty(win_name, cv2.WND_PROP_FULLSCREEN, cv2.WINDOW_NORMAL)
    except Exception:
        pass

    with dai.Device(pipeline) as device:
        print("[INFO] Device connected.")
        print("[INFO] MXID:", device.getMxId())

        try:
            calib = device.readCalibration()
            eeprom = calib.getEepromData()
            print("[INFO] EEPROM version:", eeprom.version)
            print("[INFO] Board:", eeprom.boardName, eeprom.productName)
        except Exception as e:
            print("[WARN] Não consegui ler calibração:", e)

        q_left = device.getOutputQueue("rectified_left", maxSize=4, blocking=False)
        q_right = device.getOutputQueue("rectified_right", maxSize=4, blocking=False)
        q_disp = device.getOutputQueue("disparity", maxSize=4, blocking=False)
        q_depth = device.getOutputQueue("depth", maxSize=4, blocking=False)

        while True:
            in_left = q_left.tryGet()
            in_right = q_right.tryGet()
            in_disp = q_disp.tryGet()
            in_depth = q_depth.tryGet()

            if in_left is None or in_right is None or in_disp is None or in_depth is None:
                key = cv2.waitKey(1) & 0xFF
                if key in [27, ord("q"), ord("Q")]:
                    break
                continue

            left = in_left.getCvFrame()
            right = in_right.getCvFrame()
            disp = in_disp.getCvFrame()
            depth = in_depth.getFrame()

            left_vis = draw_lines(left, 40) if args.lines else cv2.cvtColor(left, cv2.COLOR_GRAY2BGR)
            right_vis = draw_lines(right, 40) if args.lines else cv2.cvtColor(right, cv2.COLOR_GRAY2BGR)

            disp_vis = colorize_disparity(disp)
            depth_vis = colorize_depth(depth, max_depth_m=args.max_depth_m)

            left_vis = put_label(left_vis, f"rectified left ({args.left_socket})")
            right_vis = put_label(right_vis, f"rectified right ({args.right_socket})")
            disp_vis = put_label(disp_vis, "disparity")
            depth_vis = put_label(depth_vis, f"depth <= {args.max_depth_m:.1f}m")

            canvas = compose_2x2_responsive(
                left_vis,
                right_vis,
                disp_vis,
                depth_vis,
            )

            canvas_fit = fit_image_to_window(canvas, win_name)
            cv2.imshow(win_name, canvas_fit)

            key = cv2.waitKey(1) & 0xFF
            if key in [27, ord("q"), ord("Q")]:
                break

        cv2.destroyAllWindows()


if __name__ == "__main__":
    main()