import base64
import json
import os
import threading
from typing import Optional, Union

import cv2
import numpy as np
import websocket


def parse_json(raw: Optional[Union[str, bytes]]) -> Optional[dict]:
    """
    Safely parse JSON from string or bytes.

    Args:
        raw: JSON string, bytes, or None

    Returns:
        Parsed dict if successful, None if raw is None, parsing fails, or result is not a dict
    """
    if raw is None:
        return None
    if isinstance(raw, bytes):
        raw = raw.decode("utf-8", errors="replace")
    try:
        result = json.loads(raw)
        return result if isinstance(result, dict) else None
    except (json.JSONDecodeError, TypeError):
        return None


def parse_image(raw: Optional[Union[str, bytes]]) -> Optional[dict]:
    """
    Decode an image message (json with base64 data) and save it as JPEG.

    Args:
        raw: JSON string, bytes, or None

    Returns:
        Parsed dict if successful, None if raw is None, parsing fails, or result is not a dict
    """

    if raw is None:
        return None

    try:
        result = json.loads(raw)
        msg = result["msg"]
    except (json.JSONDecodeError, KeyError):
        print("[Image] Invalid JSON or missing 'msg' field.")
        return None

    data_b64 = msg.get("data")
    if not data_b64:
        print("[Image] Missing 'data' field in message.")
        return None

    # ✅ Ensure output directory exists
    os.makedirs("./camera", exist_ok=True)

    # Case 1: CompressedImage (already JPEG/PNG encoded)
    if "format" in msg and msg["format"].lower() in ["jpeg", "jpg", "png"]:
        image_bytes = base64.b64decode(data_b64)
        path = "./camera/received_image.jpeg"
        with open(path, "wb") as f:
            f.write(image_bytes)
        print(f"[Image] Saved CompressedImage to {path}")
        return result if isinstance(result, dict) else None

    # Case 2: Raw Image (rgb8, bgr8, mono8)
    height, width, encoding = msg.get("height"), msg.get("width"), msg.get("encoding")
    if not all([height, width, encoding]):
        print("[Image] Missing required fields for raw image.")
        return None

    # Decode base64 to numpy array
    image_bytes = base64.b64decode(data_b64)
    img_np = np.frombuffer(image_bytes, dtype=np.uint8)

    # Encoding handlers
    try:
        if encoding == "rgb8":
            img_cv = img_np.reshape((height, width, 3))
            img_cv = cv2.cvtColor(img_cv, cv2.COLOR_RGB2BGR)
        elif encoding == "bgr8":
            img_cv = img_np.reshape((height, width, 3))
        elif encoding == "mono8":
            img_cv = img_np.reshape((height, width))
        else:
            print(f"[Image] Unsupported encoding: {encoding}")
            return None
    except ValueError as e:
        print(f"[Image] Reshape error: {e}")
        return None

    # Save as JPEG with quality 95
    success = cv2.imwrite("./camera/received_image.jpeg", img_cv, [cv2.IMWRITE_JPEG_QUALITY, 95])
    if success:
        print("[Image] Saved raw Image to ./camera/received_image.jpeg")
        return result if isinstance(result, dict) else None
    else:
        return None


class WebSocketManager:
    def __init__(self, ip: str, port: int, default_timeout: float = 2.0):
        self.ip = ip
        self.port = port
        self.default_timeout = default_timeout
        self.ws = None
        self.lock = threading.RLock()

    def set_ip(self, ip: str, port: int):
        """
        Set the IP and port for the WebSocket connection.
        """
        self.ip = ip
        self.port = port
        print(f"[WebSocket] IP set to {self.ip}:{self.port}")

    def connect(self) -> Optional[str]:
        """
        Attempt to establish a WebSocket connection.

        Returns:
            None if successful,
            or an error message string if connection failed.
        """
        with self.lock:
            if self.ws is None or not self.ws.connected:
                try:
                    url = f"ws://{self.ip}:{self.port}"
                    self.ws = websocket.create_connection(url, timeout=self.default_timeout)
                    print(f"[WebSocket] Connected ({self.default_timeout}s timeout)")
                    return None  # no error
                except Exception as e:
                    error_msg = f"[WebSocket] Connection error: {e}"
                    print(error_msg)
                    self.ws = None
                    return error_msg
            return None  # already connected, no error

    def send(self, message: dict) -> Optional[str]:
        """
        Send a JSON-serializable message over WebSocket.

        Returns:
            None if successful,
            or an error message string if send failed.
        """
        with self.lock:
            conn_error = self.connect()
            if conn_error:
                return conn_error  # failed to connect

            if self.ws:
                try:
                    json_msg = json.dumps(message)  # ensure it's JSON-serializable
                    self.ws.send(json_msg)
                    return None  # no error
                except TypeError as e:
                    error_msg = f"[WebSocket] JSON serialization error: {e}"
                    print(error_msg)
                    self.close()
                    return error_msg
                except Exception as e:
                    error_msg = f"[WebSocket] Send error: {e}"
                    print(error_msg)
                    self.close()
                    return error_msg

            return "[WebSocket] Not connected, send aborted."

    def receive(self, timeout: Optional[float] = None) -> Optional[Union[str, bytes]]:
        """
        Receive a single message from rosbridge within the given timeout.

        Args:
            timeout (Optional[float]): Seconds to wait before timing out.
                                     If None, uses the default timeout.

        Returns:
            Optional[str]: JSON string received from rosbridge, or None if timeout/error.
        """
        with self.lock:
            self.connect()
            if self.ws:
                try:
                    # Use default timeout if none specified
                    actual_timeout = timeout if timeout is not None else self.default_timeout
                    # Temporarily set the receive timeout
                    self.ws.settimeout(actual_timeout)
                    raw = self.ws.recv()  # rosbridge sends JSON as a string
                    return raw
                except Exception as e:
                    print(f"[WebSocket] Receive error or timeout: {e}")
                    self.close()
                    return None
            return None

    def request(self, message: dict, timeout: Optional[float] = None) -> dict:
        """
        Send a request to Rosbridge and return the response.

        Args:
            message (dict): The Rosbridge message dictionary to send.
            timeout (Optional[float]): Seconds to wait for a response.
                                     If None, uses the default timeout.

        Returns:
            dict:
                - Parsed JSON response if successful.
                - {"error": "<error message>"} if connection/send/receive fails.
                - {"error": "invalid_json", "raw": <response>} if decoding fails.
        """
        # Attempt to send the message (connect() is called internally in send())
        send_error = self.send(message)
        if send_error:
            return {"error": send_error}

        # Attempt to receive a response (connect() is called internally in receive())
        response = self.receive(timeout=timeout)
        if response is None:
            return {"error": "no response or timeout from rosbridge"}

        # Attempt to parse JSON
        parsed_response = parse_json(response)
        if parsed_response is None:
            print(f"[WebSocket] JSON decode error for response: {response}")
            return {"error": "invalid_json", "raw": response}
        return parsed_response

    def close(self):
        with self.lock:
            if self.ws and self.ws.connected:
                try:
                    self.ws.close()
                    print("[WebSocket] Closed")
                except Exception as e:
                    print(f"[WebSocket] Close error: {e}")
                finally:
                    self.ws = None

    def __enter__(self):
        """Context manager entry - automatically connects."""
        # Don't connect here since we want to maintain the existing pattern
        # where request() handles connection automatically
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        """Context manager exit - automatically closes the connection."""
        self.close()
