import { anthropic } from '@ai-sdk/anthropic';
import { openai, createOpenAI } from '@ai-sdk/openai';
import type { LanguageModel } from 'ai';

export type ModelProvider = 'anthropic' | 'openai' | 'fireworks';

export interface ModelConfig {
  provider: ModelProvider;
  modelId: string;
  maxTokens?: number;
  temperature?: number;
}

/**
 * Model name to provider configuration mapping
 */
const MODEL_MAP: Record<string, ModelConfig> = {
  // Anthropic Claude models
  'claude-sonnet-4-20250514': { provider: 'anthropic', modelId: 'claude-sonnet-4-20250514' },
  'claude-opus-4-20250514': { provider: 'anthropic', modelId: 'claude-opus-4-20250514' },
  'claude-3-5-haiku-20241022': { provider: 'anthropic', modelId: 'claude-3-5-haiku-20241022' },
  'claude-3-5-sonnet-20241022': { provider: 'anthropic', modelId: 'claude-3-5-sonnet-20241022' },

  // Shorthand aliases
  'sonnet': { provider: 'anthropic', modelId: 'claude-sonnet-4-20250514' },
  'opus': { provider: 'anthropic', modelId: 'claude-opus-4-20250514' },
  'haiku': { provider: 'anthropic', modelId: 'claude-3-5-haiku-20241022' },

  // OpenAI models
  'gpt-4o': { provider: 'openai', modelId: 'gpt-4o' },
  'gpt-4o-mini': { provider: 'openai', modelId: 'gpt-4o-mini' },
  'gpt-4-turbo': { provider: 'openai', modelId: 'gpt-4-turbo' },
  'o1': { provider: 'openai', modelId: 'o1' },
  'o1-mini': { provider: 'openai', modelId: 'o1-mini' },

  // Fireworks models (Llama, etc.)
  'llama-v3p1-405b-instruct': {
    provider: 'fireworks',
    modelId: 'accounts/fireworks/models/llama-v3p1-405b-instruct',
  },
  'llama-v3p1-70b-instruct': {
    provider: 'fireworks',
    modelId: 'accounts/fireworks/models/llama-v3p1-70b-instruct',
  },
};

// Lazy-initialized Fireworks client
let fireworksClient: ReturnType<typeof createOpenAI> | null = null;

function getFireworksClient(): ReturnType<typeof createOpenAI> {
  if (!fireworksClient) {
    const apiKey = process.env.FIREWORKS_API_KEY;
    if (!apiKey) {
      throw new Error('FIREWORKS_API_KEY environment variable is required for Fireworks models');
    }
    fireworksClient = createOpenAI({
      apiKey,
      baseURL: 'https://api.fireworks.ai/inference/v1',
    });
  }
  return fireworksClient;
}

/**
 * Get a language model instance by name
 *
 * Supports models from:
 * - Anthropic (Claude): claude-sonnet-4-20250514, claude-opus-4-20250514, claude-3-5-haiku-20241022
 * - OpenAI: gpt-4o, gpt-4o-mini, gpt-4-turbo, o1, o1-mini
 * - Fireworks: llama-v3p1-405b-instruct, llama-v3p1-70b-instruct
 *
 * @param modelName - The model name or alias
 * @returns AI SDK LanguageModel instance
 */
export function getModel(modelName: string): LanguageModel {
  const config = MODEL_MAP[modelName];
  if (!config) {
    throw new Error(
      `Unknown model: ${modelName}. Available models: ${Object.keys(MODEL_MAP).join(', ')}`
    );
  }

  switch (config.provider) {
    case 'anthropic':
      return anthropic(config.modelId);

    case 'openai':
      return openai(config.modelId);

    case 'fireworks':
      return getFireworksClient()(config.modelId);

    default:
      throw new Error(`Unknown provider: ${config.provider}`);
  }
}

/**
 * Get model configuration by name
 */
export function getModelConfig(modelName: string): ModelConfig | undefined {
  return MODEL_MAP[modelName];
}

/**
 * Check if a model name is valid
 */
export function isValidModel(modelName: string): boolean {
  return modelName in MODEL_MAP;
}

/**
 * List all available model names
 */
export function listModels(): string[] {
  return Object.keys(MODEL_MAP);
}
