import { useCallback, useEffect, useRef, useState } from "react";
import type { ClientMessage, ServerMessage } from "./types";

const SESSION_KEY = "web_surface_session";
const USER_ID_KEY = "web_surface_user_id";

function getSavedSession(): { sessionToken: string; userId: string } | null {
  const token = localStorage.getItem(SESSION_KEY);
  const uid = localStorage.getItem(USER_ID_KEY);
  if (token && uid) return { sessionToken: token, userId: uid };
  return null;
}

function saveSession(token: string, userId: string) {
  localStorage.setItem(SESSION_KEY, token);
  localStorage.setItem(USER_ID_KEY, userId);
}

function clearSession() {
  localStorage.removeItem(SESSION_KEY);
  localStorage.removeItem(USER_ID_KEY);
}

export interface WSState {
  status: "disconnected" | "connecting" | "connected";
  userId: string | null;
  send: (msg: ClientMessage) => void;
  reconnect: () => void;
}

export function useWebSocket(
  onMessage: (msg: ServerMessage) => void
): WSState {
  const [status, setStatus] = useState<WSState["status"]>("disconnected");
  const [userId, setUserId] = useState<string | null>(null);
  const wsRef = useRef<WebSocket | null>(null);
  const savedRef = useRef(getSavedSession());
  const reconnectRef = useRef<() => void>(() => {});

  const connect = useCallback(() => {
    if (wsRef.current?.readyState === WebSocket.OPEN) return;
    setStatus("connecting");

    const protocol = location.protocol === "https:" ? "wss:" : "ws:";
    const wsUrl = `${protocol}//${location.host}/ws`;
    const ws = new WebSocket(wsUrl);
    wsRef.current = ws;

    ws.onopen = () => {
      setStatus("connected");
      const saved = savedRef.current;
      const authMsg: Record<string, string> = { type: "auth" };
      if (saved) {
        authMsg.session_token = saved.sessionToken;
      }
      ws.send(JSON.stringify(authMsg));
    };

    ws.onmessage = (event) => {
      try {
        const data: ServerMessage = JSON.parse(event.data);
        if (data.type === "connected") {
          savedRef.current = {
            sessionToken: data.session_token,
            userId: data.user_id,
          };
          saveSession(data.session_token, data.user_id);
          setUserId(data.user_id);
        }
        onMessage(data);
      } catch {
        // ignore parse errors
      }
    };

    ws.onclose = () => {
      setStatus("disconnected");
      wsRef.current = null;
    };

    ws.onerror = () => {
      ws.close();
    };
  }, [onMessage]);

  const send = useCallback((msg: ClientMessage) => {
    if (wsRef.current?.readyState === WebSocket.OPEN) {
      wsRef.current.send(JSON.stringify(msg));
    }
  }, []);

  const reconnect = useCallback(() => {
    clearSession();
    savedRef.current = null;
    setUserId(null);
    connect();
  }, [connect]);

  useEffect(() => {
    connect();
    return () => {
      wsRef.current?.close();
    };
  }, [connect]);

  return { status, userId, send, reconnect };
}
