import type { Context, Middleware } from "./context.ts";

export interface SpanRecord {
  name: string;
  startTime: number;
  endTime?: number;
  durationMs?: number;
  status?: "ok" | "error";
  attributes: Record<string, string | number | boolean>;
  error?: unknown;
}

export interface Tracer {
  startSpan(name: string, attributes?: SpanRecord["attributes"]): Span;
  records(): readonly SpanRecord[];
}

export interface Span {
  setAttribute(name: string, value: string | number | boolean): void;
  end(status?: "ok" | "error", error?: unknown): SpanRecord;
}

export function createTracer(clock: () => number = () => performance.now()): Tracer {
  const spans: SpanRecord[] = [];
  return {
    startSpan(name, attributes = {}) {
      const record: SpanRecord = { name, startTime: clock(), attributes: { ...attributes } };
      spans.push(record);
      let ended = false;
      return {
        setAttribute(key, value) {
          record.attributes[key] = value;
        },
        end(status = "ok", error) {
          if (!ended) {
            ended = true;
            record.endTime = clock();
            record.durationMs = record.endTime - record.startTime;
            record.status = status;
            record.error = error;
          }
          return record;
        },
      };
    },
    records: () => spans,
  };
}

export async function withSpan<T>(
  tracer: Tracer,
  name: string,
  run: (span: Span) => T | Promise<T>,
  attributes?: SpanRecord["attributes"],
): Promise<T> {
  const span = tracer.startSpan(name, attributes);
  try {
    const result = await run(span);
    span.end("ok");
    return result;
  } catch (error) {
    span.end("error", error);
    throw error;
  }
}

export interface TracingMiddlewareOptions {
  /** Include W3C Server-Timing response headers. Defaults to true. */
  serverTiming?: boolean;
  /** Fraction of requests to trace, from 0 to 1. Defaults to 1. */
  sampleRate?: number;
  /** Called after a traced response completes. */
  onComplete?: (ctx: Context, records: readonly SpanRecord[]) => void | Promise<void>;
}

export function tracingMiddleware(
  tracerFactory: (ctx: Context) => Tracer = () => createTracer(),
  options: TracingMiddlewareOptions = {},
): Middleware {
  const sampleRate = Math.max(0, Math.min(1, options.sampleRate ?? 1));
  return async (ctx, next) => {
    if (sampleRate === 0 || (sampleRate < 1 && Math.random() > sampleRate)) return next();

    const tracer = tracerFactory(ctx);
    ctx.tracer = tracer;
    const response = await withSpan(tracer, "http.request", () => next(), {
      method: ctx.req.method,
      path: ctx.url.pathname,
    });
    const records = tracer.records();
    await options.onComplete?.(ctx, records);

    if (options.serverTiming === false) return response;

    const headers = new Headers(response.headers);
    const timings = records
      .filter((record) => record.durationMs !== undefined)
      .map(
        (record, index) =>
          `wrn${index};dur=${record.durationMs!.toFixed(2)};desc="${record.name.replace(/"/g, "")}"`,
      );
    if (timings.length) headers.set("server-timing", timings.join(", "));
    return new Response(response.body, {
      status: response.status,
      statusText: response.statusText,
      headers,
    });
  };
}