import {
  MODELS_DEV_PROVIDER_MAP,
  type SupportedProvider,
  TimeInMs,
} from "@archestra/shared";
import { z } from "zod";
import { CacheKey, cacheManager } from "@/cache-manager";
import logger from "@/logging";
import { ModelModel } from "@/models";
import {
  type CreateModel,
  type ModelInputModality,
  ModelInputModalitySchema,
  type ModelOutputModality,
  ModelOutputModalitySchema,
} from "@/types";

/**
 * Cache key for tracking when we last synced from models.dev.
 */
const MODELS_DEV_SYNC_CACHE_KEY =
  `${CacheKey.ModelsDevSync}-timestamp` as const;

/**
 * How long to wait between models.dev syncs (24 hours).
 */
const SYNC_INTERVAL_MS = 24 * TimeInMs.Hour;

/**
 * models.dev API endpoint
 */
const MODELS_DEV_API_URL = "https://models.dev/api.json";

/**
 * How long a fetched models.dev API response is served from memory before
 * hitting the network again.
 */
const FETCH_CACHE_TTL_MS = 5 * TimeInMs.Minute;

/**
 * Default timeout for a single models.dev API fetch.
 */
const DEFAULT_FETCH_TIMEOUT_MS = 30 * TimeInMs.Second;

/**
 * Retry configuration for background sync
 */
const RETRY_CONFIG = {
  maxRetries: 3,
  baseDelayMs: 1000,
  maxDelayMs: 30000,
} as const;

// ============================================================================
// Types for models.dev API response
// ============================================================================

/**
 * Cost information for a model (prices per million tokens in USD)
 *
 * Only input/output + cache_read/cache_write are persisted and used in cost
 * calculation. Intentionally NOT persisted:
 * - `reasoning`: reasoning/thinking tokens are a subset of output tokens and are
 *   billed at the output rate (reasoning === output for current models), so a
 *   separate reasoning price would double-count what output pricing already covers.
 * - `input_audio` / `output_audio`: genuinely distinct rates, but nothing
 *   captures audio token counts yet (no adapter usage parsing, interaction
 *   columns, or o11y attributes), so persisting them would be dead data. Persist
 *   these as step 0 of an end-to-end audio cost-tracking feature, not on their own.
 *
 * @public — exported for testability
 */
export type ModelsDevCost = {
  input?: number;
  output?: number;
  reasoning?: number;
  cache_read?: number;
  cache_write?: number;
  input_audio?: number;
  output_audio?: number;
};

/**
 * Token limits for a model
 * @public — exported for testability
 */
export type ModelsDevLimit = {
  context?: number;
  input?: number;
  output?: number;
};

/**
 * Input/output modalities for a model
 * @public — exported for testability
 */
export type ModelsDevModalities = {
  input?: string[];
  output?: string[];
};

/**
 * Model status indicator
 * @public — exported for testability
 */
export type ModelsDevStatus = "alpha" | "beta" | "deprecated";

/**
 * A single model from the models.dev API
 * @public — exported for testability
 */
export type ModelsDevModel = {
  id: string;
  name: string;
  family?: string;
  attachment?: boolean;
  reasoning?: boolean;
  tool_call?: boolean;
  structured_output?: boolean;
  temperature?: boolean;
  knowledge?: string;
  release_date?: string;
  last_updated?: string;
  modalities?: ModelsDevModalities;
  open_weights?: boolean;
  cost?: ModelsDevCost;
  limit?: ModelsDevLimit;
  status?: ModelsDevStatus;
};

/**
 * A provider from the models.dev API
 */
export type ModelsDevProvider = {
  id: string;
  name: string;
  npm?: string;
  env?: string[];
  doc?: string;
  api?: string | null;
  models: Record<string, ModelsDevModel>;
};

/**
 * The full models.dev API response
 */
export type ModelsDevApiResponse = Record<string, ModelsDevProvider>;

// ============================================================================
// Zod schemas for API response validation
// ============================================================================

