import type { AiProviderType } from "@server/lib/aiProviderDefaults";
import type { AiUsage } from "@server/lib/aiUsageExtraction";
import {
    getAiModelCatalog,
    type AiModelCatalogEntry,
    type CatalogProvider
} from "@server/lib/aiModelCatalog";

export type AiModelPricing = {
    inputCostPerToken: number | null;
    outputCostPerToken: number | null;
    cacheReadInputTokenCost: number | null;
    outputCostPerReasoningToken: number | null;
    // True when the match came from a different catalog provider than the
    // one mapped to this provider's type (e.g. an openRouter/custom model
    // id that only matched a global search across every provider). Costs
    // found this way are a best-effort approximation, not a guarantee the
    // upstream provider bills at the same rate.
    approximate: boolean;
};

// Each of our provider types maps to at most one catalog provider. Provider types that proxy
// arbitrary underlying models (openRouter, vercelAiGateway, custom) have no
// mapping and always fall back to a global search.
const PROVIDER_CATALOG_MAP: Record<
    Exclude<AiProviderType, "custom">,
    CatalogProvider | null
> = {
    openai: "openai",
    anthropic: "anthropic",
    googleGemini: "gemini",
    vertexAi: "vertex",
    bedrock: "bedrock",
    microsoftFoundry: "azure",
    openRouter: null,
    vercelAiGateway: null
};

// Indexed view over the in-memory catalog, rebuilt only when
// getAiModelCatalog() returns a different array instance (i.e. after a
// background refresh swaps it out), not on every lookup.
let indexedCatalog: AiModelCatalogEntry[] | null = null;
let indexedByName: Map<string, AiModelCatalogEntry[]> = new Map();

function getIndexedCatalog(): Map<string, AiModelCatalogEntry[]> {
    const catalog = getAiModelCatalog();
    if (catalog === indexedCatalog) {
        return indexedByName;
    }

    const byName = new Map<string, AiModelCatalogEntry[]>();
    for (const entry of catalog) {
        if (!entry.model) continue;
        const list = byName.get(entry.model) ?? [];
        list.push(entry);
        byName.set(entry.model, list);
    }

    indexedCatalog = catalog;
    indexedByName = byName;
    return byName;
}

function stripVendorPrefix(modelId: string): string | null {
    const idx = modelId.indexOf("/");
    if (idx === -1 || idx === modelId.length - 1) {
        return null;
    }
    return modelId.slice(idx + 1);
}

function toPricing(
    entry: AiModelCatalogEntry,
    approximate: boolean
): AiModelPricing {
    return {
        inputCostPerToken: entry.pricing.input,
        outputCostPerToken: entry.pricing.output,
        cacheReadInputTokenCost: entry.pricing.cacheRead,
        outputCostPerReasoningToken: entry.pricing.reasoningOutput,
        approximate
    };
}

function findEntry(
    byName: Map<string, AiModelCatalogEntry[]>,
    modelId: string,
    provider: CatalogProvider | null
): AiModelCatalogEntry | null {
    const candidates = [modelId, stripVendorPrefix(modelId)].filter(
        (v): v is string => v != null
    );

    for (const key of candidates) {
        const entries = byName.get(key);
        if (!entries) continue;
        const match = provider
            ? entries.find((e) => e.provider === provider)
            : entries[0];
        if (match) {
            return match;
        }
    }
    return null;
}

/**
 * Looks up per-token pricing for a model, scoped first to the catalog
 * provider that corresponds to our provider type, then falling back to a
 * global search across every provider (marked `approximate`) for provider
 * types that proxy arbitrary underlying models.
 */
export function getModelPricing(
    providerType: AiProviderType,
    modelId: string | undefined
): AiModelPricing | null {
    if (!modelId) {
        return null;
    }

    const byName = getIndexedCatalog();
    const catalogProvider =
        providerType === "custom" ? null : PROVIDER_CATALOG_MAP[providerType];

    if (catalogProvider) {
        const scoped = findEntry(byName, modelId, catalogProvider);
        if (scoped) {
            return toPricing(scoped, false);
        }
    }

    const fallback = findEntry(byName, modelId, null);
    if (fallback) {
        return toPricing(fallback, true);
    }

    return null;
}

export type AiCostBreakdown = {
    promptCost: number;
    cacheReadCost: number;
    cacheWriteCost: number;
    completionCost: number;
    reasoningCost: number;
    totalCost: number;
};

/**
 * Computes a $ cost breakdown for a usage record given a model's pricing.
 * Cache writes and reasoning tokens fall back to the normal input/output
 * rate respectively when the catalog has no dedicated rate for them (the
 * catalog has no cache-write field at all, and only some models report a
 * distinct reasoning rate).
 */
export function calculateAiCost(
    pricing: AiModelPricing | null,
    usage: AiUsage
): AiCostBreakdown | null {
    if (!pricing) {
        return null;
    }

    const inputRate = pricing.inputCostPerToken ?? 0;
    const outputRate = pricing.outputCostPerToken ?? 0;
    const cacheReadRate = pricing.cacheReadInputTokenCost ?? inputRate;
    const reasoningRate = pricing.outputCostPerReasoningToken ?? outputRate;

    const promptCost = usage.promptTokens * inputRate;
    const cacheReadCost = usage.cacheReadTokens * cacheReadRate;
    const cacheWriteCost = usage.cacheWriteTokens * inputRate;
    const completionCost = usage.completionTokens * outputRate;
    const reasoningCost = usage.reasoningTokens * reasoningRate;

    return {
        promptCost,
        cacheReadCost,
        cacheWriteCost,
        completionCost,
        reasoningCost,
        totalCost:
            promptCost +
            cacheReadCost +
            cacheWriteCost +
            completionCost +
            reasoningCost
    };
}