import time
import threading
import base64
import numpy as np


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.raw_processors = {}

    def _get_camera_spec(self, cam_id):
        cam = self.state.get_camera(cam_id)
        if cam is None:
            raise RuntimeError(f"Câmera '{cam_id}' não encontrada no state")
        return cam

    def _ensure_raw_processor_for_camera(self, cam_id):
        cam = self._get_camera_spec(cam_id)
        rp = self.raw_processors.get(cam_id)

        if (
            rp is None or
            rp.sensor_width != cam.width or
            rp.sensor_height != cam.height or
            rp.bayer_pattern.upper() != cam.bayer_pattern.upper()
        ):
            from raw_processor_core import RawProcessorCore
            rp = RawProcessorCore(
                sensor_width=cam.width,
                sensor_height=cam.height,
                bayer_pattern=cam.bayer_pattern,
            )
            self.raw_processors[cam_id] = rp

        return rp

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

        for _ in range(max_attempts):
            frames = self.camera.capture_raw_frames()

            ok = True
            for cam_id in required_sources:
                info = frames.get(cam_id)
                if info is None:
                    ok = False
                    break

                frame, width, height, channels, frame_id, frame_ts = info
                if frame is None or width <= 0 or height <= 0 or channels <= 0:
                    ok = False
                    break

            if ok:
                return frames

            last_frames = frames
            time.sleep(retry_delay_s)

        raise RuntimeError(
            f"Capture retornou frames insuficientes após {max_attempts} tentativas. "
            f"required_sources={required_sources}, received_sources={list((last_frames or {}).keys())}"
        )

    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

            required_sources = list(self.state.payload.sources)

            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_frames = self._capture_with_retry(required_sources)
                cap_end = time.perf_counter()
                dt_capture = cap_end - cap_start

                proc_start = time.perf_counter()
                frame_out, meta_extra = self._build_output_frame(raw_frames)
                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

            camera_frames_meta = {}
            for cam_id, info in raw_frames.items():
                frame, width, height, channels, cam_frame_id, cam_frame_ts = info
                cam = self._get_camera_spec(cam_id)

                camera_frames_meta[cam_id] = {
                    "camera_frame_id": int(cam_frame_id),
                    "camera_frame_ts": cam_frame_ts,
                    "width": int(width),
                    "height": int(height),
                    "channels": int(channels),
                    "bayer_pattern": cam.bayer_pattern,
                    "bit_depth": int(cam.bit_depth),
                }

            meta = {
                "frame_id": frame_id,
                "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,

                "camera_frames": camera_frames_meta,
                "capture_mode_resolved": self.state.resolve_capture_mode(),
                "payload_format_version": self.state.payload_format_version,
            }

            meta.update(meta_extra)
            return frame_out, meta

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

        if isinstance(frame, dict):
            encoded = {}
            total_size = 0

            for cam_id, arr in frame.items():
                frame_bytes = arr.tobytes()
                total_size += len(frame_bytes)
                encoded[cam_id] = {
                    "encoding": "base64",
                    "size": len(frame_bytes),
                    "data": base64.b64encode(frame_bytes).decode("ascii"),
                }

            meta["encoding"] = "base64-multi"
            meta["size"] = total_size
            meta["frames"] = encoded
            return meta

        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_frames):
        if self.state.frame_type == "RAW_BRUTO":
            return self._process_raw_bruto(raw_frames)

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

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

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

    def _convert_output_dtype(self, arr, bit_depth):
        max_sensor_value = (1 << int(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"):
                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}")

        if self.state.output_dtype == "uint16":
            if arr.dtype == "uint16":
                return arr
            if arr.dtype == "uint8":
                return (arr.astype("uint16") << 8)
            if arr.dtype == "float32":
                return (arr * 65535.0).clip(0, 65535).astype("uint16")
            if arr.dtype.kind in ("u", "i"):
                return arr.astype("uint16")
            raise RuntimeError(f"dtype não suportado para uint16: {arr.dtype}")

        if 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_frames):
        sources = list(self.state.payload.sources)

        if len(sources) == 1:
            cam_id = sources[0]
            frame, width, height, channels, cam_frame_id, cam_frame_ts = raw_frames[cam_id]
            cam = self._get_camera_spec(cam_id)

            meta_extra = {
                "frame_type": self.state.frame_type,
                "output_dtype": str(frame.dtype),
                "dtype": str(frame.dtype),
                "output_layout": "HW",
                "output_channels": 1,
                "output_channel_names": ["RAW10_PACKED"],
                "output_width": int(width),
                "output_height": int(height),
                "width": int(width),
                "height": int(height),
                "channels": 1,
                "payload_sources": [cam_id],
                "source_camera": {
                    "id": cam_id,
                    "bayer_pattern": cam.bayer_pattern,
                    "bit_depth": int(cam.bit_depth),
                },
            }
            return frame, meta_extra

        frames_out = {}
        sources_meta = []

        for cam_id in sources:
            frame, width, height, channels, cam_frame_id, cam_frame_ts = raw_frames[cam_id]
            cam = self._get_camera_spec(cam_id)
            frames_out[cam_id] = frame
            sources_meta.append({
                "id": cam_id,
                "width": int(width),
                "height": int(height),
                "bayer_pattern": cam.bayer_pattern,
                "bit_depth": int(cam.bit_depth),
            })

        meta_extra = {
            "frame_type": self.state.frame_type,
            "output_dtype": "multi",
            "dtype": "multi",
            "output_layout": "MULTI_HW",
            "output_channels": len(frames_out),
            "output_channel_names": [f"RAW_{cam_id.upper()}" for cam_id in sources],
            "output_width": None,
            "output_height": None,
            "width": None,
            "height": None,
            "channels": len(frames_out),
            "payload_sources": sources,
            "raw_sources": sources_meta,
        }

        return frames_out, meta_extra

    def _process_rgb(self, raw_frames):
        cam_id = self.state.payload.sources[0]
        packed, packed_width, packed_height, _, _, _ = raw_frames[cam_id]
        cam = self._get_camera_spec(cam_id)
        rp = self._ensure_raw_processor_for_camera(cam_id)

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

        raw16 = rp.unpack_raw10_packed(packed)
        rgb_chw = rp.build_training_rgb(raw16, bit_depth=cam.bit_depth)
        rgb_chw = self._convert_output_dtype(rgb_chw, cam.bit_depth)

        meta_extra = {
            "frame_type": self.state.frame_type,
            "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]),
            "width": int(rgb_chw.shape[2]),
            "height": int(rgb_chw.shape[1]),
            "channels": 3,
            "payload_sources": [cam_id],
            "packed_width": int(packed_width),
            "packed_height": int(packed_height),
            "source_camera": {
                "id": cam_id,
                "bayer_pattern": cam.bayer_pattern,
                "bit_depth": int(cam.bit_depth),
                "source_width": int(cam.width),
                "source_height": int(cam.height),
            },
        }

        return rgb_chw, meta_extra

    def _process_rgbnir(self, raw_frames):
        sources = list(self.state.payload.sources)
        if len(sources) < 2:
            raise RuntimeError("RGBNIR requer duas câmeras ativas")

        cam_rgb_id = sources[0]
        cam_nir_id = sources[1]

        packed_rgb, _, _, _, _, _ = raw_frames[cam_rgb_id]
        packed_nir, _, _, _, _, _ = raw_frames[cam_nir_id]

        cam_rgb = self._get_camera_spec(cam_rgb_id)
        cam_nir = self._get_camera_spec(cam_nir_id)

        rp_rgb = self._ensure_raw_processor_for_camera(cam_rgb_id)
        rp_nir = self._ensure_raw_processor_for_camera(cam_nir_id)

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

        raw16_rgb = rp_rgb.unpack_raw10_packed(packed_rgb)
        raw16_nir = rp_nir.unpack_raw10_packed(packed_nir)

        rgb_chw = rp_rgb.build_training_rgb(raw16_rgb, bit_depth=cam_rgb.bit_depth).astype("float32")
        cam2_rgb = rp_nir.build_training_rgb(raw16_nir, bit_depth=cam_nir.bit_depth).astype("float32")

        min_h = min(rgb_chw.shape[1], cam2_rgb.shape[1])
        min_w = min(rgb_chw.shape[2], cam2_rgb.shape[2])

        rgb_chw = rgb_chw[:, :min_h, :min_w]
        cam2_rgb = cam2_rgb[:, :min_h, :min_w]

        # assumindo ordem [R, G, B]
        re_single = cam2_rgb[0:1, :, :]
        nir_single = cam2_rgb[2:3, :, :]

        rgbnir = np.concatenate([rgb_chw, nir_single, re_single], axis=0)
        rgbnir = self._convert_output_dtype(rgbnir, max(cam_rgb.bit_depth, cam_nir.bit_depth))

        meta_extra = {
            "frame_type": self.state.frame_type,
            "output_dtype": self.state.output_dtype,
            "dtype": str(rgbnir.dtype),
            "output_layout": "CHW",
            "output_channels": 5,
            "output_channel_names": ["R", "G", "B", "NIR", "RE"],
            "output_width": int(rgbnir.shape[2]),
            "output_height": int(rgbnir.shape[1]),
            "width": int(rgbnir.shape[2]),
            "height": int(rgbnir.shape[1]),
            "channels": 5,
            "payload_sources": [cam_rgb_id, cam_nir_id],
            "source_cameras": [
                {
                    "id": cam_rgb_id,
                    "role": cam_rgb.role,
                    "bayer_pattern": cam_rgb.bayer_pattern,
                    "bit_depth": int(cam_rgb.bit_depth),
                },
                {
                    "id": cam_nir_id,
                    "role": cam_nir.role,
                    "bayer_pattern": cam_nir.bayer_pattern,
                    "bit_depth": int(cam_nir.bit_depth),
                },
            ],
            "rgbnir_note": "RGB da cam0; NIR extraído do canal B da cam1; RE extraído do canal R da cam1",
        }

        return rgbnir, meta_extra