import { injectable } from "@codemation/core";
import type { TelemetrySpanUpsert } from "../../domain/telemetry/TelemetryContracts";
import type { TelemetrySpanPublisher } from "./TelemetrySpanPublisher";
import type { SseSpanSink } from "./SseSpanSink.types";

@injectable()
export class TelemetrySpanSseFanout implements TelemetrySpanPublisher {
  private readonly sinks = new Map<string, Set<SseSpanSink>>();

  register(runId: string, sink: SseSpanSink): void {
    const set = this.sinks.get(runId) ?? new Set<SseSpanSink>();
    set.add(sink);
    this.sinks.set(runId, set);
  }

  unregister(runId: string, sink: SseSpanSink): void {
    const set = this.sinks.get(runId);
    if (!set) return;
    set.delete(sink);
    if (set.size === 0) this.sinks.delete(runId);
  }

  async publishSpan(span: TelemetrySpanUpsert): Promise<void> {
    const set = this.sinks.get(span.runId);
    if (!set) return;
    for (const sink of set) {
      sink.send(span);
    }
  }
}