import { BaseProvider } from '~/lib/modules/llm/base-provider';
import type { ModelInfo } from '~/lib/modules/llm/types';
import type { IProviderSetting } from '~/types/model';
import type { LanguageModelV1 } from 'ai';
import { createOpenAI } from '@ai-sdk/openai';

export default class AstraChatProvider extends BaseProvider {
  name = 'AstraChat';
  getApiKeyLink = 'https://integrate.api.nvidia.com/';

  config = {
    apiTokenKey: 'NVAPI_KEY',
  };

  staticModels: ModelInfo[] = [
    { name: 'Kimi K2 0905', label: 'Kimi K2 0905', provider: 'AstraChat', maxTokenAllowed: 4096 },
    { name: 'Grok Code Fast 1', label: 'Grok Code Fast 1', provider: 'AstraChat', maxTokenAllowed: 4096 },
    { name: 'Stealth', label: 'Stealth', provider: 'AstraChat', maxTokenAllowed: 4096 },
  ];

  async getDynamicModels(
    apiKeys?: Record<string, string>,
    settings?: IProviderSetting,
    serverEnv?: Record<string, string>,
  ): Promise<ModelInfo[]> {
    // AstraChat only supports the Coder model
    return this.staticModels;
  }

  getModelInstance(options: {
    model: string;
    serverEnv?: Env;
    apiKeys?: Record<string, string>;
    providerSettings?: Record<string, IProviderSetting>;
  }): LanguageModelV1 {
    const { apiKeys, providerSettings, serverEnv } = options;

    if (options.model === 'Grok Code Fast 1') {
      return createOpenAI({
        baseURL: 'https://opencode.ai/zen/v1',
        apiKey: 'sk-c5bkHCKoZ9VvSRkueeSmKzUE3YDJ2NOb50Rhyeb46l84ylc0JQuyFRUXJ2Gm7YdS',
      })('grok-code');
    }

    if (options.model === 'Stealth') {
      return createOpenAI({
        baseURL: 'https://opencode.ai/zen/v1',
        apiKey: 'sk-c5bkHCKoZ9VvSRkueeSmKzUE3YDJ2NOb50Rhyeb46l84ylc0JQuyFRUXJ2Gm7YdS',
      })('big-pickle');
    }

    // For Coder model, use Nvidia
    const { apiKey } = this.getProviderBaseUrlAndKey({
      apiKeys,
      providerSettings: providerSettings?.[this.name],
      serverEnv: serverEnv as any,
      defaultBaseUrlKey: '',
      defaultApiTokenKey: 'NVAPI_KEY',
    });

    if (!apiKey) {
      throw new Error(`Missing API Key for ${this.name}`);
    }

    const actualModel = options.model === 'Kimi K2 0905' ? 'moonshotai/kimi-k2-instruct-0905' : options.model;

    return createOpenAI({
      baseURL: 'https://integrate.api.nvidia.com/v1',
      apiKey,
    })(actualModel);
  }
}