import type { GenerationStats } from "./types";
import { TURN_SUMMARY_BADGE, TURN_SUMMARY_SEPARATOR } from "./constants";
import { ShartiLogger } from "./utils/logger";

/* -------------------------------------------------------------------------- */
/*                          Statistics reporting                             */
/* -------------------------------------------------------------------------- */

/** Emit a one-line per-request performance summary (when any metric is present). */
export function logGenerationStats(logger: ShartiLogger, stats: GenerationStats): void {
  const parts: string[] = [];

  if (stats.inputTokens != null) parts.push(`📥 Input: ${stats.inputTokens} tokens`);
  if (stats.cachedTokens != null) parts.push(`🗄️ Cached: ${stats.cachedTokens} tokens`);
  if (stats.outputTokens != null) parts.push(`📤 Output: ${stats.outputTokens} tokens`);
  if (stats.promptPerSecond != null) parts.push(`⚡️ Prefill: ${stats.promptPerSecond.toFixed(1)} tok/s`);
  if (stats.timeToFirstTokenSeconds != null) parts.push(`🕐 TTFT: ${(stats.timeToFirstTokenSeconds * 1000).toFixed(0)}ms`);
  if (stats.tokensPerSecond != null) parts.push(`♻️ Speed: ${stats.tokensPerSecond.toFixed(1)} tok/s`);
  if (stats.totalTimeSeconds != null) parts.push(`✅ Done in: ${stats.totalTimeSeconds.toFixed(2)}s`);

  if (parts.length > 0) {
    logger.info(`📊 Stats: ${parts.join(" | ")}`);
  }
}

/** Raw accumulators for a single conversational turn. */
interface TurnStats {
  calls: number;
  totalInputTokens: number;
  totalCachedTokens: number;
  totalOutputTokens: number;
  totalTimeMs: number;
  totalTTFTMs: number;
  lastInputTokens: number;
  lastOutputTokens: number;
  lastHistoryLength: number;
}

/**
 * Aggregate turn statistics for a single conversational turn.
 *
 * Instances are SCOPED TO A CONVERSATION THREAD and fetched via {@link getSessionAccumulator},
 * so the same accumulator survives the multi-round (tool handoff) regenerations of one turn.
 * Distinct conversations hold distinct instances, which keeps concurrent generations isolated.
 */
export class TurnStatAccumulator {
  private acc: TurnStats | null = null;

  /** Fold one finished request's metrics into the running turn totals. */
  record(stats: GenerationStats, historyLength: number): void {
    // A shorter history implies a brand-new conversation: restart aggregation.
    if (this.acc && historyLength < this.acc.lastHistoryLength) {
      this.acc = null;
    }

    if (!this.acc) {
      this.acc = {
        calls: 0,
        totalInputTokens: 0,
        totalCachedTokens: 0,
        totalOutputTokens: 0,
        totalTimeMs: 0,
        totalTTFTMs: 0,
        lastInputTokens: 0,
        lastOutputTokens: 0,
        lastHistoryLength: historyLength,
      };
    }

    const a = this.acc;
    a.calls++;
    a.totalInputTokens += stats.inputTokens ?? 0;
    a.totalCachedTokens += stats.cachedTokens ?? 0;
    a.totalOutputTokens += stats.outputTokens ?? 0;
    a.totalTimeMs += (stats.totalTimeSeconds ?? 0) * 1000;
    a.totalTTFTMs += (stats.timeToFirstTokenSeconds ?? 0) * 1000;
    a.lastInputTokens = stats.inputTokens ?? 0;
    a.lastOutputTokens = stats.outputTokens ?? 0;
    a.lastHistoryLength = historyLength;
  }

