// Horde hub WebSocket client for Node.js (>= 22, uses the global WebSocket).
// Modeled on the browser client in reference/backend.js: nonce-tracked actions,
// action listeners, orphan handlers for the auth response, auto-reconnect.

// Mirrors the browser Stream: accumulates live content and exposes it as a
// string via getContent() once the server-side stream reaches eof.
export class Stream {
  constructor(id) {
    this.id = id;
    this.partialContent = "";
    this.eof = false;
    this.listeners = { data: [], end: [] };
    this.contentPromise = new Promise((resolve) => {
      this.on('end', () => resolve(this.partialContent));
    });
  }

  feed(chunk) {
    if (this.eof) return;
    this.partialContent += chunk;
    this.emit('data', chunk);
  }

  close() {
    if (this.eof) return;
    this.eof = true;
    this.emit('end', this.partialContent);
  }

  getContent() {
    return this.contentPromise;
  }

  emit(event, ...args) {
    for (const fn of this.listeners[event]) fn(...args);
  }

  on(event, callback) {
    this.listeners[event].push(callback);
    if (event === 'data' && this.partialContent) callback(this.partialContent);
    if (event === 'end' && this.eof) callback(this.partialContent);
  }
}

export default class Horde {
  constructor({ host, token = null }) {
    this.host = host;
    this.token = token;
    this.socket = null;
    this.nextNonce = 0;
    this.pending = new Map();          // nonce -> single { resolve, reject } or streaming { run, deferred }
    this.actionListeners = new Map();  // action -> [callback(packet)]
    this.orphanHandlers = [];          // callbacks for nonce-less messages (auth replies)
    this.connectListeners = [];
    this.reconnectTimeout = null;
    this.stopped = false;
    this.pingTimeout = null;          // Tracks the heartbeat interval
  }

  connect() {
    if (this.stopped) return;
    if (this.socket && (this.socket.readyState === WebSocket.CONNECTING || this.socket.readyState === WebSocket.OPEN)) return;

    const socket = new WebSocket(this.host + "/ws");
    this.socket = socket;

    let promiseInner = {};
    const connectPromise = new Promise((resolve, reject) => {
      promiseInner = {resolve, reject};
    });

    socket.onopen = async () => {
      console.log(`Connected to ${this.host}`);
      try {
        if (this.token) {
          socket.send(JSON.stringify({ token: this.token }));
          const response = await this.waitOrphan();
          if (response && response.error) {
            console.error('Hub authentication failed:', response.error);
            socket.close();
            return;
          }
        } else {
          console.warn('No auth token configured; the hub may reject this session');
        }
        for (const cb of this.connectListeners) {
          try { cb(); } catch (err) { console.error('Connect listener failed:', err); }
        }
      } catch (err) {
        console.error('Error during hub login:', err);
        socket.close();
      }

      // Start heartbeat
      this.schedulePing();

      if (promiseInner) {
        promiseInner.resolve();
        promiseInner = null;
      }
    };

    socket.onmessage = (event) => {
      if (typeof event.data === 'string') this.handleMessage(event.data);
    };

    socket.onclose = () => {
      this.socket = null;
      this.clearPending();
      if (!this.stopped) {
        console.log('Disconnected from hub. Reconnecting in 1s...');
        this.reconnectTimeout = setTimeout(() => this.connect(), 1000);
      }
      this.pingTimeout = null;
    };

    socket.onerror = (event) => {
      console.error('WebSocket error:', event.message || event.error || event);

      if (promiseInner) {
        promiseInner.reject(event.message || event.error || event);
        promiseInner = null;
      }
    };

    return connectPromise;
  }

  close() {
    this.stopped = true;
    if (this.reconnectTimeout) clearTimeout(this.reconnectTimeout);
    if (this.pingTimeout) clearTimeout(this.pingTimeout);
    if (this.socket) {
      this.socket.onclose = null;
      this.socket.close();
      this.socket = null;
    }
    this.clearPending();
  }

  clearPending() {
    for (const entry of this.pending.values()) {
      const reject = entry.streaming ? entry.deferred.reject : entry.reject;
      reject(new Error('Connection lost'));
    }
    this.pending.clear();
  }

  onConnect(cb) {
    this.connectListeners.push(cb);
  }

  onAction(action, cb) {
    let list = this.actionListeners.get(action);
    if (!list) {
      list = [];
      this.actionListeners.set(action, list);
    }
    list.push(cb);
  }

  waitOrphan() {
    return new Promise((resolve) => this.orphanHandlers.push(resolve));
  }

  send(value) {
    if (!this.socket || this.socket.readyState !== WebSocket.OPEN) {
      throw new Error('WebSocket is not open');
    }
    this.socket.send(JSON.stringify(value));
  }

