import type { Logger, StorageNamespace } from "@dispatch/kernel";
import { aggregateSamples, type ModelThroughput, type ThroughputSample } from "./aggregate.js";
import { dayKeyOf, type Period, resolvePeriod } from "./period.js";

export interface ThroughputReport {
  readonly period: Period;
  readonly date: string;
  /** Inclusive start, epoch-ms. */
  readonly start: number;
  /** Exclusive end, epoch-ms. */
  readonly end: number;
  readonly models: readonly ModelThroughput[];
}

export interface ThroughputQuery {
  readonly period: Period;
  readonly date: string;
}

/**
 * Cross-conversation, per-model throughput store. Records one sample per
 * completed turn and aggregates token-weighted tok/s over a day/week/month.
 * Persistence is the injected KV namespace; samples are bucketed by local day
 * key so a period query addresses only its own day buckets.
 */
export interface ThroughputStore {
  readonly record: (sample: ThroughputSample) => Promise<void>;
  readonly aggregate: (query: ThroughputQuery) => Promise<ThroughputReport>;
}

export interface ThroughputStoreDeps {
  readonly storage: StorageNamespace;
  readonly logger?: Logger;
  /** Injectable unique-id generator (default crypto.randomUUID). */
  readonly newId?: () => string;
}

/** Thrown when a query's `(period, date)` is malformed. */
export class ThroughputQueryError extends Error {}

const SAMPLE_PREFIX = "sample";

export function createThroughputStore(deps: ThroughputStoreDeps): ThroughputStore {
  const newId = deps.newId ?? (() => crypto.randomUUID());

  return {
    async record(sample) {
      // Per-sample key under the local-day bucket → write is a single set
      // (no read-modify-write, so concurrent turns can't lose a sample).
      const day = dayKeyOf(sample.ts);
      const key = `${SAMPLE_PREFIX}:${day}:${sample.ts}:${newId()}`;
      await deps.storage.set(key, JSON.stringify(sample));
    },

    async aggregate(query) {
      const resolved = resolvePeriod(query.period, query.date);
      if (!resolved.ok) throw new ThroughputQueryError(resolved.error);

      const samples: ThroughputSample[] = [];
      for (const day of resolved.dayKeys) {
        const keys = await deps.storage.keys(`${SAMPLE_PREFIX}:${day}:`);
        for (const k of keys) {
          const raw = await deps.storage.get(k);
          if (raw === null) continue;
          try {
            samples.push(JSON.parse(raw) as ThroughputSample);
          } catch {
            // Skip a malformed row rather than failing the whole query.
          }
        }
      }

      const models = aggregateSamples(samples, resolved.start, resolved.end);
      return {
        period: query.period,
        date: resolved.date,
        start: resolved.start,
        end: resolved.end,
        models,
      };
    },
  };
}
