import type { PlayExecutionGovernor, WorkLease } from './governor/governor';

export type RuntimeResourceLease = WorkLease;

export type RuntimeResourceObservation = {
  toolId?: string | null;
  providerResourceKey?: string | null;
  rowEstimatedBytes?: number | null;
  rowAdmissionWaitMs?: number | null;
  toolAdmissionWaitMs?: number | null;
  providerLatencyMs?: number | null;
  providerSuccess?: boolean;
  provider429?: boolean;
  retryAfterMs?: number | null;
  receiptLatencyMs?: number | null;
  sheetFlushLatencyMs?: number | null;
  sheetFlushBytes?: number | null;
  memoryBytes?: number | null;
};

export type RuntimeResourceSnapshot = {
  rows: {
    acquired: number;
    estimatedBytes: number;
    admissionWaitMsEwma: number | null;
  };
  tools: {
    acquired: number;
    admissionWaitMsEwma: number | null;
  };
  sheetFlush: {
    acquired: number;
    bytes: number;
    latencyMsEwma: number | null;
  };
  providers: Record<
    string,
    {
      latencyMsEwma: number | null;
      retryAfterMsEwma: number | null;
      provider429Count: number;
    }
  >;
};

export interface RuntimeResourceGovernor {
  acquireRow(input?: {
    estimatedBytes?: number | null;
    signal?: AbortSignal;
  }): Promise<RuntimeResourceLease>;
  acquireTool(input: {
    orgId?: string | null;
    providerResourceKey?: string | null;
    toolId: string;
    estimatedCost?: number | null;
    signal?: AbortSignal;
  }): Promise<RuntimeResourceLease>;
  acquireProviderPermit(input: {
    toolId: string;
    signal?: AbortSignal;
  }): Promise<RuntimeResourceLease>;
  suggestedToolParallelism(toolId: string, fallback: number): Promise<number>;
  resolveRowConcurrency(requested?: number): number;
  reportProviderBackpressure(input: {
    provider: string;
    toolId?: string;
    retryAfterMs: number;
  }): Promise<void>;
  observe(input: RuntimeResourceObservation): void;
  snapshot(): RuntimeResourceSnapshot;
}

type Ewma = {
  value: number | null;
  observe(value: number | null | undefined): void;
};

function createEwma(alpha = 0.2): Ewma {
  let value: number | null = null;
  return {
    get value() {
      return value;
    },
    observe(next) {
      if (next == null || !Number.isFinite(next)) return;
      const normalized = Math.max(0, next);
      value =
        value == null ? normalized : value * (1 - alpha) + normalized * alpha;
    },
  };
}

function elapsedSince(startedAt: number): number {
  return Math.max(0, Date.now() - startedAt);
}

function resourceKey(input?: string | null): string {
  const trimmed = input?.trim();
  return trimmed || 'unknown';
}

