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

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

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

        self._client_sock = None
        self._server_sock = None

    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)

    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 = header["payload_size_raw"]
                            if len(payload) != expected:
                                raise ValueError(
                                    f"Tamanho descomprimido inválido: {len(payload)} != {expected}"
                                )
                            
                            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")

                            dtype = self._numpy_dtype_from_header(header)

                            arr = np.frombuffer(payload, dtype=dtype)

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

                            elif layout == "CHW":
                                frame = arr.reshape(channels, height, width)

                            elif layout == "HWC":
                                frame = arr.reshape(height, width, channels)

                            else:
                                raise RuntimeError(f"Layout não suportado: {layout}")

                            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


    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

    
    def _numpy_dtype_from_header(self, header: dict):
        dtype_str = header.get("dtype") or header.get("output_dtype") or "uint8"

        mapping = {
            "uint8": np.uint8,
            "float32": np.float32,
            "uint16": np.uint16,
        }

        if dtype_str not in mapping:
            raise RuntimeError(f"dtype não suportado: {dtype_str}")

        return mapping[dtype_str]