import type { TimelineItem } from "./streaming-v2";

export interface Metrics {
  context: { consumed: number; budget: number } | null;
  cached: { cached: number; prompt: number; ratio: number } | null;
  tokPerSec: number | null;
  reported: boolean;
  route: string | null;
}

export function readMetrics(items: TimelineItem[]): Metrics {
  if (!Array.isArray(items) || items.length === 0) {
    return {
      context: null,
      cached: null,
      tokPerSec: null,
      reported: true,
      route: null,
    };
  }

  let newestReceipt: { consumed: number; budget: number } | null = null;
  const usageRecords: Array<{
    tokensReported: boolean;
    routeModel: string | null;
    promptTokens?: number;
    cachedPromptTokens?: number;
    completionTokens?: number;
    generationMs?: number;
  }> = [];

  for (const item of items) {
    if (item.lane !== "lifecycle") continue;
    const entry = item.data as any;
    if (!entry || typeof entry !== "object") continue;

    const kind = entry.kind || entry.type;
    const d = (entry.data && typeof entry.data === "object" ? entry.data : entry) as Record<string, any>;
    const innerKind = d.kind || d.type;

    // 1. Context receipt: check data.consumed_tokens, data.context_receipt.consumed_tokens, data.first_change.consumed_tokens
    if (
      kind === "context.receipt" ||
      innerKind === "context_receipt" ||
      innerKind === "context.receipt" ||
      d.context_receipt !== undefined
    ) {
      const cr = d.context_receipt;
      const fc = d.first_change;

      let consumed: number | undefined;
      if (typeof d.consumed_tokens === "number") consumed = d.consumed_tokens;
      else if (typeof cr?.consumed_tokens === "number") consumed = cr.consumed_tokens;
      else if (typeof fc?.consumed_tokens === "number") consumed = fc.consumed_tokens;

      let budget: number | undefined;
      if (typeof d.total_budget_tokens === "number") budget = d.total_budget_tokens;
      else if (typeof cr?.total_budget_tokens === "number") budget = cr.total_budget_tokens;
      else if (typeof fc?.total_budget_tokens === "number") budget = fc.total_budget_tokens;
      else if (typeof d.budget_tokens === "number") budget = d.budget_tokens;
      else if (typeof cr?.budget_tokens === "number") budget = cr.budget_tokens;

      if (
        typeof consumed === "number" &&
        typeof budget === "number" &&
        !Number.isNaN(consumed) &&
        !Number.isNaN(budget)
      ) {
        newestReceipt = { consumed, budget };
      }
    }

    // 2. Model usage: request_usage or model.usage payload
    if (
      kind === "model.usage" ||
      innerKind === "request_usage" ||
      innerKind === "model.usage" ||
      d.request_sha256 !== undefined ||
      d.usage !== undefined
    ) {
      const usage = (d.usage && typeof d.usage === "object" ? d.usage : d) as Record<string, any>;

      let tokensReported = true;
      if (
        usage.tokens_reported === false ||
        d.tokens_reported === false ||
        d.usage_reported === false
      ) {
        tokensReported = false;
      }

      let routeModel: string | null = null;
      const routeObj = d.route ?? usage.route;
      if (routeObj && typeof routeObj === "object" && typeof routeObj.model === "string" && routeObj.model.trim()) {
        routeModel = routeObj.model.trim();
      } else if (typeof routeObj === "string" && routeObj.trim()) {
        routeModel = routeObj.trim();
      }

      const promptTokens = typeof usage.prompt_tokens === "number" ? usage.prompt_tokens : undefined;
      const cachedPromptTokens =
        typeof usage.cached_prompt_tokens === "number"
          ? usage.cached_prompt_tokens
          : typeof usage.prompt_cache_hit_tokens === "number"
          ? usage.prompt_cache_hit_tokens
          : undefined;
      const completionTokens = typeof usage.completion_tokens === "number" ? usage.completion_tokens : undefined;
      const generationMs = typeof usage.generation_ms === "number" ? usage.generation_ms : undefined;

      usageRecords.push({
        tokensReported,
        routeModel,
        promptTokens,
        cachedPromptTokens,
        completionTokens,
        generationMs,
      });
    }
  }

  // Session-scoped: sum the usage records of the session for the cache figures
  let sumPrompt = 0;
  let sumCached = 0;
  let anyPromptReported = false;
  let anyCacheReported = false;

  for (const rec of usageRecords) {
    if (!rec.tokensReported) continue;
    if (typeof rec.promptTokens === "number") {
      sumPrompt += rec.promptTokens;
      anyPromptReported = true;
    }
    if (typeof rec.cachedPromptTokens === "number") {
      sumCached += rec.cachedPromptTokens;
      anyCacheReported = true;
    }
  }

  const cached =
    anyCacheReported && anyPromptReported && sumPrompt > 0
      ? { cached: sumCached, prompt: sumPrompt, ratio: sumCached / sumPrompt }
      : null;

  // Newest usage for tok/s and route; reported: false when newest usage tokens_reported === false
  let tokPerSec: number | null = null;
  let reported = true;
  let route: string | null = null;

  if (usageRecords.length > 0) {
    const newest = usageRecords[usageRecords.length - 1];
    reported = newest.tokensReported;
    route = newest.routeModel;

    if (
      reported &&
      typeof newest.completionTokens === "number" &&
      typeof newest.generationMs === "number" &&
      newest.generationMs > 0
    ) {
      tokPerSec = (newest.completionTokens / newest.generationMs) * 1000;
    }
  }

  return {
    context: newestReceipt,
    cached,
    tokPerSec,
    reported,
    route,
  };
}