import type { UIConfigField } from '../../../config/types.js';
import { getConfiguredModelProviderById } from '../../../config/serverRegistry.js';
import BaseModelProvider from '../../base/provider.js';
import type { Model, ModelList, ProviderMetadata } from '../../types.js';
import type BaseLLM from '../../base/llm.js';
import type BaseEmbedding from '../../base/embedding.js';
import OllamaLLM from './ollamaLLM.js';
import OllamaEmbedding from './ollamaEmbedding.js';

interface OllamaConfig {
  baseURL: string;
  embeddingBaseURL?: string;
}

const providerConfigFields: UIConfigField[] = [
  {
    type: 'string',
    name: 'Base URL',
    key: 'baseURL',
    description: 'The base URL for the Ollama',
    required: true,
    placeholder:
      process.env.DOCKER === '1'
        ? 'http://host.docker.internal:11434'
        : 'http://localhost:11434',
    env: 'OLLAMA_BASE_URL',
    scope: 'server',
  },
];

class OllamaProvider extends BaseModelProvider<OllamaConfig> {
  constructor(id: string, name: string, config: OllamaConfig) {
    super(id, name, config);
  }

  private async fetchModels(baseURL: string): Promise<Model[]> {
    const res = await fetch(`${baseURL}/api/tags`, {
      method: 'GET',
      headers: { 'Content-type': 'application/json' },
    });
    const data = (await res.json()) as { models?: { name?: string; model?: string }[] };
    return (data.models ?? []).map((m) => ({
      name: m.model ?? m.name ?? '',
      key: m.model ?? m.name ?? '',
    }));
  }

  async getDefaultModels(): Promise<ModelList> {
    try {
      const [chatModels, embeddingModels] = await Promise.all([
        this.fetchModels(this.config.baseURL),
        this.fetchModels(
          this.config.embeddingBaseURL ?? this.config.baseURL,
        ),
      ]);
      return { chat: chatModels, embedding: embeddingModels };
    } catch (err) {
      if (err instanceof TypeError) {
        throw new Error(
          'Error connecting to Ollama API. Please ensure the base URL is correct and the Ollama server is running.',
        );
      }
      throw err;
    }
  }

  async getModelList(): Promise<ModelList> {
    const defaultModels = await this.getDefaultModels();
    const configProvider = getConfiguredModelProviderById(this.id);
    if (!configProvider) {
      return defaultModels;
    }
    return {
      embedding: [
        ...defaultModels.embedding,
        ...configProvider.embeddingModels,
      ],
      chat: [...defaultModels.chat, ...configProvider.chatModels],
    };
  }

  async loadChatModel(key: string): Promise<BaseLLM<unknown>> {
    const modelList = await this.getModelList();
    const exists = modelList.chat.find((m) => m.key === key);
    if (!exists) {
      throw new Error(
        'Error Loading Ollama Chat Model. Invalid Model Selected',
      );
    }
    return new OllamaLLM({
      baseURL: this.config.baseURL,
      model: key,
    });
  }

  async loadEmbeddingModel(key: string): Promise<BaseEmbedding<unknown>> {
    const modelList = await this.getModelList();
    const exists = modelList.embedding.find((m) => m.key === key);
    if (!exists) {
      throw new Error(
        'Error Loading Ollama Embedding Model. Invalid Model Selected.',
      );
    }
    return new OllamaEmbedding({
      model: key,
      baseURL: this.config.embeddingBaseURL ?? this.config.baseURL,
    });
  }

  static parseAndValidate(raw: unknown): OllamaConfig {
    if (!raw || typeof raw !== 'object')
      throw new Error('Invalid config provided. Expected object');
    const obj = raw as Record<string, unknown>;
    if (!obj.baseURL)
      throw new Error('Invalid config provided. Base URL must be provided');
    return {
      baseURL: String(obj.baseURL),
      embeddingBaseURL: obj.embeddingBaseURL
        ? String(obj.embeddingBaseURL)
        : undefined,
    };
  }

  static getProviderConfigFields(): UIConfigField[] {
    return providerConfigFields;
  }

  static getProviderMetadata(): ProviderMetadata {
    return { key: 'ollama', name: 'Ollama' };
  }
}

export default OllamaProvider;
