import time
import threading
import base64


class FrameService:
    def __init__(self, state, trigger_manager, camera_manager):
        self.state = state
        self.trigger = trigger_manager
        self.camera = camera_manager
        self._capture_lock = threading.RLock()

        self._ensure_raw_processor()

    def _ensure_raw_processor(self):
        if (
            not hasattr(self, "raw_processor") or
            self.raw_processor.sensor_width != self.state.width or
            self.raw_processor.sensor_height != self.state.height or
            self.raw_processor.bayer_pattern.upper() != self.state.source_bayer_pattern.upper()
        ):
            from raw_processor_core import RawProcessorCore
            self.raw_processor = RawProcessorCore(
                sensor_width=self.state.width,
                sensor_height=self.state.height,
                bayer_pattern=self.state.source_bayer_pattern,
            )

    def _capture_with_retry(self, max_attempts=10, retry_delay_s=0.02):
        last_info = None

        for _ in range(max_attempts):
            info = self.camera.capture_raw_frame()

            if (
                info is not None and
                info.get("frame") is not None and
                info.get("width", 0) > 0 and
                info.get("height", 0) > 0 and
                info.get("channels", 0) > 0
            ):
                return info

            last_info = info
            time.sleep(retry_delay_s)

        if last_info is None:
            raise RuntimeError(f"Capture retornou frame vazio após {max_attempts} tentativas")

        raise RuntimeError(
            f"Capture retornou frame vazio após {max_attempts} tentativas: "
            f"width={last_info.get('width')}, "
            f"height={last_info.get('height')}, "
            f"channels={last_info.get('channels')}, "
            f"camera_frame_id={last_info.get('frame_id')}, "
            f"camera_frame_ts={last_info.get('timestamp')}"
        )

    def capture_frame_raw(self):
        with self._capture_lock:
            if not self.state.initialized:
                raise RuntimeError("Módulo não inicializado")

            t0_perf = time.perf_counter()
            t0_unix = time.time()

            dt_trigger = 0.0
            dt_settle = 0.0
            dt_capture = 0.0
            dt_process = 0.0

            try:
                if self.state.trigger_enabled:
                    trig_start = time.perf_counter()
                    self.trigger.pulse()
                    trig_end = time.perf_counter()
                    dt_trigger = trig_end - trig_start

                    settle_ms = float(getattr(self.state, "trigger_settle_delay_ms", 0.0) or 0.0)
                    if settle_ms > 0:
                        settle_start = time.perf_counter()
                        time.sleep(settle_ms / 1000.0)
                        settle_end = time.perf_counter()
                        dt_settle = settle_end - settle_start

                cap_start = time.perf_counter()
                raw_frame_info = self._capture_with_retry()
                cap_end = time.perf_counter()
                dt_capture = cap_end - cap_start

                proc_start = time.perf_counter()
                if self.state.frame_type != "RAW_BRUTO":
                    self._ensure_raw_processor()
                frame_out, meta_extra = self._build_output_frame(raw_frame_info)
                proc_end = time.perf_counter()
                dt_process = proc_end - proc_start

            except Exception as e:
                raise RuntimeError(f"Falha durante captura/processamento de frame: {e}") from e

            if frame_out is None:
                raise RuntimeError("Frame processado retornou nulo")

            total_end = time.perf_counter()

            self.state.stream_frame_id += 1
            frame_id = self.state.stream_frame_id

            meta = {
                "frame_id": frame_id,
                "camera_frame_id": int(raw_frame_info["frame_id"]),
                "camera_frame_ts": raw_frame_info["timestamp"],

                "ts_pi": t0_unix,
                "ts_pi_monotonic": t0_perf,

                "dt_trigger": dt_trigger,
                "dt_settle": dt_settle,
                "dt_capture": dt_capture,
                "dt_process": dt_process,
                "dt_total_pi": total_end - t0_perf,

                "source_bayer_pattern": self.state.source_bayer_pattern,
                "source_bit_depth": self.state.source_bit_depth,
            }

            meta.update(meta_extra)

            return frame_out, meta

    def capture_frame_base64(self):
        frame, meta = self.capture_frame_raw()

        frame_bytes = frame.tobytes()
        meta["encoding"] = "base64"
        meta["size"] = len(frame_bytes)
        meta["data"] = base64.b64encode(frame_bytes).decode("ascii")
        return meta

    def _build_output_frame(self, raw_frame_info):
        if self.state.frame_type == "RAW_BRUTO":
            return self._process_raw_bruto(raw_frame_info)

        if self.state.frame_type == "RGB":
            return self._process_rgb(raw_frame_info)

        if self.state.frame_type == "RGBNIR":
            return self._process_rgbnir(raw_frame_info)

        raise RuntimeError(f"frame_type inválido: {self.state.frame_type}")

    def _convert_output_dtype(self, arr):
        max_sensor_value = (1 << int(self.state.source_bit_depth)) - 1
        if self.state.output_dtype == "uint8":
            if arr.dtype == "uint8":
                return arr
            if arr.dtype == "float32":
                return (arr * 255.0).clip(0, 255).astype("uint8")
            if arr.dtype.kind in ("u", "i"):
                # assume 10-bit/16-bit vindo do pipeline
                max_val = arr.max() if arr.size > 0 else 0
                if max_val <= 255:
                    return arr.astype("uint8")
                return ((arr.astype("float32") / max_sensor_value) * 255.0).clip(0, 255).astype("uint8")
            raise RuntimeError(f"dtype não suportado para uint8: {arr.dtype}")

        elif self.state.output_dtype == "float32":
            if arr.dtype == "float32":
                return arr
            if arr.dtype == "uint8":
                return arr.astype("float32") / 255.0
            if arr.dtype.kind in ("u", "i"):
                max_val = arr.max() if arr.size > 0 else 0
                if max_val <= 255:
                    return arr.astype("float32") / 255.0
                return arr.astype("float32") / max_sensor_value
            raise RuntimeError(f"dtype não suportado para float32: {arr.dtype}")

        raise RuntimeError(f"output_dtype inválido: {self.state.output_dtype}")
        
    def _process_raw_bruto(self, raw_frame_info):
        frame = raw_frame_info["frame"]

        meta_extra = {
            "frame_type": self.state.frame_type,
            "payload_format_version": self.state.payload_format_version,
            "output_dtype": str(frame.dtype),
            "dtype": str(frame.dtype),
            "output_layout": "HW",
            "output_channels": 1,
            "output_channel_names": ["RAW10_PACKED"],
            "output_width": int(raw_frame_info["width"]),
            "output_height": int(raw_frame_info["height"]),
            "source_width": int(self.state.width),
            "source_height": int(self.state.height),
            "packed_width": int(raw_frame_info["width"]),
            "packed_height": int(raw_frame_info["height"]),
            "width": int(raw_frame_info["width"]),
            "height": int(raw_frame_info["height"]),
            "channels": 1,
        }

        return frame, meta_extra

    def _process_rgb(self, raw_frame_info):
        packed = raw_frame_info["frame"]
        b_depth = self.state.source_bit_depth

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

        raw16 = self.raw_processor.unpack_raw10_packed(packed)
        rgb_chw = self.raw_processor.build_training_rgb(raw16, bit_depth=b_depth)

        rgb_chw = self._convert_output_dtype(rgb_chw)

        meta_extra = {
            "frame_type": self.state.frame_type,
            "payload_format_version": self.state.payload_format_version,
            "output_dtype": self.state.output_dtype,
            "dtype": str(rgb_chw.dtype),
            "output_layout": "CHW",
            "output_channels": 3,
            "output_channel_names": ["R", "G", "B"],
            "output_width": int(rgb_chw.shape[2]),
            "output_height": int(rgb_chw.shape[1]),
            "source_width": int(self.state.width),
            "source_height": int(self.state.height),
            "width": int(rgb_chw.shape[2]),
            "height": int(rgb_chw.shape[1]),
            "channels": 3,
            "packed_width": int(raw_frame_info["width"]),
            "packed_height": int(raw_frame_info["height"]),
        }
        
        return rgb_chw, meta_extra

    def _process_rgbnir(self, raw_frame_info):
        raise NotImplementedError("RGBNIR ainda não implementado")