export function createRuntimeResourceGovernor(input: {
  executionGovernor: PlayExecutionGovernor;
}): RuntimeResourceGovernor {
  const rowAdmissionWait = createEwma();
  const toolAdmissionWait = createEwma();
  const sheetFlushLatency = createEwma();
  const providers = new Map<
    string,
    {
      latency: Ewma;
      retryAfter: Ewma;
      provider429Count: number;
    }
  >();
  const rowCounters = {
    acquired: 0,
    estimatedBytes: 0,
  };
  const toolCounters = {
    acquired: 0,
  };
  const sheetFlushCounters = {
    acquired: 0,
    bytes: 0,
  };

  const providerStats = (key: string) => {
    const normalized = resourceKey(key);
    let stats = providers.get(normalized);
    if (!stats) {
      stats = {
        latency: createEwma(),
        retryAfter: createEwma(),
        provider429Count: 0,
      };
      providers.set(normalized, stats);
    }
    return stats;
  };

  const observe = (observation: RuntimeResourceObservation): void => {
    rowAdmissionWait.observe(observation.rowAdmissionWaitMs);
    toolAdmissionWait.observe(observation.toolAdmissionWaitMs);
    sheetFlushLatency.observe(observation.sheetFlushLatencyMs);
    if (observation.sheetFlushLatencyMs != null) {
      sheetFlushCounters.acquired += 1;
    }
    if (observation.rowEstimatedBytes != null) {
      rowCounters.estimatedBytes += Math.max(
        0,
        Math.floor(observation.rowEstimatedBytes),
      );
    }
    if (observation.sheetFlushBytes != null) {
      sheetFlushCounters.bytes += Math.max(
        0,
        Math.floor(observation.sheetFlushBytes),
      );
    }
    if (
      observation.providerResourceKey ||
      observation.providerLatencyMs != null ||
      observation.provider429 ||
      observation.retryAfterMs != null
    ) {
      const stats = providerStats(observation.providerResourceKey ?? 'unknown');
      stats.latency.observe(observation.providerLatencyMs);
      stats.retryAfter.observe(observation.retryAfterMs);
      if (observation.provider429) {
        stats.provider429Count += 1;
      }
      if (
        observation.providerSuccess === true &&
        (observation.toolId || observation.providerResourceKey)
      ) {
        if (observation.toolId) {
          input.executionGovernor.observeToolSuccess({
            toolId: observation.toolId,
            latencyMs: observation.providerLatencyMs,
          });
        } else {
          input.executionGovernor.observeProviderSuccess({
            provider: observation.providerResourceKey!,
            latencyMs: observation.providerLatencyMs,
          });
        }
      }
    }
  };

  return {
    async acquireRow(rowInput) {
      const startedAt = Date.now();
      const lease = await input.executionGovernor.acquireRowSlot({
        signal: rowInput?.signal,
      });
      rowCounters.acquired += 1;
      observe({
        rowAdmissionWaitMs: elapsedSince(startedAt),
        rowEstimatedBytes: rowInput?.estimatedBytes ?? null,
      });
      return lease;
    },

    async acquireTool(toolInput) {
      const startedAt = Date.now();
      const lease = await input.executionGovernor.acquireToolSlot(
        toolInput.toolId,
        { signal: toolInput.signal },
      );
      toolCounters.acquired += 1;
      observe({
        providerResourceKey: toolInput.providerResourceKey ?? toolInput.toolId,
        toolAdmissionWaitMs: elapsedSince(startedAt),
      });
      return lease;
    },

    acquireProviderPermit(providerInput) {
      return input.executionGovernor.acquireProviderPermit(
        providerInput.toolId,
        { signal: providerInput.signal },
      );
    },

    suggestedToolParallelism(toolId, fallback) {
      return input.executionGovernor.suggestedParallelism(toolId, fallback);
    },

    resolveRowConcurrency(requested) {
      return input.executionGovernor.resolveRowConcurrency(requested);
    },

    async reportProviderBackpressure(backpressure) {
      await input.executionGovernor.reportProviderBackpressure(backpressure);
      observe({
        providerResourceKey: backpressure.provider,
        provider429: true,
        retryAfterMs: backpressure.retryAfterMs,
      });
    },

    observe,

    snapshot() {
      return {
        rows: {
          acquired: rowCounters.acquired,
          estimatedBytes: rowCounters.estimatedBytes,
          admissionWaitMsEwma: rowAdmissionWait.value,
        },
        tools: {
          acquired: toolCounters.acquired,
          admissionWaitMsEwma: toolAdmissionWait.value,
        },
        sheetFlush: {
          acquired: sheetFlushCounters.acquired,
          bytes: sheetFlushCounters.bytes,
          latencyMsEwma: sheetFlushLatency.value,
        },
        providers: Object.fromEntries(
          [...providers.entries()].map(([key, stats]) => [
            key,
            {
              latencyMsEwma: stats.latency.value,
              retryAfterMsEwma: stats.retryAfter.value,
              provider429Count: stats.provider429Count,
            },
          ]),
        ),
      };
    },
  };
}