import { DEFAULT_CREDITS_PER_USD, POLICY_CAPS } from "../config/defaults.js";

export { DEFAULT_CREDITS_PER_USD };
import type { BudgetState, Policy, PriceCatalog, PriceRate } from "../types.js";

export const BUDGET_ALLOCATION = { explorationImplementation: 0.55, verification: 0.2, escalation: 0.15, finalSynthesis: 0.1 } as const;

export function budgetAllocationFor(policy: Policy): Record<keyof typeof BUDGET_ALLOCATION, number> {
  const hardCap = POLICY_CAPS[policy].hardCap;
  return Object.fromEntries(Object.entries(BUDGET_ALLOCATION).map(([phase, share]) => [phase, hardCap * share])) as Record<keyof typeof BUDGET_ALLOCATION, number>;
}

export function estimateCredits(catalog: PriceCatalog, model: string, inputTokens: number, cachedInputTokens: number, outputTokens: number): number {
  const rate = catalog.models[model];
  if (!rate) throw new Error(`No price catalog entry for ${model}`);
  return (inputTokens * rate.input + cachedInputTokens * rate.cached + outputTokens * rate.output) / 1_000_000;
}

export interface UsdModelCost { input: number; cacheRead: number; output: number }

/** A registry cost is usable when it prices the two rates every call pays. */
export function usableUsdCost(cost: UsdModelCost | undefined): cost is UsdModelCost {
  if (!cost) return false;
  const finite = [cost.input, cost.cacheRead, cost.output].every((value) => Number.isFinite(value) && value >= 0);
  return finite && cost.input > 0 && cost.output > 0;
}

export function derivedRate(cost: UsdModelCost, creditsPerUsd = DEFAULT_CREDITS_PER_USD): PriceRate {
  return { input: cost.input * creditsPerUsd, cached: cost.cacheRead * creditsPerUsd, output: cost.output * creditsPerUsd };
}

/**
 * The catalogue a run actually prices against: declared entries, plus derived ones for models the
 * document does not name. Declared always wins, so deriving can never override a price the operator
 * wrote down — including a deliberate zero.
 */
export function resolveCatalog(catalog: PriceCatalog, derived: ReadonlyMap<string, PriceRate>): PriceCatalog {
  if (derived.size === 0) return catalog;
  return { ...catalog, models: { ...Object.fromEntries(derived), ...catalog.models } };
}

export function usdBudgetForCredits(catalog: PriceCatalog, models: readonly { model: string; cost: UsdModelCost }[], credits: number): number | undefined {
  if (!Number.isFinite(credits) || credits <= 0 || models.length === 0) return undefined;
  let creditsPerUsd = 0;
  for (const { model, cost } of models) {
    const rate = catalog.models[model];
    if (!rate || !usableUsdCost(cost)) return undefined;
    // A zero cached rate is real and common in the registry, so it drops out of the comparison
    // instead of voiding the whole conversion.
    const ratios = [rate.input / cost.input, rate.output / cost.output, ...(cost.cacheRead > 0 ? [rate.cached / cost.cacheRead] : [])];
    creditsPerUsd = Math.max(creditsPerUsd, ...ratios);
  }
  return creditsPerUsd > 0 ? credits / creditsPerUsd : undefined;
}

export function budgetFor(policy: Policy, spentCredits = 0, weeklyCreditBudget?: number, dailyCreditBudget?: number): BudgetState {
  const cap = POLICY_CAPS[policy];
  return { spentCredits, hardCap: cap.hardCap, softCap: cap.softCap, internalStopTarget: cap.internalStopTarget, weeklyCreditBudget, dailyCreditBudget };
}

export type BudgetPhase = "normal" | "no-speculation" | "finish-only" | "safe-turn" | "stop";

export function budgetPhase(state: BudgetState): BudgetPhase {
  const ratio = state.spentCredits / state.hardCap;
  if (ratio >= 1) return "stop";
  if (ratio >= 0.9) return "safe-turn";
  if (ratio >= 0.75) return "finish-only";
  if (ratio >= 0.6) return "no-speculation";
  return "normal";
}

export function maySpawn(state: BudgetState, estimatedCredits: number, expectedUsefulProbability: number, estimatedImpact: number, coordinationPenalty: number, latencyPenalty: number, speculative = true): boolean {
  if (!Number.isFinite(estimatedCredits) || estimatedCredits < 0) return false;
  const phase = budgetPhase(state);
  if (phase === "stop" || phase === "safe-turn" || (speculative && phase !== "normal")) return false;
  if (state.spentCredits + estimatedCredits > state.internalStopTarget) return false;
  return expectedUsefulProbability * estimatedImpact - estimatedCredits - coordinationPenalty - latencyPenalty > 0;
}