const ModelsDevCostSchema = z
  .object({
    input: z.number().optional(),
    output: z.number().optional(),
    reasoning: z.number().optional(),
    cache_read: z.number().optional(),
    cache_write: z.number().optional(),
    input_audio: z.number().optional(),
    output_audio: z.number().optional(),
  })
  .optional();

const ModelsDevLimitSchema = z
  .object({
    context: z.number().optional(),
    input: z.number().optional(),
    output: z.number().optional(),
  })
  .optional();

const ModelsDevModalitiesSchema = z
  .object({
    input: z.array(z.string()).optional(),
    output: z.array(z.string()).optional(),
  })
  .optional();

const ModelsDevModelSchema = z.object({
  id: z.string(),
  name: z.string(),
  family: z.string().optional(),
  attachment: z.boolean().optional(),
  reasoning: z.boolean().optional(),
  tool_call: z.boolean().optional(),
  structured_output: z.boolean().optional(),
  temperature: z.boolean().optional(),
  knowledge: z.string().optional(),
  release_date: z.string().optional(),
  last_updated: z.string().optional(),
  modalities: ModelsDevModalitiesSchema,
  open_weights: z.boolean().optional(),
  cost: ModelsDevCostSchema,
  limit: ModelsDevLimitSchema,
  status: z.enum(["alpha", "beta", "deprecated"]).optional(),
});

const ModelsDevProviderSchema = z.object({
  id: z.string(),
  name: z.string(),
  npm: z.string().optional(),
  env: z.array(z.string()).optional(),
  doc: z.string().optional(),
  api: z.string().nullable().optional(),
  models: z.record(z.string(), ModelsDevModelSchema),
});

const ModelsDevApiResponseSchema = z.record(
  z.string(),
  ModelsDevProviderSchema,
);

/**
 * Convert a models.dev `cost` object (prices per million tokens) into our
 * per-token string representation. Returns null for any price the registry omits.
 * @public — shared by the registry client and provider model sync.
 */
export function modelsDevCostToPerToken(cost: ModelsDevCost | undefined): {
  promptPricePerToken: string | null;
  completionPricePerToken: string | null;
  cacheReadPricePerToken: string | null;
  cacheWritePricePerToken: string | null;
} {
  // Round to the `numeric(20, 12)` column precision so the per-token string is
  // free of float-division noise (e.g. 0.8 / 1e6 = 8.000000000000001e-7) and
  // matches exactly what the database stores.
  const perToken = (perMillion: number | undefined): string | null =>
    perMillion !== undefined
      ? Number.parseFloat((perMillion / 1_000_000).toFixed(12)).toString()
      : null;

  return {
    promptPricePerToken: perToken(cost?.input),
    completionPricePerToken: perToken(cost?.output),
    cacheReadPricePerToken: perToken(cost?.cache_read),
    cacheWritePricePerToken: perToken(cost?.cache_write),
  };
}

/**
 * Sanitize a models.dev output-token limit into a stored `outputLength`.
 * Keeps only positive integers; drops 0/null/negative/non-integer garbage so a
 * reseller's malformed row cannot poison the model's output-token budget.
 * @public — shared by the registry client and provider model sync.
 */
export function sanitizeOutputLimit(
  output: number | null | undefined,
): number | null {
  return typeof output === "number" && Number.isInteger(output) && output > 0
    ? output
    : null;
}

// ============================================================================
// Client implementation
// ============================================================================

/**
 * models.dev Model Registry Client.
 *
 * Fetches model metadata from models.dev API and syncs it to our database.
 * Validated fetch results are cached in memory for a short TTL, and
 * concurrent callers share a single in-flight request. The raw fallback used
 * when schema validation fails and the `{}` error result are never cached,
 * so a transient bad response does not stick for the TTL.
 *
 * @public — exported so tests can construct instances with custom options
 */
