import fs from "node:fs";
import path from "node:path";
import { APP_PATH } from "@server/lib/consts";
import type { AiProviderType } from "@server/lib/aiProviderDefaults";
import type { AiUsage } from "@server/lib/aiUsageExtraction";
import logger from "@server/logger";

const MODELS_JSON_PATH = path.join(APP_PATH, "models.json");

export type CatalogProvider =
    | "openai"
    | "anthropic"
    | "gemini"
    | "vertex"
    | "azure"
    | "bedrock";

export type AiModelCatalogEntry = {
    provider: CatalogProvider;
    model: string;
    pricing: {
        input: number | null;
        output: number | null;
        cacheRead: number | null;
        reasoningOutput: number | null;
    };
};

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
};

let modelsByName: Map<string, AiModelCatalogEntry[]> | null = null;

function loadModels(): Map<string, AiModelCatalogEntry[]> {
    if (modelsByName) {
        return modelsByName;
    }

    const byName = new Map<string, AiModelCatalogEntry[]>();
    try {
        if (fs.existsSync(MODELS_JSON_PATH)) {
            const raw = fs.readFileSync(MODELS_JSON_PATH, "utf-8");
            const parsed = JSON.parse(raw) as { data: AiModelCatalogEntry[] };
            for (const entry of parsed.data ?? []) {
                if (!entry.model) continue;
                const list = byName.get(entry.model) ?? [];
                list.push(entry);
                byName.set(entry.model, list);
            }
        } else {
            logger.debug(
                `AI model pricing file not found at ${MODELS_JSON_PATH}; cost calculation will fall back to unknown pricing`
            );
        }
    } catch (error) {
        logger.warn("Failed to load AI model pricing file", { error });
    }

    modelsByName = 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 = loadModels();
    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
    };
}