import cv2
import numpy as np
from typing import Optional


class RawProcessorPreview:
    def __init__(self, sensor_width: int, sensor_height: int, bayer_pattern: str = "GBRG"):
        self.sensor_width = sensor_width
        self.sensor_height = sensor_height
        self.bayer_pattern = bayer_pattern.upper()

    def raw16_to_vis8(
        self, raw16: np.ndarray, 
        black_level: Optional[int] = None, 
        white_level: Optional[int] = None, 
        gamma: float = 2.2, 
        bit_depth: int = 10
    ) -> np.ndarray:
        """
        Conversão para visualização:
        - auto-level
        - gamma
        """
        max_val = float((1 << bit_depth) - 1)

        raw = raw16.astype(np.float32)

        if black_level is None:
            black_level = float(raw.min())
        if white_level is None:
            white_level = float(raw.max())

        if white_level <= black_level:
            norm = raw / max_val
        else:
            norm = (raw - black_level) / (white_level - black_level)

        norm = np.clip(norm, 0.0, 1.0)

        if gamma is not None and gamma > 0:
            norm = np.power(norm, 1.0 / gamma)

        return (norm * 255.0).clip(0, 255).astype(np.uint8)

    def _debayer_code(self):
        mapping = {
            # Troque de BayerGB para BayerGR para inverter R e B
            "GBRG": cv2.COLOR_BayerGR2BGR, 
            "GRBG": cv2.COLOR_BayerGB2BGR,
            "RGGB": cv2.COLOR_BayerBG2BGR,
            "BGGR": cv2.COLOR_BayerRG2BGR,
        }
        if self.bayer_pattern not in mapping:
            raise ValueError(f"Padrão Bayer não suportado: {self.bayer_pattern}")
        return mapping[self.bayer_pattern]

    def apply_preview_white_balance(self, bgr: np.ndarray, strength: float = 1.0) -> np.ndarray:
        """
        Gray-world simples para deixar o preview mais agradável.
        Não usar no raw de treino.
        """
        img = bgr.astype(np.float32)

        mean_b = float(img[:, :, 0].mean())
        mean_g = float(img[:, :, 1].mean())
        mean_r = float(img[:, :, 2].mean())

        mean_gray = (mean_b + mean_g + mean_r) / 3.0

        eps = 1e-6
        gain_b = mean_gray / max(mean_b, eps)
        gain_g = mean_gray / max(mean_g, eps)
        gain_r = mean_gray / max(mean_r, eps)

        # strength=1 aplica total, strength=0 não aplica
        gain_b = 1.0 + (gain_b - 1.0) * strength
        gain_g = 1.0 + (gain_g - 1.0) * strength
        gain_r = 1.0 + (gain_r - 1.0) * strength

        img[:, :, 0] *= gain_b
        img[:, :, 1] *= gain_g
        img[:, :, 2] *= gain_r

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

    def apply_preview_contrast(self, bgr: np.ndarray, alpha: float = 1.08, beta: float = 0.0) -> np.ndarray:
        """
        Ajuste leve de contraste/brilho para preview.
        """
        out = cv2.convertScaleAbs(bgr, alpha=alpha, beta=beta)
        return out

    def raw16_to_preview_bgr(
        self,
        raw16: np.ndarray,
        gamma: float = 2.2,
        wb_strength: float = 0.8,
        apply_wb: bool = True,
        apply_contrast: bool = True,
        bit_depth: int = 10,
    ) -> np.ndarray:
        """
        Pipeline de preview bonito:
        1. auto-level + gamma no mosaico
        2. demosaic
        3. white balance simples
        4. leve contraste final
        """
        vis8 = self.raw16_to_vis8(raw16, gamma=gamma, bit_depth=bit_depth)
        bgr = cv2.cvtColor(vis8, self._debayer_code())

        if apply_wb:
            bgr = self.apply_preview_white_balance(bgr, strength=wb_strength)

        if apply_contrast:
            bgr = self.apply_preview_contrast(bgr, alpha=1.08, beta=0.0)

        return bgr

    def raw16_to_preview_jpg_bytes(self, raw16: np.ndarray, jpeg_quality: int = 95) -> bytes:
        bgr = self.raw16_to_preview_bgr(raw16)
        ok, enc = cv2.imencode(".jpg", bgr, [int(cv2.IMWRITE_JPEG_QUALITY), int(jpeg_quality)])
        if not ok:
            raise RuntimeError("Falha ao codificar preview JPG")
        return enc.tobytes()