connectionState.ts

import { action, notifications, runtime, tabs } from "webextension-polyfill";

import { getEndpoint, loadSettings } from "#/settings";

type ConnectionState = "connected" | "disconnected" | "unknown";
type ConnectionSource = "app" | "fallback" | null;

interface WalletInfo {
  accounts: string[];
  chainId: string;
  balance: string;
}

let globalConnectionState: ConnectionState = "unknown";
let globalConnectionSource: ConnectionSource = null;
let hasShownNotification = false;

const NOTIFICATION_ID = "ethui-connection-status";

export function resetConnectionState() {
  globalConnectionState = "unknown";
  globalConnectionSource = null;
  hasShownNotification = false;
  updateBadge();

  // Notify popup to re-check connection
  runtime
    .sendMessage({
      type: "connection-state",
      state: "unknown",
      source: null,
    })
    .catch(() => {
      // Popup may not be open, ignore error
    });
}

export function setConnectionState(
  state: ConnectionState,
  source: ConnectionSource = null,
) {
  const previousState = globalConnectionState;
  globalConnectionState = state;

  if (state === "connected") {
    globalConnectionSource = source;
  } else if (state === "disconnected") {
    globalConnectionSource = null;
  }

  // Broadcast state change to any open popups
  runtime
    .sendMessage({
      type: "connection-state",
      state: globalConnectionState,
      source: globalConnectionSource,
    })
    .catch(() => {
      // Popup may not be open, ignore error
    });

  updateBadge();

  // Show notification on first disconnection
  if (
    state === "disconnected" &&
    previousState !== "disconnected" &&
    !hasShownNotification
  ) {
    showDisconnectedNotification();
    hasShownNotification = true;
  }

  // Reset notification flag when connected
  if (state === "connected") {
    hasShownNotification = false;
  }
}

function updateBadge() {
  if (globalConnectionState === "disconnected") {
    action.setBadgeText({ text: "!" });
    action.setBadgeBackgroundColor({ color: "#ef4444" });
  } else if (
    globalConnectionState === "connected" &&
    globalConnectionSource === "fallback"
  ) {
    // Show indicator when using fallback
    action.setBadgeText({ text: "F" });
    action.setBadgeBackgroundColor({ color: "#f59e0b" }); // amber/warning color
  } else {
    action.setBadgeText({ text: "" });
  }
}

function showDisconnectedNotification() {
  if (!notifications?.create) {
    return;
  }

  notifications.create(NOTIFICATION_ID, {
    type: "basic",
    iconUrl: runtime.getURL("icons/ethui-black-128.png"),
    title: "ethui Desktop Not Running",
    message:
      "The ethui desktop app doesn't appear to be running. Click the extension icon for more info.",
  });
}

async function checkConnection(): Promise<ConnectionState> {
  const settings = await loadSettings();
  const endpoint = getEndpoint(settings);

  return new Promise((resolve) => {
    let resolved = false;
    const done = (state: ConnectionState) => {
      if (resolved) return;
      resolved = true;
      clearTimeout(timeout);
      setConnectionState(state, state === "connected" ? "app" : null);
      resolve(state);
    };

    const ws = new WebSocket(endpoint);
    const timeout = setTimeout(() => {
      ws.close();
      done("disconnected");
    }, 3000);

    ws.onopen = () => {
      done("connected");
      ws.close();
    };

    ws.onerror = () => {
      ws.close();
      done("disconnected");
    };

    ws.onclose = () => {
      done("disconnected");
    };
  });
}

async function fetchWalletInfo(): Promise<WalletInfo | null> {
  const settings = await loadSettings();
  const endpoint = getEndpoint(settings, globalConnectionSource);

  return new Promise((resolve) => {
    const ws = new WebSocket(endpoint);
    let requestId = 1;
    const pending = new Map<
      number,
      { resolve: (result: unknown) => void; reject: (err: Error) => void }
    >();

    const rejectAllPending = (reason: string) => {
      const error = new Error(reason);
      for (const { reject } of pending.values()) {
        reject(error);
      }
      pending.clear();
    };

    const sendRequest = (method: string, params: unknown[] = []) => {
      const id = requestId++;
      return new Promise<unknown>((res, rej) => {
        pending.set(id, { resolve: res, reject: rej });
        ws.send(JSON.stringify({ jsonrpc: "2.0", id, method, params }));
      });
    };

    const timeout = setTimeout(() => {
      rejectAllPending("Request timed out");
      ws.close();
      resolve(null);
    }, 5000);

    ws.onopen = async () => {
      try {
        const [accounts, chainId] = await Promise.all([
          sendRequest("eth_accounts"),
          sendRequest("eth_chainId"),
        ]);

        const accountsArray = accounts as string[];
        let balance = "0x0";

        if (accountsArray.length > 0) {
          balance = (await sendRequest("eth_getBalance", [
            accountsArray[0],
            "latest",
          ])) as string;
        }

        clearTimeout(timeout);
        ws.close();
        resolve({
          accounts: accountsArray,
          chainId: chainId as string,
          balance,
        });
      } catch {
        clearTimeout(timeout);
        ws.close();
        resolve(null);
      }
    };

    ws.onmessage = (event) => {
      try {
        const data = JSON.parse(event.data);
        if (data.id && pending.has(data.id)) {
          const { resolve: res } = pending.get(data.id)!;
          pending.delete(data.id);
          res(data.result);
        }
      } catch {
        // Ignore parse errors
      }
    };

    ws.onerror = () => {
      clearTimeout(timeout);
      rejectAllPending("WebSocket error");
      resolve(null);
    };

    ws.onclose = () => {
      clearTimeout(timeout);
      rejectAllPending("WebSocket closed");
    };
  });
}

export function setupConnectionStateListener() {
  const handleMessage = (
    message: unknown,
    _sender: unknown,
    sendResponse: (r: unknown) => void,
  ): true | undefined => {
    if (
      typeof message !== "object" ||
      message === null ||
      !("type" in message)
    ) {
      return;
    }

    const msg = message as { type: string };

    if (msg.type === "get-connection-state") {
      // If state is unknown, check connection before responding
      if (globalConnectionState === "unknown") {
        checkConnection().then((state) => {
          sendResponse({
            type: "connection-state",
            state,
            source: globalConnectionSource,
          });
        });
      } else {
        sendResponse({
          type: "connection-state",
          state: globalConnectionState,
          source: globalConnectionSource,
        });
      }
      return true;
    }

    if (msg.type === "get-wallet-info") {
      fetchWalletInfo().then((info) => {
        sendResponse({ type: "wallet-info", info });
      });
      return true;
    }

    if (msg.type === "check-connection") {
      checkConnection().then((state) => {
        sendResponse({
          type: "connection-state",
          state,
          source: globalConnectionSource,
        });
      });
      return true;
    }
  };

  runtime.onMessage.addListener(
    handleMessage as Parameters<typeof runtime.onMessage.addListener>[0],
  );

  // Handle notification click - open ethui.dev
  notifications.onClicked.addListener((notificationId) => {
    if (notificationId === NOTIFICATION_ID) {
      tabs.create({ url: "https://ethui.dev" });
    }
  });
}