  /**
   * If the turn appears to be concluding (not handing off to a tool call), emit an aggregated
   * summary to the console and, optionally, to the chat, then reset.
   */
  emitSummaryIfFinal(logger: ShartiLogger, finishReason: string | null, showInChat: boolean): void {
    const isFinal = finishReason !== "tool_calls";
    const a = this.acc;
    if (!isFinal || !a || a.calls === 0) return;

    const totalTimeSec = a.totalTimeMs / 1000;
    const avgTTFTms = Math.round(a.totalTTFTMs / a.calls);
    const decodeTimeMs = a.totalTimeMs - a.totalTTFTMs;
    const decodeSpeed = decodeTimeMs > 0 ? a.totalOutputTokens / (decodeTimeMs / 1000) : 0;

    const body = `Calls: ${a.calls} | Tokens: ${a.totalInputTokens} in / ${a.totalOutputTokens} out | TTFT: ${avgTTFTms}ms | Speed: ${decodeSpeed.toFixed(1)} tps | Time: ${totalTimeSec.toFixed(2)}s`;

    logger.notice(TURN_SUMMARY_BADGE + body);
    if (showInChat) {
      // Emits exactly TURN_SUMMARY_MARKER (SEPARATOR + BADGE) so ingestion-side stripping agrees.
      logger.notifyUser(TURN_SUMMARY_SEPARATOR + TURN_SUMMARY_BADGE + body);
    }
    this.reset();
  }

  /** Discard accumulated state. */
  reset(): void {
    this.acc = null;
  }
}

/* -------------------------------------------------------------------------- */
/*                       Cross-invocation session store                       */
/*                                                                            */
/* One accumulator per conversation thread, held in module scope so the        */
/* multi-round (tool handoff) accounting survives across separate generate()   */
/* invocations. Bounded by a small LRU so a long-lived process cannot leak.    */
/* -------------------------------------------------------------------------- */

const MAX_TRACKED_SESSIONS = 32;
const sessionStore = new Map<string, TurnStatAccumulator>();

/** Fetch (creating on miss) the accumulator for a conversation-thread key. */
export function getSessionAccumulator(key: string): TurnStatAccumulator {
  const existing = sessionStore.get(key);
  if (existing) {
    // Touch → move to most-recently-used end of the ordered map.
    sessionStore.delete(key);
    sessionStore.set(key, existing);
    return existing;
  }

  if (sessionStore.size >= MAX_TRACKED_SESSIONS) {
    const oldest = sessionStore.keys().next().value;
    if (oldest !== undefined) sessionStore.delete(oldest);
  }

  const created = new TurnStatAccumulator();
  sessionStore.set(key, created);
  return created;
}

import type { GenerationStats } from "./types";
import { TURN_SUMMARY_BADGE, TURN_SUMMARY_SEPARATOR } from "./constants";
import { ShartiLogger } from "./utils/logger";

/* -------------------------------------------------------------------------- */
/*                          Statistics reporting                             */
/* -------------------------------------------------------------------------- */

/** Emit a one-line per-request performance summary (when any metric is present). */
export function logGenerationStats(logger: ShartiLogger, stats: GenerationStats): void {
  const parts: string[] = [];

  if (stats.inputTokens != null) parts.push(`📥 Input: ${stats.inputTokens} tokens`);
  if (stats.cachedTokens != null) parts.push(`🗄️ Cached: ${stats.cachedTokens} tokens`);
  if (stats.outputTokens != null) parts.push(`📤 Output: ${stats.outputTokens} tokens`);
  if (stats.promptPerSecond != null) parts.push(`⚡️ Prefill: ${stats.promptPerSecond.toFixed(1)} tok/s`);
  if (stats.timeToFirstTokenSeconds != null) parts.push(`🕐 TTFT: ${(stats.timeToFirstTokenSeconds * 1000).toFixed(0)}ms`);
  if (stats.tokensPerSecond != null) parts.push(`♻️ Speed: ${stats.tokensPerSecond.toFixed(1)} tok/s`);
  if (stats.totalTimeSeconds != null) parts.push(`✅ Done in: ${stats.totalTimeSeconds.toFixed(2)}s`);

  if (parts.length > 0) {
    logger.info(`📊 Stats: ${parts.join(" | ")}`);
  }
}

/** Raw accumulators for a single conversational turn. */
interface TurnStats {
  calls: number;
  totalInputTokens: number;
  totalCachedTokens: number;
  totalOutputTokens: number;
  totalTimeMs: number;
  totalTTFTMs: number;
  lastInputTokens: number;
  lastOutputTokens: number;
  lastHistoryLength: number;
}

