import { type LlamaSwapConfig } from "./util";
import { parseSse, type SseMessage } from "./sse";
import { type FeedEvent, type InflightRequest, type ModelState } from "./inflight-tracker";

export function decodeEvent(msg: SseMessage): FeedEvent | null {
  try {
    const outer = JSON.parse(msg.data) as { type?: string; data?: string };
    if (!outer.data) return null;
    const inner = JSON.parse(outer.data) as Record<string, unknown>;

    if (outer.type === "activity") {
      const id = typeof inner.id === "number" ? inner.id : Number.NaN;
      if (Number.isFinite(id)) return { type: "activity", id };
      return null;
    }

    if (outer.type === "inflight") {
      const operation = inner.operation as string;
      if (operation === "remove") {
        const id = typeof inner.id === "string" ? inner.id : undefined;
        if (id) return { type: "inflight", operation: "remove", id };
        return null;
      }
      if (operation !== "snapshot" && operation !== "upsert" && operation !== "add") return null;
      const requests: InflightRequest[] = Array.isArray(inner.requests)
        ? (inner.requests as InflightRequest[])
        : Array.isArray(inner.request)
          ? (inner.request as InflightRequest[])
          : inner.request && typeof inner.request === "object"
            ? [inner.request as InflightRequest]
            : [];
      return {
        type: "inflight",
        operation: operation === "snapshot" ? "snapshot" : "add",
        requests,
      };
    }

    if (outer.type === "modelStatus") {
      if (!Array.isArray(inner)) return null;
      return { type: "modelStatus", models: inner as unknown as ModelState[] };
    }

    return null;
  } catch {
    return null;
  }
}

const MAX_RETRY_MS = 30000;
const INITIAL_RETRY_MS = 1000;

export class EventFeed {
  private controller?: AbortController;
  private closed = false;
  private retryMs = INITIAL_RETRY_MS;
  private statusHandler?: (connected: boolean) => void;

  constructor(
    private cfg: LlamaSwapConfig,
    private onEvent: (ev: FeedEvent) => void,
  ) {}

  setStatusHandler(fn: (connected: boolean) => void): void {
    this.statusHandler = fn;
  }

  start(): void {
    void this.connect();
  }

  stop(): void {
    this.closed = true;
    this.controller?.abort();
  }

  private async connect(): Promise<void> {
    while (!this.closed) {
      this.controller = new AbortController();
      try {
        const headers: Record<string, string> = { Accept: "text/event-stream" };
        if (this.cfg.apiKey) headers["Authorization"] = `Bearer ${this.cfg.apiKey}`;
        const res = await fetch(`${this.cfg.baseUrl}/api/events`, {
          headers,
          signal: this.controller.signal,
        });
        if (!res.ok || !res.body) throw new Error(`events HTTP ${res.status}`);
        this.statusHandler?.(true);
        this.retryMs = INITIAL_RETRY_MS;
        await this.readStream(res.body.getReader());
      } catch {
        if (this.closed) return;
      }
      this.statusHandler?.(false);
      if (this.closed) return;
      await sleep(this.retryMs);
      this.retryMs = Math.min(this.retryMs * 2, MAX_RETRY_MS);
    }
  }

  private async readStream(reader: ReadableStreamDefaultReader<Uint8Array>): Promise<void> {
    const decoder = new TextDecoder();
    let buffer = "";
    try {
      while (!this.closed) {
        const { done, value } = await reader.read();
        if (done) break;
        buffer += decoder.decode(value, { stream: true });
        const { messages, rest } = parseSse(buffer);
        buffer = rest;
        for (const msg of messages) {
          const ev = decodeEvent(msg);
          if (ev) this.onEvent(ev);
        }
      }
    } finally {
      reader.releaseLock();
    }
  }
}

function sleep(ms: number): Promise<void> {
  return new Promise((resolve) => setTimeout(resolve, ms));
}