import socket
import json
import numpy as np
import base64
import time


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()

    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:
            pass

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

        self.file = None
        self.sock = None

    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())

    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):
        return self._send_command({"cmd": "begin"})

    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_resolution(self, width: int, height: int):
        return self._send_command({
            "cmd": "set_resolution",
            "width": width,
            "height": height
        })
    
    def capture_jpg_disk(self):
        return self._send_command({"cmd": "capture_jpg_disk"})
    
    def capture_jpg_bytes(self):
        return self._send_command({"cmd": "capture_jpg_bytes"})
    
    def capture_jpg_base64(self):
        resp = self._send_command({"cmd": "capture_jpg_base64"})
        if not resp.get("ok"):
            raise RuntimeError(resp.get("error", "Falha ao capturar JPG"))
        return base64.b64decode(resp["data"])
    
    def capture_frame_array(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"))
        
        raw = base64.b64decode(resp["data"])
        width = resp["width"]
        height = resp["height"]
        channels = resp["channels"]
        arr = np.frombuffer(raw, dtype=np.uint8).reshape(height, width, channels)

        t1 = time.perf_counter()

        meta = {
            "frame_type": resp.get("frame_type"),
            "width": width,
            "height": height,
            "channels": channels,
            "dtype": resp.get("dtype"),
            "size": resp.get("size"),

            # 👇 TELEMETRIA PI
            "ts_pi": resp.get("ts_pi"),
            "dt_trigger": resp.get("dt_trigger"),
            "dt_settle": resp.get("dt_settle"),
            "dt_capture": resp.get("dt_capture"),
            "dt_total_pi": resp.get("dt_total_pi"),

            # 👇 TELEMETRIA PC
            "dt_total_pc": t1 - t0,
        }

        return arr, meta

    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"})