/**
 * Aggregate turn statistics for a single conversational turn.
 *
 * Instances are SCOPED TO A CONVERSATION THREAD and fetched via {@link getSessionAccumulator},
 * so the same accumulator survives the multi-round (tool handoff) regenerations of one turn.
 * Distinct conversations hold distinct instances, which keeps concurrent generations isolated.
 */
export class TurnStatAccumulator {
  private acc: TurnStats | null = null;

  /** Fold one finished request's metrics into the running turn totals. */
  record(stats: GenerationStats, historyLength: number): void {
    // A shorter history implies a brand-new conversation: restart aggregation.
    if (this.acc && historyLength < this.acc.lastHistoryLength) {
      this.acc = null;
    }

    if (!this.acc) {
      this.acc = {
        calls: 0,
        totalInputTokens: 0,
        totalCachedTokens: 0,
        totalOutputTokens: 0,
        totalTimeMs: 0,
        totalTTFTMs: 0,
        lastInputTokens: 0,
        lastOutputTokens: 0,
        lastHistoryLength: historyLength,
      };
    }

    const a = this.acc;
    a.calls++;
    a.totalInputTokens += stats.inputTokens ?? 0;
    a.totalCachedTokens += stats.cachedTokens ?? 0;
    a.totalOutputTokens += stats.outputTokens ?? 0;
    a.totalTimeMs += (stats.totalTimeSeconds ?? 0) * 1000;
    a.totalTTFTMs += (stats.timeToFirstTokenSeconds ?? 0) * 1000;
    a.lastInputTokens = stats.inputTokens ?? 0;
    a.lastOutputTokens = stats.outputTokens ?? 0;
    a.lastHistoryLength = historyLength;
  }

  /**
   * If the turn appears to be concluding (not handing off to a tool call), emit an aggregated
   * summary to the console and, optionally, to the chat, then reset.
   */
  emitSummaryIfFinal(logger: ShartiLogger, finishReason: string | null, showInChat: boolean): void {
    const isFinal = finishReason !== "tool_calls";
    const a = this.acc;
    if (!isFinal || !a || a.calls === 0) return;

    const totalTimeSec = a.totalTimeMs / 1000;
    const avgTTFTms = Math.round(a.totalTTFTMs / a.calls);
    const decodeTimeMs = a.totalTimeMs - a.totalTTFTMs;
    const decodeSpeed = decodeTimeMs > 0 ? a.totalOutputTokens / (decodeTimeMs / 1000) : 0;

    const body = `Calls: ${a.calls} | Tokens: ${a.totalInputTokens} in / ${a.totalOutputTokens} out | TTFT: ${avgTTFTms}ms | Speed: ${decodeSpeed.toFixed(1)} tps | Time: ${totalTimeSec.toFixed(2)}s`;

    logger.notice(TURN_SUMMARY_BADGE + body);
    if (showInChat) {
      // Emits exactly TURN_SUMMARY_MARKER (SEPARATOR + BADGE) so ingestion-side stripping agrees.
      logger.notifyUser(TURN_SUMMARY_SEPARATOR + TURN_SUMMARY_BADGE + body);
    }
    this.reset();
  }

  /** Discard accumulated state. */
  reset(): void {
    this.acc = null;
  }
}

/* -------------------------------------------------------------------------- */
/*                       Cross-invocation session store                       */
/*                                                                            */
/* One accumulator per conversation thread, held in module scope so the        */
/* multi-round (tool handoff) accounting survives across separate generate()   */
/* invocations. Bounded by a small LRU so a long-lived process cannot leak.    */
/* -------------------------------------------------------------------------- */

const MAX_TRACKED_SESSIONS = 32;
const sessionStore = new Map<string, TurnStatAccumulator>();

/** Fetch (creating on miss) the accumulator for a conversation-thread key. */
export function getSessionAccumulator(key: string): TurnStatAccumulator {
  const existing = sessionStore.get(key);
  if (existing) {
    // Touch → move to most-recently-used end of the ordered map.
    sessionStore.delete(key);
    sessionStore.set(key, existing);
    return existing;
  }

  if (sessionStore.size >= MAX_TRACKED_SESSIONS) {
    const oldest = sessionStore.keys().next().value;
    if (oldest !== undefined) sessionStore.delete(oldest);
  }

  const created = new TurnStatAccumulator();
  sessionStore.set(key, created);
  return created;
}