import { useEffect, useRef, useState } from 'react';
import type { ServerTextMessage } from '../lib/ssh-console-types';

export type ConnState =
  | { kind: 'no_session' }
  | { kind: 'connecting' }
  | { kind: 'replaying' }
  | { kind: 'connected'; canWrite: boolean; cols: number; rows: number }
  | { kind: 'disconnected'; reason?: string };

export interface ConsoleSessionApi {
  state: ConnState;
  /** Subscribe to PTY output bytes (binary). */
  onOutput(cb: (data: Uint8Array) => void): () => void;
  /** Subscribe to text-frame notices. */
  onNotice(cb: (msg: { severity: string; msg: string }) => void): () => void;
  send(input: string): void;
  sendResize(cols: number, rows: number): void;
  close(): void;
}

/**
 * WS client for the shared SSH console.
 *
 * Auto-reconnects with exponential backoff (1s → 30s cap) on any close,
 * because the server returns no_session if the AI hasn't called
 * SshConsoleEnsure yet — the user opening the SSH tab before the agent
 * is ready is a normal case, not an error. The hook silently keeps
 * retrying until the session appears, at which point the next attempt
 * succeeds and the terminal "comes alive".
 */
export function useConsoleSession(taskId: string | number): ConsoleSessionApi {
  const wsRef = useRef<WebSocket | null>(null);
  const [state, setState] = useState<ConnState>({ kind: 'no_session' });
  const outputListeners = useRef(new Set<(d: Uint8Array) => void>());
  const noticeListeners = useRef(new Set<(m: any) => void>());
  const lastAttachRef = useRef<{ canWrite: boolean; cols: number; rows: number } | null>(null);

  useEffect(() => {
    let cancelled = false;
    let retryDelayMs = 1000;          // start: 1s; doubles each failure
    const MAX_RETRY_MS = 30_000;      // cap at 30s
    let retryTimer: ReturnType<typeof setTimeout> | null = null;

    const url = `${location.origin.replace(/^http/, 'ws')}/api/local/tasks/${encodeURIComponent(String(taskId))}/console/ws`;

    const connect = (): void => {
      if (cancelled) return;
      setState({ kind: 'connecting' });
      const ws = new WebSocket(url);
      ws.binaryType = 'arraybuffer';
      wsRef.current = ws;

      ws.onopen = () => {
        // Successful handshake means the server didn't reject us. Reset
        // the backoff so the next legitimate disconnect retries fast.
        retryDelayMs = 1000;
      };

      ws.onmessage = (ev) => {
        if (typeof ev.data === 'string') {
          let msg: ServerTextMessage;
          try { msg = JSON.parse(ev.data); } catch { return; }
          if (msg.type === 'attach') {
            lastAttachRef.current = { canWrite: msg.can_write, cols: msg.cols, rows: msg.rows };
            setState({ kind: 'replaying' });
          } else if (msg.type === 'replay_end') {
            const a = lastAttachRef.current;
            setState({
              kind: 'connected',
              canWrite: a?.canWrite ?? true,
              cols: a?.cols ?? 80,
              rows: a?.rows ?? 24,
            });
          } else if (msg.type === 'notice') {
            noticeListeners.current.forEach((l) => l(msg));
          } else if (msg.type === 'close') {
            setState({ kind: 'disconnected', reason: msg.reason });
          }
        } else {
          const buf = new Uint8Array(ev.data as ArrayBuffer);
          outputListeners.current.forEach((l) => l(buf));
        }
      };

      ws.onclose = () => {
        if (cancelled) return;
        // If we never got `attach` (server rejected at upgrade) the close
        // looks indistinguishable from "session not ready yet". Always
        // retry — the user opening the tab before the agent is ready is
        // the most common case.
        setState((p) => (p.kind === 'disconnected' ? p : { kind: 'disconnected' }));
        retryTimer = setTimeout(connect, retryDelayMs);
        retryDelayMs = Math.min(retryDelayMs * 2, MAX_RETRY_MS);
      };

      ws.onerror = () => {
        setState({ kind: 'disconnected', reason: 'ws_error' });
        // onerror is followed by onclose; the close handler does the retry.
      };
    };

    connect();

    return () => {
      cancelled = true;
      if (retryTimer) clearTimeout(retryTimer);
      try { wsRef.current?.close(); } catch {}
    };
  }, [taskId]);

  return {
    state,
    onOutput(cb) {
      outputListeners.current.add(cb);
      return () => { outputListeners.current.delete(cb); };
    },
    onNotice(cb) {
      noticeListeners.current.add(cb);
      return () => { noticeListeners.current.delete(cb); };
    },
    send(input) {
      const ws = wsRef.current;
      if (ws && ws.readyState === ws.OPEN) {
        ws.send(new TextEncoder().encode(input));
      }
    },
    sendResize(cols, rows) {
      const ws = wsRef.current;
      if (ws && ws.readyState === ws.OPEN) {
        ws.send(JSON.stringify({ type: 'resize', cols, rows }));
      }
    },
    close() {
      try { wsRef.current?.close(); } catch {}
    },
  };
}