  // Sends an action. Without a callback this is a single-response request.
  // With a callback it is a multiple-result action (like the browser client's
  // sendAction): the callback runs once per packet carrying this nonce, and the
  // returned promise settles only when the callback calls deferred.resolve()
  // (or deferred.reject()). This is the shape used by read-stream and
  // prompt-agent, which both emit many packets per request.
  sendAction(action, body = {}, callback = null) {
    const nonce = this.nextNonce++;

    if (callback) {
      let resolveInner, rejectInner;
      const promise = new Promise((res, rej) => { resolveInner = res; rejectInner = rej; });
      const deferred = { promise, resolve: resolveInner, reject: rejectInner };
      this.pending.set(nonce, {
        streaming: true,
        run: (packet) => callback(packet, deferred),
        deferred,
      });
      // Ensure cleanup runs whichever way the action settles, and swallow the
      // rejection on the derived chain so it never surfaces as unhandled.
      promise.catch(() => {}).finally(() => this.pending.delete(nonce));
      this.send({ action, nonce, body });
      return promise;
    }

    const promise = new Promise((resolve, reject) => {
      this.pending.set(nonce, { streaming: false, resolve, reject });
    });
    promise.catch(() => {}); // avoid unhandled rejections when nobody awaits
    this.send({ action, nonce, body });
    return promise;
  }

  // Reads a server-side stream into a Stream object. Content arrives as a
  // sequence of read-stream packets sharing one nonce; eof closes the stream.
  readStream(streamId) {
    const stream = new Stream(streamId);
    this.sendAction("read-stream", { streamId }, (packet, deferred) => {
      if (packet.eof) {
        stream.close();
        deferred.resolve(stream.partialContent);
      } else if (packet.data != null) {
        stream.feed(packet.data);
      }
    }).catch(() => stream.close());
    return stream;
  }

  // Runs a prompt and yields every generated inference message in arrival
  // order. prompt-agent is a multiple-result action; assistant/reasoning
  // content is an open stream ({streamId}) which is read to completion before
  // the message is yielded, so callers always see plain string content.
  async *prompt(agentId, body) {
    const queue = [];
    let notify = null;
    let finished = false;
    let failure = null;

    const wake = () => {
      if (notify) { const resolve = notify; notify = null; resolve(); }
    };

    this.sendAction(
      "prompt-agent",
      { agentId, ...body },
      (packet, deferred) => {
        if (packet.eof) {
          finished = true;
          deferred.resolve(packet);
          wake();
        } else {
          queue.push(packet.data);
          wake();
        }
      }
    ).catch((err) => {
      failure = err;
      finished = true;
      wake();
    });

    while (true) {
      if (failure) throw failure;
      if (queue.length) {
        yield await this.resolveStreams(queue.shift());
        continue;
      }
      if (finished) return;
      await new Promise((resolve) => { notify = resolve; });
    }
  }

  // Resolves any open Stream in a message's content to its final text.
  async resolveStreams(msg) {
    const content = msg && msg.content;
    if (content && typeof content === 'object' && content.streamId != null) {
      const stream = this.readStream(content.streamId);
      const full = await stream.getContent();
      return { ...msg, content: full };
    }
    return msg;
  }

  handleMessage(raw) {
    let packet;
    try {
      packet = JSON.parse(raw);
    } catch (err) {
      console.error('Failed to parse hub message:', raw, err);
      return;
    }

    if (packet.nonce != null) {
      const entry = this.pending.get(packet.nonce);
      if (!entry) {
        console.warn(`Orphan response for nonce ${packet.nonce}`);
        return;
      }

      if (entry.streaming) {
        // Multiple-result action: keep the entry alive and hand every packet
        // (including the terminal eof) to the callback. The callback settles
        // the promise once the action is complete.
        if (packet.error) entry.deferred.reject(new Error(packet.error));
        else entry.run(packet);
        return;
      }

      this.pending.delete(packet.nonce);
      if (packet.error) entry.reject(new Error(packet.error));
      else entry.resolve(packet);
      return;
    }

    if (packet.action) {
      const listeners = this.actionListeners.get(packet.action);
      if (listeners && listeners.length) {
        for (const cb of listeners) {
          try { cb(packet); } catch (err) {
            console.error(`Listener for '${packet.action}' failed:`, err);
          }
        }
        return;
      }
    }

    if (this.orphanHandlers.length) {
      this.orphanHandlers.shift()(packet);
      return;
    }

    if (packet.error) console.error('Hub error:', packet.error);
    else {
      //console.warn('Ignored hub message:', packet.action);
    }
  }

  schedulePing() {
    if (this.pingTimeout) clearTimeout(this.pingTimeout);
    this.pingTimeout = setTimeout(() => {
      if (this.stopped) {
        this.pingTimeout = null;
        return;
      }
      if (this.socket && this.socket.readyState === WebSocket.OPEN) {
        this.sendAction("ping");
        this.schedulePing();
      } else {
        this.pingTimeout = null;
      }
    }, 30_000);
  }
}