export class ModelsDevClient {
  private readonly fetchTimeoutMs: number;
  // Hand-rolled instead of cacheManager/LRUCacheManager: single-flight
  // coalescing needs the in-flight promise as instance state, and the
  // multi-MB registry payload does not belong in the Postgres-backed cache.
  private cachedResponse: {
    data: ModelsDevApiResponse;
    fetchedAt: number;
  } | null = null;
  private inflightFetch: Promise<ModelsDevApiResponse> | null = null;
  // Bumped by clearFetchCache so a fetch started before the clear cannot
  // repopulate the cache with pre-clear data when it settles.
  private fetchGeneration = 0;

  constructor(opts?: { fetchTimeoutMs?: number }) {
    this.fetchTimeoutMs = opts?.fetchTimeoutMs ?? DEFAULT_FETCH_TIMEOUT_MS;
  }

  /**
   * Fetches all providers and models from models.dev API.
   * Validates the response against the expected schema.
   * Serves a cached response within the TTL and deduplicates concurrent calls.
   */
  async fetchModelsFromApi(): Promise<ModelsDevApiResponse> {
    if (
      this.cachedResponse &&
      Date.now() - this.cachedResponse.fetchedAt < FETCH_CACHE_TTL_MS
    ) {
      return this.cachedResponse.data;
    }

    if (this.inflightFetch) {
      return this.inflightFetch;
    }

    const fetchPromise = this.doFetchModelsFromApi();
    this.inflightFetch = fetchPromise;
    try {
      return await fetchPromise;
    } finally {
      if (this.inflightFetch === fetchPromise) {
        this.inflightFetch = null;
      }
    }
  }

  /**
   * Drops the in-memory fetch cache and detaches any in-flight fetch so the
   * next call hits the network. Tests use it to isolate cases.
   */
  clearFetchCache(): void {
    this.fetchGeneration++;
    this.cachedResponse = null;
    this.inflightFetch = null;
  }

  /**
   * Maps a models.dev provider ID to an Archestra provider.
   * Returns null if the provider is not supported.
   */
  mapProvider(providerId: string): SupportedProvider | null {
    const mappedProvider = MODELS_DEV_PROVIDER_MAP[providerId];
    if (mappedProvider === undefined) {
      logger.debug({ providerId }, "Unknown models.dev provider, skipping");
      return null;
    }
    return mappedProvider;
  }

  /**
   * Converts a models.dev model to our CreateModel format.
   * Returns null if the model's provider is not supported.
   */
  convertToModel(
    providerId: string,
    model: ModelsDevModel,
  ): CreateModel | null {
    const provider = this.mapProvider(providerId);
    if (!provider) {
      return null;
    }

    // Map input modalities using Zod schema for validation
    const inputModalities: ModelInputModality[] = [];
    for (const mod of model.modalities?.input ?? []) {
      const result = ModelInputModalitySchema.safeParse(mod);
      if (result.success) {
        inputModalities.push(result.data);
      }
    }
    if (inputModalities.length === 0) {
      inputModalities.push("text");
    }

    // Map output modalities using Zod schema for validation
    const outputModalities: ModelOutputModality[] = [];
    for (const mod of model.modalities?.output ?? []) {
      const result = ModelOutputModalitySchema.safeParse(mod);
      if (result.success) {
        outputModalities.push(result.data);
      }
    }
    if (outputModalities.length === 0) {
      outputModalities.push("text");
    }

    // Convert cost from per-million to per-token (store as string for precision)
    const prices = modelsDevCostToPerToken(model.cost);

    return {
      externalId: `${providerId}/${model.id}`,
      provider,
      modelId: model.id,
      description: model.name,
      contextLength: model.limit?.context ?? null,
      outputLength: sanitizeOutputLimit(model.limit?.output),
      inputModalities,
      outputModalities,
      supportsToolCalling: model.tool_call ?? false,
      promptPricePerToken: prices.promptPricePerToken,
      completionPricePerToken: prices.completionPricePerToken,
      cacheReadPricePerToken: prices.cacheReadPricePerToken,
      cacheWritePricePerToken: prices.cacheWritePricePerToken,
      lastSyncedAt: new Date(),
    };
  }

