import json
import socket
import threading
import time
from numcodecs import Blosc
import numpy as np


class StreamReceiver:
    def __init__(self, host="0.0.0.0", port=6001):
        self.host = host
        self.port = port

        self._server_sock = None
        self._client_sock = None
        self._thread = None
        self._running = False

        self.last_frame = None
        self.last_meta = None
        self.last_receive_ts = None

        self._codec = None
        self._codec_signature = None

    @property
    def is_running(self):
        return self._running

    # =========================================================
    # Lifecycle
    # =========================================================

    def start(self):
        if self._running:
            return

        self._running = True
        self._thread = threading.Thread(target=self._worker, daemon=True)
        self._thread.start()

    def stop(self):
        self._running = False

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

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

        if self._thread is not None:
            self._thread.join(timeout=2.0)
            self._thread = None

        self._client_sock = None
        self._server_sock = None

    # =========================================================
    # Socket helpers
    # =========================================================

    def _recv_exact(self, sock: socket.socket, n: int) -> bytes:
        chunks = []
        remaining = n

        while remaining > 0:
            chunk = sock.recv(remaining)
            if not chunk:
                raise ConnectionError("Conexão encerrada durante recv")
            chunks.append(chunk)
            remaining -= len(chunk)

        return b"".join(chunks)

    # =========================================================
    # Worker principal
    # =========================================================

    def _worker(self):
        try:
            with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as server:
                server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
                server.bind((self.host, self.port))
                server.listen(1)
                server.settimeout(1.0)

                self._server_sock = server
                print(f"[INFO] StreamReceiver ouvindo em {self.host}:{self.port}")

                while self._running:
                    try:
                        client, addr = server.accept()
                    except socket.timeout:
                        continue

                    print(f"[INFO] StreamReceiver conectado por {addr}")
                    self._client_sock = client

                    with client:
                        while self._running:
                            header_len = int.from_bytes(self._recv_exact(client, 4), "big")
                            header_bytes = self._recv_exact(client, header_len)
                            header = json.loads(header_bytes.decode("utf-8"))

                            payload_len = int.from_bytes(self._recv_exact(client, 4), "big")
                            payload_comp = self._recv_exact(client, payload_len)

                            self._ensure_codec(header)
                            if header.get("codec_family") == "none":
                                payload = payload_comp
                            else:
                                payload = self._codec.decode(payload_comp)

                            expected = int(header["payload_size_raw"])
                            if len(payload) != expected:
                                raise ValueError(
                                    f"Tamanho descomprimido inválido: {len(payload)} != {expected}"
                                )

                            payload_kind = header.get("payload_kind")
                            if header.get("multi_payload", False) or payload_kind == "multi_array":
                                frame = self._decode_multi_payload(payload, header)
                            else:
                                frame = self._decode_single_payload(payload, header)

                            self.last_frame = frame
                            self.last_meta = header
                            self.last_receive_ts = time.perf_counter()

                    print("[INFO] StreamReceiver cliente desconectado")
                    self._client_sock = None

        except Exception as e:
            print(f"[WARN] StreamReceiver encerrado com erro: {e}")

        finally:
            self._running = False
            self._client_sock = None
            self._server_sock = None

    # =========================================================
    # Decode helpers
    # =========================================================

    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: {dtype_str}")

        return mapping[dtype_str]

    def _reshape_from_shape(self, payload: bytes, dtype_str: str, shape):
        dtype = self._numpy_dtype_from_string(dtype_str)
        arr = np.frombuffer(payload, dtype=dtype)
        return arr.reshape(tuple(shape))

    def _reshape_from_layout(self, payload: bytes, dtype_str: str, layout: str, width: int, height: int, channels: int):
        dtype = self._numpy_dtype_from_string(dtype_str)
        arr = np.frombuffer(payload, dtype=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: {layout}")

    def _decode_single_payload(self, payload: bytes, header: dict):
        payload_parts = header.get("payload_parts", []) or []
        first_part = payload_parts[0] if payload_parts else {}

        dtype_str = first_part.get("dtype") or header.get("dtype") or header.get("output_dtype") or "uint8"
        shape = first_part.get("shape")

        if shape:
            return self._reshape_from_shape(payload, dtype_str, shape)

        height = int(header.get("output_height", header.get("height")))
        width = int(header.get("output_width", header.get("width")))
        channels = int(header.get("output_channels", header.get("channels", 1)))
        layout = header.get("output_layout", "HWC")

        return self._reshape_from_layout(
            payload=payload,
            dtype_str=dtype_str,
            layout=layout,
            width=width,
            height=height,
            channels=channels,
        )

    def _decode_multi_payload(self, payload: bytes, header: dict):
        payload_parts = header.get("payload_parts", []) or []
        camera_frames = header.get("camera_frames", {}) or {}

        frames = {}
        offset = 0

        for part in payload_parts:
            cam_id = part.get("camera_id")
            if not cam_id:
                raise RuntimeError("payload_parts sem camera_id")

            if offset + 4 > len(payload):
                raise RuntimeError("Payload multi truncado ao ler tamanho da parte")

            part_size = int.from_bytes(payload[offset:offset + 4], "big")
            offset += 4

            if offset + part_size > len(payload):
                raise RuntimeError(f"Payload multi truncado ao ler dados de {cam_id}")

            part_bytes = payload[offset:offset + part_size]
            offset += part_size

            cam_meta = camera_frames.get(cam_id, {})

            dtype_str = part.get("dtype") or header.get("dtype") or header.get("output_dtype")
            shape = part.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:
                frame = self._reshape_from_shape(part_bytes, dtype_str, shape)
            else:
                width = int(part.get("width", cam_meta.get("width", 0)))
                height = int(part.get("height", cam_meta.get("height", 0)))
                channels = int(part.get("channels", cam_meta.get("channels", 1)))

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

                frame = self._reshape_from_layout(
                    payload=part_bytes,
                    dtype_str=dtype_str,
                    layout=layout,
                    width=width,
                    height=height,
                    channels=channels,
                )

            frames[cam_id] = frame

        if offset != len(payload):
            raise RuntimeError(
                f"Payload multi com bytes sobrando: consumidos={offset}, total={len(payload)}"
            )

        return frames

    # =========================================================
    # Codec
    # =========================================================

    def _normalize_shuffle(self, shuffle_value):
        if isinstance(shuffle_value, int):
            return shuffle_value

        mapping = {
            "NOSHUFFLE": Blosc.NOSHUFFLE,
            "SHUFFLE": Blosc.SHUFFLE,
            "BITSHUFFLE": Blosc.BITSHUFFLE,
        }

        key = str(shuffle_value).upper()
        if key not in mapping:
            raise ValueError(f"shuffle inválido: {shuffle_value}")

        return mapping[key]

    def _build_codec_from_header(self, header: dict):
        family = header.get("codec_family")
        name = header.get("codec_name")
        params = dict(header.get("codec_params", {}))

        if family != "numcodecs":
            raise ValueError(f"Família de codec não suportada: {family}")

        if name == "blosc":
            params["shuffle"] = self._normalize_shuffle(params.get("shuffle", "SHUFFLE"))
            return Blosc(**params)

        raise ValueError(f"Codec numcodecs não suportado: {name}")

    def _get_codec_signature_from_header(self, header: dict):
        return (
            header.get("codec_family"),
            header.get("codec_name"),
            tuple(sorted(dict(header.get("codec_params", {})).items()))
        )

    def _ensure_codec(self, header: dict):
        family = header.get("codec_family")

        if family == "none":
            self._codec = None
            self._codec_signature = ("none", None, ())
            return

        sig = self._get_codec_signature_from_header(header)
        if self._codec is None or self._codec_signature != sig:
            self._codec = self._build_codec_from_header(header)
            self._codec_signature = sig