import { IMetric, isIMetric, isMetricValue, ITrace, ITraceWith, MetricValue, SideEffect } from '@emprespresso/pengueno';

export type MetricsTraceSupplier =
    | ITraceWith<IMetric | MetricValue | undefined>
    | Array<ITraceWith<IMetric | MetricValue | undefined>>;
export const isMetricsTraceSupplier = (t: unknown): t is MetricsTraceSupplier =>
    isMetricValue(t) || isIMetric(t) || (Array.isArray(t) && t.every((_m) => isMetricValue(_m) || isIMetric(_m)));

export class MetricsTrace implements ITrace<MetricsTraceSupplier> {
    constructor(
        private readonly metricConsumer: SideEffect<Array<MetricValue>>,
        private readonly activeTraces: ReadonlyMap<IMetric, number> = new Map(),
        private readonly completedTraces: ReadonlySet<IMetric> = new Set(),
    ) {}

    public traceScope(trace: MetricsTraceSupplier): MetricsTrace {
        const now = Date.now();
        const metricsToTrace = (Array.isArray(trace) ? trace : [trace]).filter(isIMetric);

        const initialTraces = new Map(metricsToTrace.map((metric) => [metric, now]));

        return new MetricsTrace(this.metricConsumer, initialTraces);
    }

    public trace(metrics: MetricsTraceSupplier): MetricsTrace {
        if (!metrics || typeof metrics === 'string') {
            return this;
        }

        const now = Date.now();
        const allMetrics = Array.isArray(metrics) ? metrics : [metrics];

        // partition the incoming metrics
        const valuesToEmit = allMetrics.filter(isMetricValue);
        const traceableMetrics = allMetrics.filter(isIMetric);

        const metricsToStart = traceableMetrics.filter((m) => !this.activeTraces.has(m));
        const metricsToEnd = traceableMetrics.filter((m) => this.activeTraces.has(m) && !this.completedTraces.has(m));

        // the new metrics to emit based on traces ending *now*
        const endedMetricValues = metricsToEnd.flatMap((metric) => [
            metric.count.withValue(1.0),
            metric.time.withValue(now - this.activeTraces.get(metric)!),
        ]);

        const allMetricsToEmit = [...valuesToEmit, ...endedMetricValues];
        if (allMetricsToEmit.length > 0) {
            this.metricConsumer(allMetricsToEmit);
        }

        // the next immutable state
        const nextActiveTraces = new Map([
            ...this.activeTraces,
            ...metricsToStart.map((m): [IMetric, number] => [m, now]),
        ]);
        const nextCompletedTraces = new Set([...this.completedTraces, ...metricsToEnd]);
        return new MetricsTrace(this.metricConsumer, nextActiveTraces, nextCompletedTraces);
    }
}