"use client";

import { useCallback, useEffect, useRef, useState } from "react";
import { useAccount } from "wagmi";
import { getRuntimeConfig } from "@/lib/config";
import type {
  PendingRequest,
  OutboundMessage,
} from "@/lib/types";

export function useProxySocket() {
  const { address, isConnected } = useAccount();
  const [pendingRequests, setPendingRequests] = useState<PendingRequest[]>([]);
  const [connected, setConnected] = useState(false);
  const [requiredChainId, setRequiredChainId] = useState<number | null>(null);
  const wsRef = useRef<WebSocket | null>(null);
  const reconnectTimer = useRef<ReturnType<typeof setTimeout> | null>(null);
  const addressRef = useRef(address);
  addressRef.current = address;

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

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

    const { wsUrl } = await getRuntimeConfig();
    const ws = new WebSocket(wsUrl);

    ws.onopen = () => {
      setConnected(true);
      if (addressRef.current) {
        ws.send(
          JSON.stringify({ type: "account_connected", address: addressRef.current } as OutboundMessage)
        );
      }
    };

    ws.onmessage = (event) => {
      try {
        const data = JSON.parse(event.data);
        if (data.type === "signing_request") {
          setPendingRequests((prev) => [
            ...prev,
            { ...data, receivedAt: Date.now() } as PendingRequest,
          ]);
        } else if (data.type === "chain_config") {
          const chainId = parseInt(data.chainId, 16);
          setRequiredChainId(chainId);
        }
      } catch {
        console.warn("Invalid message from proxy:", event.data);
      }
    };

    ws.onclose = () => {
      setConnected(false);
      wsRef.current = null;
      reconnectTimer.current = setTimeout(connect, 2000);
    };

    ws.onerror = () => {
      ws.close();
    };

    wsRef.current = ws;
  }, []);

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

  useEffect(() => {
    if (!connected) return;
    if (isConnected && address) {
      send({ type: "account_connected", address });
    } else {
      send({ type: "account_disconnected" });
    }
  }, [isConnected, address, connected, send]);

  const removeRequest = useCallback((id: string) => {
    setPendingRequests((prev) => prev.filter((r) => r.id !== id));
  }, []);

  const respond = useCallback(
    (id: string, result?: string, error?: { code: number; message: string }) => {
      if (error) {
        send({ id, type: "signing_response", error });
      } else {
        send({ id, type: "signing_response", result: result! });
      }
      removeRequest(id);
    },
    [send, removeRequest]
  );

  return {
    pendingRequests,
    proxyConnected: connected,
    requiredChainId,
    respond,
    removeRequest,
  };
}
