import type { AssistantMessage } from "@earendil-works/pi-ai";
import type { SessionEntry, SessionManager } from "@earendil-works/pi-coding-agent";
import { addToTotals, cacheTtlMsForApi, computeCacheHitPercent, emptyTotals, isIdleExpired } from "./cache-math.js";
import type { AssistantUsageMetric, CacheSessionMetrics } from "./types.js";

function isAssistantMessageEntry(entry: SessionEntry): entry is Extract<SessionEntry, { type: "message" }> & {
  message: AssistantMessage;
} {
  return entry.type === "message" && entry.message.role === "assistant";
}

interface ReadUsage {
  input: number;
  output: number;
  cacheRead: number;
  cacheWrite: number;
  totalTokens: number;
  cost: { input: number; cacheRead: number; cacheWrite: number };
}

/**
 * Reads usage defensively. Restored or legacy sessions can carry assistant
 * messages without a usage object, or with partial fields; a message that
 * reports no tokens at all is skipped instead of crashing the command.
 */
function readUsage(message: AssistantMessage): ReadUsage | undefined {
  const usage = message.usage;
  if (!usage) return undefined;
  const input = usage.input ?? 0;
  const output = usage.output ?? 0;
  const cacheRead = usage.cacheRead ?? 0;
  const cacheWrite = usage.cacheWrite ?? 0;
  if (input <= 0 && output <= 0 && cacheRead <= 0 && cacheWrite <= 0) return undefined;
  return {
    input,
    output,
    cacheRead,
    cacheWrite,
    totalTokens: usage.totalTokens ?? input + output + cacheRead + cacheWrite,
    cost: {
      input: usage.cost?.input ?? 0,
      cacheRead: usage.cost?.cacheRead ?? 0,
      cacheWrite: usage.cost?.cacheWrite ?? 0,
    },
  };
}

type SessionReader = Pick<SessionManager, "getEntries" | "getBranch">;

/** Minimal model lookup for cache-miss pricing and cache-TTL selection (pi's cache-stats.js math). */
type ModelLookup = {
  find(provider: string, modelId: string): { api?: string; cost?: { cacheRead?: number } } | undefined;
};

/** Same TTL and noise floor as pi's built-in cache diagnostics (cache-stats.js). */
const NOISE_FLOOR_TOKENS = 1024;
/**
 * Miss-ratio threshold at 100k cached tokens (10%); scales down as the cached
 * prefix grows: ~4.5% at 500k, ~3.2% at 1M. Clamped 2%..30%.
 *
 * The miss ratio is the ONLY meaningful cache-break indicator: pi's context
 * is a three-part prefix (system prompt + tool definitions + messages), and
 * new content is always appended at the end so the prefix stays cacheable.
 * Anything that rewrites the prefix (tool set, context files, extension
 * injection, mid-prefix edits) shows up as a large miss ratio — appending
 * fresh information (searches, file reads) does NOT, and needs no special
 * handling.
 */
const MISS_RATIO_AT_100K = 0.1;
const MISS_RATIO_MIN = 0.02;
const MISS_RATIO_MAX = 0.3;

