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<string, number> = new Map(),
        private readonly completedTraces: ReadonlySet<string> = 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.name, now]));

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

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

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

        const valuesToEmit = allMetrics.filter(isMetricValue);
        const traceableMetrics = allMetrics.filter(isIMetric);

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

        const endedMetricValues = metricsToEnd.flatMap((metric) => [
            metric.count.withValue(1.0),
            metric.time.withValue(now - this.activeTraces.get(metric.name)!),
        ]);

        const parentBasedMetrics = metricsToStart.filter((metric) => {
            const parent = metric.parent;
            return parent && this.activeTraces.has(parent.name);
        });

        const parentBasedValues = parentBasedMetrics.flatMap((metric) => {
            const parentStart = this.activeTraces.get(metric.parent!.name)!;
            return [
                metric.count.withValue(1.0),
                metric.time.withValue(now - parentStart),
            ];
        });

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

        const nextActiveTraces = new Map([
            ...this.activeTraces,
            ...metricsToStart.map((m): [string, number] => [m.name, now]),
        ]);
        const nextCompletedTraces = new Set([
            ...this.completedTraces,
            ...metricsToEnd.map((m) => m.name),
            ...parentBasedMetrics.map((m) => m.name),
        ]);

        return new MetricsTrace(this.metricConsumer, nextActiveTraces, nextCompletedTraces);
    }
}