  /**
   * Checks if we need to sync based on the last sync timestamp.
   */
  async shouldSync(): Promise<boolean> {
    const lastSyncTime = await cacheManager.get<number>(
      MODELS_DEV_SYNC_CACHE_KEY,
    );

    if (!lastSyncTime) {
      return true;
    }

    const timeSinceLastSync = Date.now() - lastSyncTime;
    return timeSinceLastSync >= SYNC_INTERVAL_MS;
  }

  /**
   * Syncs model metadata from models.dev to our database.
   * Only syncs if the cache has expired (24h by default).
   *
   * @param force - If true, bypass cache and sync immediately
   * @returns Number of models synced
   */
  async syncModelMetadata(force = false): Promise<number> {
    if (!force && !(await this.shouldSync())) {
      logger.debug(
        "models.dev model metadata sync skipped (cache still valid)",
      );
      return 0;
    }

    logger.info("Starting models.dev model metadata sync");

    const apiResponse = await this.fetchModelsFromApi();
    const providerIds = Object.keys(apiResponse);

    if (providerIds.length === 0) {
      logger.warn("No providers returned from models.dev API");
      return 0;
    }

    logger.info(
      { totalProviders: providerIds.length },
      "Fetched providers from models.dev API",
    );

    const modelsToSync: CreateModel[] = [];
    const skippedProviders = new Set<string>();
    let totalModels = 0;

    for (const providerId of providerIds) {
      const provider = apiResponse[providerId];
      if (!provider.models) {
        continue;
      }

      for (const modelId of Object.keys(provider.models)) {
        totalModels++;
        const model = provider.models[modelId];
        const converted = this.convertToModel(providerId, model);

        if (converted) {
          modelsToSync.push(converted);
        } else {
          skippedProviders.add(providerId);
        }
      }
    }

    // Deduplicate models by (provider, modelId) key - this can happen when
    // multiple models.dev providers map to the same Archestra provider
    // (e.g., both google and google-vertex map to gemini).
    //
    // We use a deterministic priority system: prefer direct API providers over
    // derived/vertex providers. The externalId contains the original models.dev
    // provider (e.g., "google/gemini-2.5-flash" vs "google-vertex/gemini-2.5-flash").
    const preferredSourcePrefixes: Record<SupportedProvider, string[]> = {
      gemini: ["google/"], // Prefer google over google-vertex
      openai: ["openai/"], // Prefer direct providers over aggregators
      openrouter: ["openrouter/"],
      anthropic: ["anthropic/"],
      cohere: ["cohere/"],
      cerebras: ["cerebras/"],
      mistral: ["mistral/"],
      perplexity: ["perplexity/"],
      groq: ["groq/"],
      xai: ["xai/"],
      bedrock: ["amazon-bedrock/"],
      ollama: ["ollama/"],
      vllm: ["vllm/"],
      zhipuai: ["zhipuai/"],
      deepseek: ["deepseek/"],
      minimax: ["minimax/"],
      azure: ["azure/"],
      // Not synced via models.dev (subscription-dependent /models endpoint)
      "github-copilot": [],
    };

    const getSourcePriority = (model: CreateModel): number => {
      const prefixes = preferredSourcePrefixes[model.provider] ?? [];
      for (let i = 0; i < prefixes.length; i++) {
        if (model.externalId.startsWith(prefixes[i])) {
          return i; // Lower index = higher priority
        }
      }
      return prefixes.length; // Not in preferred list = lowest priority
    };

    const deduplicatedModels = new Map<string, CreateModel>();
    for (const model of modelsToSync) {
      const key = `${model.provider}:${model.modelId}`;
      const existing = deduplicatedModels.get(key);

      if (!existing) {
        deduplicatedModels.set(key, model);
      } else {
        // Keep the model with higher priority (lower priority number)
        const modelPriority = getSourcePriority(model);
        const existingPriority = getSourcePriority(existing);

        if (modelPriority < existingPriority) {
          deduplicatedModels.set(key, model);
        }
        // Otherwise keep existing (same priority = first occurrence wins)
      }
    }
    const uniqueModelsToSync = Array.from(deduplicatedModels.values());

    logger.info(
      {
        totalModels,
        modelsToSync: modelsToSync.length,
        uniqueModelsToSync: uniqueModelsToSync.length,
        duplicatesRemoved: modelsToSync.length - uniqueModelsToSync.length,
        skippedProviders: Array.from(skippedProviders),
      },
      "Filtered models for sync",
    );

    if (uniqueModelsToSync.length > 0) {
      await ModelModel.bulkUpsert(uniqueModelsToSync);
    }

    await this.updateSyncTimestamp();

    logger.info(
      { syncedCount: uniqueModelsToSync.length },
      "models.dev model metadata sync completed",
    );

    return uniqueModelsToSync.length;
  }

