import socket
import json
import numpy as np
import base64
import time
import os
from typing import Optional


class MultiSpectralService:
    def __init__(self, host="192.168.105.6", port=5000, timeout=5):
        self.host = host
        self.port = port
        self.timeout = timeout
        self.sock = None
        self.file = None

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

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

    # =========================================================
    # Conexão
    # =========================================================

    def connect(self):
        if self.sock is not None:
            return

        self.sock = socket.create_connection((self.host, self.port), timeout=self.timeout)
        self.sock.settimeout(self.timeout)
        self.file = self.sock.makefile("r", encoding="utf-8")

    def disconnect(self):
        try:
            if self.file:
                self.file.close()
        except Exception:
            pass

        try:
            if self.sock:
                self.sock.close()
        except Exception:
            pass

        self.file = None
        self.sock = None

    def check_connection(self, timeout: float = None) -> bool:
        old_timeout = self.timeout
        old_sock_timeout = None

        try:
            if timeout is not None:
                self.timeout = timeout

            self.connect()

            if timeout is not None and self.sock is not None:
                old_sock_timeout = self.sock.gettimeout()
                self.sock.settimeout(timeout)

            resp = self._send_command({"cmd": "ping"})
            return resp.get("ok") and resp.get("reply") == "pong"

        except Exception:
            return False

        finally:
            # restaura timeout do socket
            if old_sock_timeout is not None and self.sock is not None:
                try:
                    self.sock.settimeout(old_sock_timeout)
                except Exception:
                    pass

            # restaura timeout do serviço
            self.timeout = old_timeout

    def _send_command(self, payload: dict) -> dict:
        if self.sock is None:
            self.connect()

        data = (json.dumps(payload) + "\n").encode("utf-8")
        self.sock.sendall(data)

        line = self.file.readline()
        if not line:
            self.disconnect()
            raise RuntimeError("Conexão encerrada pelo servidor")

        return json.loads(line.strip())

    # =========================================================
    # Helpers numpy
    # =========================================================

    def _numpy_dtype_from_string(self, dtype_str: str):
        mapping = {
            "uint8": np.uint8,
            "uint16": np.uint16,
            "float32": np.float32,
        }
        if dtype_str not in mapping:
            raise RuntimeError(f"dtype não suportado recebido do Pi: {dtype_str}")
        return mapping[dtype_str]

    def _reshape_array_from_shape(self, raw_bytes: bytes, dtype_str: str, shape: list | tuple):
        np_dtype = self._numpy_dtype_from_string(dtype_str)
        arr = np.frombuffer(raw_bytes, dtype=np_dtype)
        return arr.reshape(tuple(shape))

    def _reshape_array(self, raw_bytes: bytes, dtype_str: str, layout: str, width: int, height: int, channels: int):
        np_dtype = self._numpy_dtype_from_string(dtype_str)
        arr = np.frombuffer(raw_bytes, dtype=np_dtype)

        if layout == "HW":
            return arr.reshape(height, width)

        if layout == "CHW":
            return arr.reshape(channels, height, width)

        if layout == "HWC":
            return arr.reshape(height, width, channels)

        raise RuntimeError(f"Layout não suportado recebido do Pi: {layout}")

    def _decode_single_array(self, raw: bytes, resp: dict):
        dtype_str = resp.get("dtype") or resp.get("output_dtype") or "uint8"
        shape = resp.get("shape")

        if shape:
            return self._reshape_array_from_shape(raw, dtype_str, shape)

        output_layout = resp.get("output_layout", "HW")
        width = int(resp.get("output_width", resp.get("width", 0)))
        height = int(resp.get("output_height", resp.get("height", 0)))
        channels = int(resp.get("output_channels", resp.get("channels", 1)))

        return self._reshape_array(
            raw_bytes=raw,
            dtype_str=dtype_str,
            layout=output_layout,
            width=width,
            height=height,
            channels=channels,
        )

    def _decode_multi_frames_base64(self, resp: dict):
        frames_resp = resp.get("frames", {})
        payload_parts = resp.get("payload_parts", []) or []
        camera_frames = resp.get("camera_frames", {}) or {}

        parts_by_cam = {}
        for part in payload_parts:
            cam_id = part.get("camera_id")
            if cam_id:
                parts_by_cam[cam_id] = part

        frames = {}

        for cam_id, item in frames_resp.items():
            raw = base64.b64decode(item["data"])
            part_meta = parts_by_cam.get(cam_id, {})
            cam_meta = camera_frames.get(cam_id, {})

            dtype_str = part_meta.get("dtype") or resp.get("dtype") or resp.get("output_dtype")
            shape = part_meta.get("shape")

            if dtype_str == "multi" or dtype_str is None:
                # fallback conservador:
                # no protocolo atual do Pi, payload multi nativo trafega bytes crus,
                # inclusive mono packed, então o mais seguro é assumir uint8.
                dtype_str = "uint8"

            if shape:
                arr = self._reshape_array_from_shape(raw, dtype_str, shape)
            else:
                width = int(part_meta.get("width", cam_meta.get("width", 0)))
                height = int(part_meta.get("height", cam_meta.get("height", 0)))
                channels = int(part_meta.get("channels", cam_meta.get("channels", 1)))

                layout = "HWC" if channels > 1 else "HW"

                arr = self._reshape_array(
                    raw_bytes=raw,
                    dtype_str=dtype_str,
                    layout=layout,
                    width=width,
                    height=height,
                    channels=channels,
                )

            frames[cam_id] = arr

        return frames

    def _extract_meta(self, resp: dict, t0: float) -> dict:
        dtype_str = resp.get("dtype") or resp.get("output_dtype") or "uint8"

        return {
            "frame_id": resp.get("frame_id"),
            "frame_type": resp.get("frame_type"),
            "payload_format_version": resp.get("payload_format_version"),

            "output_dtype": resp.get("output_dtype"),
            "dtype": dtype_str,
            "output_layout": resp.get("output_layout"),
            "output_channels": resp.get("output_channels"),
            "output_channel_names": resp.get("output_channel_names"),
            "output_width": resp.get("output_width"),
            "output_height": resp.get("output_height"),

            "payload_sources": resp.get("payload_sources"),
            "payload_complete": resp.get("payload_complete"),

            "source_camera": resp.get("source_camera"),
            "source_cameras": resp.get("source_cameras"),
            "camera_frames": resp.get("camera_frames"),

            "multi_payload": resp.get("multi_payload", False),
            "payload_kind": resp.get("payload_kind"),
            "payload_parts": resp.get("payload_parts"),

            "packed_width": resp.get("packed_width"),
            "packed_height": resp.get("packed_height"),
            "source_width": resp.get("source_width"),
            "source_height": resp.get("source_height"),
            "source_bayer_pattern": resp.get("source_bayer_pattern"),
            "source_bit_depth": resp.get("source_bit_depth"),

            "size": resp.get("size"),
            "ts_pi": resp.get("ts_pi"),
            "ts_pi_monotonic": resp.get("ts_pi_monotonic"),

            "dt_trigger": resp.get("dt_trigger"),
            "dt_settle": resp.get("dt_settle"),
            "dt_capture": resp.get("dt_capture"),
            "dt_process": resp.get("dt_process"),
            "dt_total_pi": resp.get("dt_total_pi"),

            "dt_total_pc": time.perf_counter() - t0,
        }

    # =========================================================
    # Comandos básicos
    # =========================================================

    def ping(self):
        return self._send_command({"cmd": "ping"})

    def get_status(self):
        return self._send_command({"cmd": "get_status"})

    def get_config(self):
        return self._send_command({"cmd": "get_config"})

    def begin(self, frame_type: str = "RAW_BRUTO", output_dtype: str = "uint8", capture_mode: str = "AUTO"):
        return self._send_command({
            "cmd": "begin",
            "frame_type": frame_type,
            "output_dtype": output_dtype,
            "capture_mode": capture_mode,
        })

    def stop(self):
        return self._send_command({"cmd": "stop"})

    def set_fps(self, fps: int):
        return self._send_command({"cmd": "set_fps", "value": fps})

    def set_jpeg_quality(self, quality: int):
        return self._send_command({"cmd": "set_jpeg_quality", "value": quality})

    def set_frame_type(self, frame_type: str):
        return self._send_command({"cmd": "set_frame_type", "value": frame_type})

    def set_output_dtype(self, output_dtype: str):
        return self._send_command({"cmd": "set_output_dtype", "value": output_dtype})

    def set_capture_mode(self, capture_mode: str):
        return self._send_command({"cmd": "set_capture_mode", "value": capture_mode})

    def set_camera_enabled(self, index: int, enabled: bool):
        return self._send_command({
            "cmd": "set_camera_enabled",
            "index": index,
            "enabled": bool(enabled)
        })

    def set_camera_bayer(self, index: int, bayer_pattern: str):
        return self._send_command({
            "cmd": "set_camera_bayer",
            "index": index,
            "pattern": bayer_pattern
        })

    def set_camera_resolution(self, index: int, width: int, height: int):
        return self._send_command({
            "cmd": "set_camera_resolution",
            "index": index,
            "width": width,
            "height": height
        })

    # =========================================================
    # Captura
    # =========================================================

    def capture_frame(self):
        t0 = time.perf_counter()

        resp = self._send_command({"cmd": "capture_frame"})
        if not resp.get("ok"):
            raise RuntimeError(resp.get("error", "Falha ao capturar frame"))

        encoding = resp.get("encoding", "base64")
        meta = self._extract_meta(resp, t0)

        if encoding == "base64":
            raw = base64.b64decode(resp["data"])
            arr = self._decode_single_array(raw, resp)

            meta.update({
                "shape": list(arr.shape),
                "channels": int(arr.shape[0]) if arr.ndim == 3 and resp.get("output_layout") == "CHW"
                           else (int(arr.shape[2]) if arr.ndim == 3 else 1),
                "width": int(arr.shape[2]) if arr.ndim == 3 and resp.get("output_layout") == "CHW"
                         else (int(arr.shape[1]) if arr.ndim == 3 else int(arr.shape[1])),
                "height": int(arr.shape[1]) if arr.ndim == 3 and resp.get("output_layout") == "CHW"
                          else int(arr.shape[0]),
            })
            return arr, meta

        if encoding == "base64-multi":
            frames = self._decode_multi_frames_base64(resp)
            meta.update({
                "frames_meta": resp.get("camera_frames", {}),
                "decoded_shapes": {cam_id: list(arr.shape) for cam_id, arr in frames.items()},
            })
            return frames, meta

        raise RuntimeError(f"encoding não suportado recebido do Pi: {encoding}")

    def capture_frame_array(self):
        return self.capture_frame()

    # =========================================================
    # Stream
    # =========================================================

    def start_stream(self, host: str, port: int, fps: float):
        return self._send_command({
            "cmd": "start_stream",
            "host": host,
            "port": port,
            "fps": fps
        })

    def stop_stream(self):
        return self._send_command({"cmd": "stop_stream"})

    # =========================================================
    # Controles de câmera
    # =========================================================

    def get_camera_controls(self, camera_id: str):
        return self._send_command({
            "cmd": "get_camera_controls",
            "camera_id": camera_id,
        })

    def set_ae_enable(self, camera_id: str, value: bool):
        return self._send_command({
            "cmd": "set_ae_enable",
            "camera_id": camera_id,
            "value": bool(value),
        })

    def set_awb_enable(self, camera_id: str, value: bool):
        return self._send_command({
            "cmd": "set_awb_enable",
            "camera_id": camera_id,
            "value": bool(value),
        })

    def set_exposure_time(self, camera_id: str, exposure_time_us: Optional[int] = None):
        return self._send_command({
            "cmd": "set_exposure_time",
            "camera_id": camera_id,
            "value": exposure_time_us,
        })

    def clear_exposure_time(self, camera_id: str):
        return self._send_command({
            "cmd": "clear_exposure_time",
            "camera_id": camera_id,
        })

    def set_analogue_gain(self, camera_id: str, gain: Optional[float]):
        return self._send_command({
            "cmd": "set_analogue_gain",
            "camera_id": camera_id,
            "value": gain,
        })

    def clear_analogue_gain(self, camera_id: str):
        return self._send_command({
            "cmd": "clear_analogue_gain",
            "camera_id": camera_id,
        })

    def set_colour_gains(self, camera_id: str, r_gain: float, b_gain: float):
        return self._send_command({
            "cmd": "set_colour_gains",
            "camera_id": camera_id,
            "r_gain": r_gain,
            "b_gain": b_gain,
        })

    def clear_colour_gains(self, camera_id: str):
        return self._send_command({
            "cmd": "clear_colour_gains",
            "camera_id": camera_id,
        })

    def get_sensor_modes(self):
        return self._send_command({"cmd": "get_sensor_modes"})
    
    
    def load_camera_params_json(self, path: str) -> dict:
        if not path or not os.path.isfile(path):
            raise FileNotFoundError(f"Arquivo de parâmetros das câmeras não encontrado: {path}")

        with open(path, "r", encoding="utf-8") as f:
            data = json.load(f)

        settings = data.get("camera_settings")
        if not isinstance(settings, dict):
            raise RuntimeError("JSON de calibração inválido: chave camera_settings ausente")

        return settings

    def apply_camera_settings(self, camera_settings: dict) -> dict:
        applied = {}

        for cam_id in ("cam0", "cam1", "cam2"):
            ctrl = camera_settings.get(cam_id)
            if not ctrl:
                applied[cam_id] = {
                    "ok": False,
                    "error": "sem parâmetros no JSON",
                }
                continue

            ae = bool(ctrl.get("ae_enable", False))
            awb = bool(ctrl.get("awb_enable", False))
            exp = ctrl.get("exposure_time_us")
            gain = ctrl.get("analogue_gain")
            colour_gains = ctrl.get("colour_gains")

            result = {}

            result["ae"] = self.set_ae_enable(cam_id, ae)

            if cam_id == "cam2":
                result["awb"] = self.set_awb_enable(cam_id, awb)

            if not ae:
                if exp is not None:
                    result["exposure"] = self.set_exposure_time(cam_id, int(exp))
                if gain is not None:
                    result["gain"] = self.set_analogue_gain(cam_id, float(gain))

            if cam_id == "cam2" and not awb and colour_gains is not None:
                r_gain, b_gain = colour_gains
                result["colour_gains"] = self.set_colour_gains(cam_id, float(r_gain), float(b_gain))

            applied[cam_id] = {
                "ok": True,
                "requested": ctrl,
                "responses": result,
            }

        return applied

    def apply_camera_params_json(self, path: str) -> dict:
        if path is None:
            return {
                "ok": False,
                "path": None,
                "camera_settings": None,
                "applied": None,
            }

        camera_settings = self.load_camera_params_json(path)
        applied = self.apply_camera_settings(camera_settings)

        return {
            "ok": True,
            "path": path,
            "camera_settings": camera_settings,
            "applied": applied,
        }