import time
import json

import numpy as np

from multispectral_service import MultiSpectralService
from stream_receiver import StreamReceiver
from pi.raw_processor_core import RawProcessorCore
from pi.raw_processor_preview import RawProcessorPreview
from radiometric_controller import RadiometricController


class MultiSpectralClient:
    def __init__(
        self,
        pi_host="192.168.105.6",
        pc_host="192.168.105.5",
        server_port=5000,
        stream_port=6001,
        timeout=10,
        width=640,
        height=480,
        bayer="GBRG",
        fps=15,
        frame_type="RAW_BRUTO",
        output_dtype="uint8",
        capture_mode="AUTO",
        raw_policy="allow_single",
        module_calibration_json=None,
        radiometric_enabled=False,
    ):
        self.pi_host = pi_host
        self.pc_host = pc_host
        self.server_port = server_port
        self.stream_port = stream_port

        self.width = width
        self.height = height
        self.bayer = bayer
        self.fps = fps
        self.frame_type = frame_type
        self.output_dtype = output_dtype
        self.capture_mode = capture_mode
        self.raw_policy = raw_policy
        self.module_calibration_json = module_calibration_json

        self.svc = MultiSpectralService(
            host=pi_host,
            port=server_port,
            timeout=timeout,
        )

        self.receiver = StreamReceiver(
            host="0.0.0.0",
            port=stream_port,
        )

        self.core = RawProcessorCore(
            sensor_width=width,
            sensor_height=height,
            bayer_pattern=bayer,
            calibration_json_path=module_calibration_json,
        )

        self.preview = RawProcessorPreview(
            sensor_width=width,
            sensor_height=height,
            bayer_pattern=bayer,
        )

        self.last_frame_id = None
        self.status = None
        self.begin_resp = None
        self.applied_params = None

        self.radiometric_enabled = bool(radiometric_enabled)
        self.radiometric_controller = None

    def __enter__(self):
        self.start()
        return self

    def __exit__(self, exc_type, exc, tb):
        self.stop()

    def start(self, print_debug=True):
        if print_debug:
            print(f"[INFO] Verificando módulo em {self.pi_host}:{self.server_port}...")

        if not self.svc.check_connection(2):
            raise RuntimeError("Módulo não encontrado ou não respondeu ao ping.")

        self.receiver.start()
        time.sleep(0.5)

        self.svc.connect()

        if print_debug:
            print("[OK] Módulo conectado.")

        self._configure_module(print_debug=print_debug)
        self._begin_module(print_debug=print_debug)
        self._apply_module_params(print_debug=print_debug)
        self._start_stream(print_debug=print_debug)

        if self.radiometric_enabled:
            self.enable_radiometric_controller()

        return self

    def enable_radiometric_controller(self):
        self.radiometric_controller = RadiometricController(
            client=self,
            enabled=True,
            config_json_path=self.module_calibration_json,
        )
        self.radiometric_controller.sync_from_camera_controls(self.applied_params["camera_settings"])

        return self.radiometric_controller

    def update_radiometry(self, decoded, meta=None):
        if self.radiometric_controller is None:
            return None

        return self.radiometric_controller.update(decoded, meta)

    def _configure_module(self, print_debug=True):
        r0 = self.svc.set_resolution(self.width, self.height)
        r1 = self.svc.set_bayer(self.bayer)
        r2 = self.svc.set_fps(self.fps)
        r3 = self.svc.set_capture_mode(self.capture_mode)
        r4 = self.svc.set_frame_type(self.frame_type)
        r5 = self.svc.set_output_dtype(self.output_dtype)
        if print_debug:
            print("SET RES:", r0)
            print("SET BAYER:", r1)
            print("SET FPS:", r2)
            print("SET CAPTURE MODE:", r3)
            print("SET FRAME TYPE:", r4)
            print("SET OUTPUT DTYPE:", r5)

    def _begin_module(self, print_debug=True):
        self.begin_resp = self.svc.begin(
            frame_type=self.frame_type,
            output_dtype=self.output_dtype,
            capture_mode=self.capture_mode,
        )

        self.status = self.svc.get_status()

        self.svc.validate_module_ready(
            self.status,
            self.frame_type,
            self.raw_policy,
            self.capture_mode,
        )

        if print_debug:
            print("BEGIN:", self.begin_resp)
            print("STATUS:", json.dumps({
                "status": self.status.get("status"),
                "detected_mode": self.status.get("detected_mode"),
                "camera_count_active": self.status.get("camera_count_active"),
                "active_camera_ids": self.status.get("active_camera_ids"),
            }, ensure_ascii=False))

    def _apply_module_params(self, print_debug=True):
        if not self.module_calibration_json:
            return None

        self.applied_params = self.svc.apply_camera_params_json(
            self.module_calibration_json
        )

        if print_debug:
            print("[OK] Parâmetros do módulo aplicados:")
            print(json.dumps(self.applied_params.get("applied"), ensure_ascii=False, indent=2))

        return self.applied_params

    def _start_stream(self, print_debug=True):
        resp = self.svc.start_stream(
            self.pc_host,
            self.stream_port,
            fps=self.fps,
        )

        if print_debug:
            print("START STREAM:", resp)

        return resp

    def get_latest(self):
        return self.receiver.last_frame, self.receiver.last_meta

    def get_next_frame(self, timeout=2.0):
        t0 = time.perf_counter()

        while time.perf_counter() - t0 < timeout:
            frame = self.receiver.last_frame
            meta = self.receiver.last_meta

            if frame is None or meta is None:
                time.sleep(0.001)
                continue

            frame_id = meta.get("frame_id")

            if frame_id != self.last_frame_id:
                self.last_frame_id = frame_id
                return frame, meta

            time.sleep(0.001)

        raise TimeoutError("Timeout aguardando novo frame do stream.")

    def get_next_decoded(self, timeout=2.0, update_radiometry=True):
        frame, meta = self.get_next_frame(timeout=timeout)
        decoded = self.core.decode_stream_cameras(frame, meta)

        if update_radiometry:
            self.update_radiometry(decoded, meta)

        return frame, meta, decoded

    def build_infer_tensor(self, frame, meta, channels_expected, target_size=None):
        return self.core.build_infer_tensor_from_stream(
            frame,
            meta,
            channels_expected=channels_expected,
            target_size=target_size,
        )

    def build_infer_tensor_from_decoded(self, decoded, meta, channels_expected, target_size=None):
        tensor = self.core.fuse_multispec_cameras(decoded, meta, channels_expected)
        return self.core.resize_tensor_chw(tensor, target_size=target_size)

    def build_preview_from_raw_payload(self, frame, meta: dict):
        """
        Gera preview priorizando a câmera RGB (cam2).
        Se cam2 não estiver presente, cai para fallback usando a primeira câmera mono disponível.
        Retorna:
            preview_bgr
            payload_float_preview
            preview_source_id
        """
        payload_sources = meta.get("payload_sources", []) or []

        # Caso multi-payload: tenta usar cam2 primeiro
        if isinstance(frame, dict):
            if "cam2" in frame:
                rgb_frame = frame["cam2"]

                if rgb_frame.ndim != 3 or rgb_frame.shape[2] != 3:
                    raise RuntimeError(f"cam2 recebida mas inválida para preview RGB: shape={rgb_frame.shape}")

                preview_bgr = rgb_frame.copy()
                payload_float = rgb_frame[:, :, ::-1].astype(np.float32) / 255.0
                payload_float = np.transpose(payload_float, (2, 0, 1))

                return preview_bgr, payload_float, "cam2"

            # fallback: usa a primeira câmera mono disponível
            fallback_id = None
            for cid in ("cam0", "cam1"):
                if cid in frame:
                    fallback_id = cid
                    break

            if fallback_id is None:
                raise RuntimeError("Nenhuma câmera disponível no payload para gerar preview")

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

            cam_frames = meta.get("camera_frames", {}) or {}
            cam_meta = cam_frames.get(fallback_id, {})
            bit_depth = int(cam_meta.get("bit_depth", 10))

            raw16 = self.core.unpack_raw10_packed(packed)
            preview_bgr = self.preview.raw16_to_preview_bgr(raw16, bit_depth=bit_depth)

            payload_float = self.core.build_training_rgb(
                raw16,
                output_dtype="float32",
                bit_depth=bit_depth,
            )

            return preview_bgr, payload_float, fallback_id

        # Caso single-payload
        if isinstance(frame, np.ndarray):
            # Se vier HWC/3ch, tratamos como RGB USB
            if frame.ndim == 3 and frame.shape[2] == 3:
                preview_bgr = frame.copy()
                payload_float = frame[:, :, ::-1].astype(np.float32) / 255.0
                payload_float = np.transpose(payload_float, (2, 0, 1))
                return preview_bgr, payload_float, "cam2"

            # Se vier mono packed, fallback antigo
            packed = frame
            if packed.ndim == 3 and packed.shape[2] == 1:
                packed = packed[:, :, 0]

            source_camera = meta.get("source_camera") or {}
            bit_depth = int(source_camera.get("bit_depth", meta.get("source_bit_depth", 10)))

            raw16 = self.core.unpack_raw10_packed(packed)
            preview_bgr = self.preview.raw16_to_preview_bgr(raw16, bit_depth=bit_depth)

            payload_float = self.core.build_training_rgb(
                raw16,
                output_dtype="float32",
                bit_depth=bit_depth,
            )

            return preview_bgr, payload_float, source_camera.get("id", "unknown")

        raise RuntimeError(f"Tipo de frame não suportado para preview: {type(frame)}")


    def stop(self):
        try:
            self.svc.stop_stream()
        except Exception:
            pass

        try:
            self.svc.stop()
        except Exception:
            pass

        try:
            self.svc.disconnect()
        except Exception:
            pass

        try:
            self.receiver.stop()
        except Exception:
            pass