import { EventFeed } from "./event-feed";
import { InflightTracker } from "./inflight-tracker";
import { MetricsPoller } from "./metrics-poller";
import { StatsCache } from "./stats-cache";
import { type UsageStats } from "./stats";
import { type LlamaSwapConfig } from "./util";

class Runtime {
  readonly tracker = new InflightTracker();
  offline = true;
  private cfg?: LlamaSwapConfig;
  private feed?: EventFeed;
  private poller?: MetricsPoller;
  private statsCache?: StatsCache;
  private listeners = new Set<() => void>();

  ensureConnections(cfg: LlamaSwapConfig): void {
    const changed = !this.cfg || this.cfg.baseUrl !== cfg.baseUrl || this.cfg.apiKey !== cfg.apiKey;
    if (changed) {
      this.feed?.stop();
      this.poller?.stop();
      this.feed = undefined;
      this.poller = undefined;
      this.statsCache?.setConfig(cfg);
      this.cfg = cfg;
    }
    if (!this.feed) {
      this.feed = new EventFeed(cfg, (ev) => {
        this.tracker.apply(ev);
        if (ev.type === "activity") this.statsCache?.scheduleRefresh();
        this.emit();
      });
      this.feed.setStatusHandler((connected) => {
        this.offline = !connected;
        this.emit();
      });
      this.feed.start();
    }
    if (!this.poller) {
      this.poller = new MetricsPoller(cfg);
      this.poller.on(() => this.emit());
      this.poller.start();
    }
    if (!this.statsCache) {
      this.statsCache = new StatsCache(cfg);
      this.statsCache.onChange(() => this.emit());
    }
  }

  get pollerInstance(): MetricsPoller | undefined {
    return this.poller;
  }

  getStats(modelId: string): UsageStats | undefined {
    return this.statsCache?.get(modelId);
  }

  watchStats(modelId: string): () => void {
    this.statsCache?.register(modelId);
    return () => this.statsCache?.unregister(modelId);
  }

  subscribe(listener: () => void): () => void {
    this.listeners.add(listener);
    return () => {
      this.listeners.delete(listener);
    };
  }

  private emit(): void {
    for (const listener of this.listeners) listener();
  }
}

export const runtime = new Runtime();