import { AppState, AppStateStatus } from "react-native";
import { WsOutgoing } from "../types";

type WebSocketMessage = Record<string, unknown>;
type MessageCallback = (data: WebSocketMessage) => void;
type StatusCallback = () => void;
type ErrorCallback = (error: Event) => void;

interface WebSocketClientOptions {
  onMessage?: MessageCallback;
  onOpen?: StatusCallback;
  onClose?: StatusCallback;
  onError?: ErrorCallback;
}

const INITIAL_RECONNECT_DELAY = 1000;
const MAX_RECONNECT_DELAY = 30000;
const RECONNECT_MULTIPLIER = 2;
const LOCAL_TIMEOUT = 2500;
const HEARTBEAT_INTERVAL = 30000; // 30s ping to detect zombie sockets

export class WebSocketClient {
  private ws: WebSocket | null = null;
  private urls: string[] = [];
  private urlIndex: number = 0;
  private reconnectDelay: number = INITIAL_RECONNECT_DELAY;
  private reconnectTimer: ReturnType<typeof setTimeout> | null = null;
  private localTimer: ReturnType<typeof setTimeout> | null = null;
  private heartbeatTimer: ReturnType<typeof setInterval> | null = null;
  private lastMessageAt: number = 0; // timestamp of last received message (any type)
  private shouldReconnect: boolean = false;
  private connected: boolean = false;
  private callbacks: WebSocketClientOptions = {};

  constructor() {
    // When app comes back to foreground, verify the connection
    AppState.addEventListener("change", (state: AppStateStatus) => {
      if (state === "active" && this.shouldReconnect) {
        if (!this.ws || this.ws.readyState === WebSocket.CLOSED || this.ws.readyState === WebSocket.CLOSING) {
          // Socket is definitively dead — force immediate reconnect
          this.reconnectDelay = INITIAL_RECONNECT_DELAY;
          this.urlIndex = 0;
          this.tryUrl();
        } else if (this.connected) {
          // Socket might be alive — send a ping and restart heartbeat
          // to give it a fresh 2-interval window to prove liveness
          this.startHeartbeat();
          this.sendPing();
        }
      }
    });
  }

  setCallbacks(callbacks: WebSocketClientOptions) {
    this.callbacks = callbacks;
  }

  connect(urls: string[]) {
    this.urls = urls.filter(Boolean);
    if (this.urls.length === 0) return;
    this.shouldReconnect = true;
    this.reconnectDelay = INITIAL_RECONNECT_DELAY;
    this.urlIndex = 0;
    this.connected = false;
    this.tryUrl();
  }

  private cleanup() {
    if (this.localTimer) { clearTimeout(this.localTimer); this.localTimer = null; }
    if (this.reconnectTimer) { clearTimeout(this.reconnectTimer); this.reconnectTimer = null; }
    this.stopHeartbeat();
    if (this.ws) {
      const old = this.ws;
      this.ws = null;
      old.onopen = null;
      old.onclose = null;
      old.onerror = null;
      old.onmessage = null;
      try { old.close(); } catch { /* ignore */ }
    }
  }

  private startHeartbeat() {
    this.stopHeartbeat();
    this.lastMessageAt = Date.now();
    this.heartbeatTimer = setInterval(() => {
      const silentMs = Date.now() - this.lastMessageAt;
      if (silentMs > HEARTBEAT_INTERVAL * 2) {
        // No message (including pong) for 2 full intervals — zombie socket
        this.connected = false;
        this.callbacks.onClose?.();
        this.reconnectDelay = INITIAL_RECONNECT_DELAY;
        this.urlIndex = 0;
        this.tryUrl();
        return;
      }
      // Send a ping so the server has something to respond to
      this.sendPing();
    }, HEARTBEAT_INTERVAL);
  }

  private stopHeartbeat() {
    if (this.heartbeatTimer) { clearInterval(this.heartbeatTimer); this.heartbeatTimer = null; }
  }

  private sendPing() {
    if (this.ws && this.ws.readyState === WebSocket.OPEN) {
      try {
        this.ws.send(JSON.stringify({ type: "ping" }));
      } catch {
        // Send failed — socket is dead
        this.connected = false;
        this.callbacks.onClose?.();
        this.reconnectDelay = INITIAL_RECONNECT_DELAY;
        this.urlIndex = 0;
        this.tryUrl();
      }
    }
  }

  private tryUrl() {
    this.cleanup();

    const url = this.urls[this.urlIndex];
    if (!url) return;

    const ws = new WebSocket(url);
    this.ws = ws;

    // If trying local (index 0) and we have a remote fallback,
    // give local 2.5s before switching to remote
    if (this.urlIndex === 0 && this.urls.length > 1) {
      this.localTimer = setTimeout(() => {
        this.localTimer = null;
        if (this.connected) return; // already connected, ignore
        // Local didn't connect in time — try remote
        this.urlIndex = 1;
        this.tryUrl();
      }, LOCAL_TIMEOUT);
    }

    ws.onopen = () => {
      if (ws !== this.ws) return; // stale
      this.connected = true;
      if (this.localTimer) { clearTimeout(this.localTimer); this.localTimer = null; }
      this.reconnectDelay = INITIAL_RECONNECT_DELAY;
      this.startHeartbeat();
      this.callbacks.onOpen?.();
    };

    ws.onmessage = (event) => {
      if (ws !== this.ws) return; // stale
      this.lastMessageAt = Date.now(); // any message proves the connection is alive
      try {
        const data = JSON.parse(event.data) as WebSocketMessage;
        if (data.type === "pong") return; // heartbeat response, don't forward
        this.callbacks.onMessage?.(data);
      } catch {
        this.callbacks.onMessage?.({ type: "text", content: String(event.data) });
      }
    };

    ws.onclose = () => {
      if (ws !== this.ws) return; // stale
      this.connected = false;
      this.callbacks.onClose?.();
      if (this.shouldReconnect) {
        this.scheduleReconnect();
      }
    };

    ws.onerror = () => {
      if (ws !== this.ws) return; // stale
      // Don't do anything here — onclose always fires after onerror
      // and handles reconnect. Just swallow the error event.
    };
  }

  private scheduleReconnect() {
    if (this.reconnectTimer) clearTimeout(this.reconnectTimer);
    this.reconnectTimer = setTimeout(() => {
      this.reconnectTimer = null;
      this.reconnectDelay = Math.min(
        this.reconnectDelay * RECONNECT_MULTIPLIER,
        MAX_RECONNECT_DELAY
      );
      // Alternate between URLs on each reconnect attempt
      if (this.urls.length > 1) {
        this.urlIndex = this.urlIndex === 0 ? 1 : 0;
      }
      this.tryUrl();
    }, this.reconnectDelay);
  }

  disconnect() {
    this.shouldReconnect = false;
    this.connected = false;
    this.cleanup();
  }

  send(message: WsOutgoing) {
    if (this.ws && this.ws.readyState === WebSocket.OPEN) {
      this.ws.send(JSON.stringify(message));
      return true;
    }
    return false;
  }

  get readyState(): number {
    return this.ws?.readyState ?? WebSocket.CLOSED;
  }

  get isConnected(): boolean {
    return this.connected;
  }

  get currentUrl(): string {
    return this.urls[this.urlIndex] ?? "";
  }
}

export const wsClient = new WebSocketClient();