export function collectCacheSessionMetrics(
  sessionManager: SessionReader,
  models?: ModelLookup,
): CacheSessionMetrics {
  const allEntries = sessionManager.getEntries();
  const activeBranchIds = new Set(sessionManager.getBranch().map((entry) => entry.id));

  const treeTotals = emptyTotals();
  const activeBranchTotals = emptyTotals();
  const allMessages: AssistantUsageMetric[] = [];

  let sequence = 0;
  let activeBranchSequence = 0;
  let prevAssistant: { modelKey: string; timestamp: number; promptTokens: number } | undefined;
  let sawCompactionSincePrev = false;
  let sawThinkingChangeSincePrev = false;

  for (const entry of allEntries) {
    if (entry.type === "compaction" || entry.type === "branch_summary") {
      sawCompactionSincePrev = true;
      continue;
    }
    if (entry.type === "thinking_level_change") {
      sawThinkingChangeSincePrev = true;
      continue;
    }
    if (!isAssistantMessageEntry(entry)) continue;
    const usage = readUsage(entry.message);
    if (!usage) continue;

    sequence += 1;

    const entryTime = Date.parse(entry.timestamp);
    const modelKey = `${entry.message.provider}/${entry.message.model}`;
    const promptTokens = usage.input + usage.cacheRead + usage.cacheWrite;
    const idleMs = prevAssistant ? Math.max(0, entryTime - prevAssistant.timestamp) : undefined;
    const idleExpired =
      idleMs !== undefined && isIdleExpired(idleMs, cacheTtlMsForApi(models?.find(entry.message.provider, entry.message.model)?.api));
    const modelChanged = prevAssistant ? prevAssistant.modelKey !== modelKey : undefined;
    const afterCompact = sawCompactionSincePrev || undefined;
    const thinkingChanged = sawThinkingChangeSincePrev || undefined;
    const hitPercent = computeCacheHitPercent(usage.input, usage.cacheRead, usage.cacheWrite);

    // Same miss math as pi's cache-stats: how much of the previous prompt
    // (the cacheable prefix) was re-billed as fresh. Appended content never
    // counts here; only prefix rewrites do.
    const prevPrompt = prevAssistant?.promptTokens ?? 0;
    const missedTokens =
      prevAssistant && !afterCompact ? Math.min(prevPrompt, promptTokens) - usage.cacheRead : 0;
    const missRatio = prevPrompt > 0 ? missedTokens / prevPrompt : 0;
    const missThreshold =
      prevPrompt > 0
        ? Math.min(MISS_RATIO_MAX, Math.max(MISS_RATIO_MIN, MISS_RATIO_AT_100K * Math.sqrt(100_000 / prevPrompt)))
        : MISS_RATIO_MAX;
    const significantMiss = missedTokens > NOISE_FLOOR_TOKENS ? missedTokens : 0;
    const prefixBroken = prevPrompt > 0 && missRatio > missThreshold;
    // A prompt that shinks by more than 5% and at least the noise floor cannot
    // be explained by appending or by idle expiry — the context was rebuilt
    // differently (restart replay, restored session, changed tool set).
    const promptShrunk =
      prevPrompt > 0 &&
      promptTokens < prevPrompt * 0.95 &&
      prevPrompt - promptTokens > NOISE_FLOOR_TOKENS
        ? true
        : undefined;
    // The "context rebuilt" tag: prefix broken with no recorded cause. Unlike
    // before, an expired idle gap no longer suppresses it. Idle expiry clears
    // the whole cache block, so it explains a miss only when nothing was hit:
    // a prompt that shrank, or a significant partial cache hit (the prefix
    // matched up to a fork point and diverged after it), is decisive evidence
    // of a rebuild even while the idle clock also ran out — the tags
    // co-occur.
    // (afterCompact is not needed here: compacted turns compute missedTokens
    // as 0 above, so they can never be prefixBroken.)
    const contextInvalidated = prefixBroken &&
      !modelChanged &&
      !thinkingChanged &&
      (promptShrunk || usage.cacheRead > NOISE_FLOOR_TOKENS || !idleExpired)
      ? true
      : undefined;

    // Extra cost = missed tokens billed at the actual paid rate (input +
    // cacheWrite incl. write premium) instead of the cache-read rate.
    const paidTokens = usage.input + usage.cacheWrite;
    const paidPerToken =
      paidTokens > 0 ? (usage.cost.input + usage.cost.cacheWrite) / paidTokens : 0;
    const readPerToken =
      usage.cacheRead > 0
        ? usage.cost.cacheRead / usage.cacheRead
        : (models?.find(entry.message.provider, entry.message.model)?.cost?.cacheRead ?? 0) /
          1_000_000;
    const missedCost =
      significantMiss > 0 ? significantMiss * Math.max(0, paidPerToken - readPerToken) : 0;

    const metric: AssistantUsageMetric = {
      sequence,
      activeBranchSequence: undefined,
      entryId: entry.id,
      timestamp: entry.timestamp,
      provider: entry.message.provider,
      model: entry.message.model,
      ...usage,
      cacheHitPercent: hitPercent,
      isOnActiveBranch: activeBranchIds.has(entry.id),
      modelChanged,
      idleMs,
      idleExpired: idleExpired || undefined,
      afterCompact,
      thinkingChanged,
      promptShrunk,
      contextInvalidated,
      missFlagged: prefixBroken,
      ...(prevAssistant ? { missRatio } : {}),
      ...(significantMiss > 0 ? { missedTokens: significantMiss, missedCost } : {}),
    };

    addToTotals(treeTotals, metric);
    allMessages.push(metric);

    if (metric.isOnActiveBranch) {
      activeBranchSequence += 1;
      metric.activeBranchSequence = activeBranchSequence;
      addToTotals(activeBranchTotals, metric);
    }

    sawCompactionSincePrev = false;
    sawThinkingChangeSincePrev = false;
    prevAssistant = { modelKey, timestamp: entryTime, promptTokens };
  }

  return { allMessages, treeTotals, activeBranchTotals };
}