import React from "react";
import { WsApi, type WsMessage } from "../../api/WsApi";
import { useSnackbar } from "../../hooks/contexts_provider/SnackbarProvider";
import CircleIcon from "@mui/icons-material/Circle";
import { Tooltip } from "@mui/material";

export const WSState = {
  Closed: "Closed",
  Connected: "Connected",
  Error: "Error",
} as const;

export function MatrixWS(p: {
  onMessage: (msg: WsMessage) => void;
  onStateChange?: (state: string) => void;
}): React.ReactElement {
  const snackbar = useSnackbar();

  // Keep only the latest version of onMessage
  const cbRef = React.useRef(p.onMessage);
  React.useEffect(() => {
    cbRef.current = p.onMessage;
  }, [p.onMessage]);

  // Keep only the latest version of onStateChange
  const stateCbRef = React.useRef(p.onStateChange);
  React.useEffect(() => {
    stateCbRef.current = p.onStateChange;
  }, [p.onStateChange]);

  const [state, setState] = React.useState<string>(WSState.Closed);
  const wsId = React.useRef<number | undefined>(undefined);
  const [connCount, setConnCount] = React.useState(0);

  React.useEffect(() => {
    const id = Math.random();
    const ws = new WebSocket(WsApi.WsURL);
    wsId.current = id;

    // Open
    ws.onopen = () => {
      if (wsId.current != id) return;

      setState(WSState.Connected);
      stateCbRef.current?.(WSState.Connected);
    };

    // Error
    ws.onerror = (e) => {
      if (wsId.current != id) return;

      console.error(`WS Debug error!`, e);
      snackbar(`WebSocket error!`);
      setState(WSState.Error);
      stateCbRef.current?.(WSState.Error);

      setTimeout(() => setConnCount(connCount + 1), 500);
    };

    // Close
    ws.onclose = () => {
      if (wsId.current !== id) return;
      setState(WSState.Closed);
      stateCbRef.current?.(WSState.Closed);
      wsId.current = undefined;
    };

    // Message
    ws.onmessage = (msg) => {
      if (wsId.current !== id) return;

      const dec = JSON.parse(msg.data);
      console.info("WS message", dec);
      cbRef.current(dec);
    };

    return () => ws.close();
  }, [connCount, snackbar]);

  return (
    <Tooltip title={state}>
      <CircleIcon
        color={
          state === WSState.Connected
            ? "success"
            : state === WSState.Error
            ? "error"
            : undefined
        }
      />
    </Tooltip>
  );
}