  /**
   * Convenience method to sync if needed (non-blocking).
   * Call this in the models route to trigger background sync.
   * Uses exponential backoff retry for transient failures.
   */
  syncIfNeeded(): void {
    this.syncWithRetry().catch((error) => {
      logger.error(
        { error: error instanceof Error ? error.message : String(error) },
        "Background models.dev sync failed after all retries",
      );
    });
  }

  /**
   * Attempts to sync model metadata with exponential backoff retry.
   */
  private async syncWithRetry(): Promise<number> {
    let lastError: Error | undefined;

    for (let attempt = 0; attempt <= RETRY_CONFIG.maxRetries; attempt++) {
      try {
        return await this.syncModelMetadata();
      } catch (error) {
        lastError = error instanceof Error ? error : new Error(String(error));

        if (attempt < RETRY_CONFIG.maxRetries) {
          const delayMs = Math.min(
            RETRY_CONFIG.baseDelayMs * 2 ** attempt,
            RETRY_CONFIG.maxDelayMs,
          );

          logger.warn(
            {
              attempt: attempt + 1,
              maxRetries: RETRY_CONFIG.maxRetries,
              delayMs,
              error: lastError.message,
            },
            "models.dev sync failed, retrying",
          );

          await this.sleep(delayMs);
        }
      }
    }

    throw lastError;
  }

  /**
   * Sleeps for the specified duration.
   */
  private sleep(ms: number): Promise<void> {
    return new Promise((resolve) => setTimeout(resolve, ms));
  }

  /**
   * Performs the actual network fetch. Only validated responses are cached:
   * the raw validation-fallback and the `{}` error result are returned
   * uncached so retries refetch instead of reusing a bad payload.
   */
  private async doFetchModelsFromApi(): Promise<ModelsDevApiResponse> {
    const generation = this.fetchGeneration;
    try {
      const response = await fetch(MODELS_DEV_API_URL, {
        signal: AbortSignal.timeout(this.fetchTimeoutMs),
      });
      if (!response.ok) {
        throw new Error(`HTTP ${response.status}: ${response.statusText}`);
      }

      const json = await response.json();
      const parseResult = ModelsDevApiResponseSchema.safeParse(json);

      if (!parseResult.success) {
        logger.warn(
          { errors: parseResult.error.format() },
          "models.dev API response validation failed, using partial data",
        );
        // Fall back to casting if validation fails - the API may have added new fields
        return json as ModelsDevApiResponse;
      }

      if (generation === this.fetchGeneration) {
        this.cachedResponse = { data: parseResult.data, fetchedAt: Date.now() };
      }
      return parseResult.data;
    } catch (error) {
      logger.error(
        { error: error instanceof Error ? error.message : String(error) },
        "Error fetching models from models.dev API",
      );
      return {};
    }
  }

  /**
   * Updates the last sync timestamp in cache.
   */
  private async updateSyncTimestamp(): Promise<void> {
    await cacheManager.set(
      MODELS_DEV_SYNC_CACHE_KEY,
      Date.now(),
      SYNC_INTERVAL_MS,
    );
  }
}

export const modelsDevClient = new ModelsDevClient();
