// *****************************************************************************
// Copyright (C) 2024 TypeFox GmbH.
//
// This program and the accompanying materials are made available under the
// terms of the Eclipse Public License v. 2.0 which is available at
// http://www.eclipse.org/legal/epl-2.0.
//
// This Source Code may also be made available under the following Secondary
// Licenses when the conditions for such availability set forth in the Eclipse
// Public License v. 2.0 are satisfied: GNU General Public License, version 2
// with the GNU Classpath Exception which is available at
// https://www.gnu.org/software/classpath/license.html.
//
// SPDX-License-Identifier: EPL-2.0 OR GPL-2.0-only WITH Classpath-exception-2.0
// *****************************************************************************

import { LanguageModelRegistry, LanguageModelStatus } from '@theia/ai-core';
import { getProxyUrl } from '@theia/ai-core/lib/node';
import { inject, injectable } from '@theia/core/shared/inversify';
import { OllamaModel } from './ollama-language-model';
import { OllamaLanguageModelsManager, OllamaModelDescription } from '../common';

@injectable()
export class OllamaLanguageModelsManagerImpl implements OllamaLanguageModelsManager {

    protected _host: string | undefined;
    protected _proxyUrl: string | undefined;

    @inject(LanguageModelRegistry)
    protected readonly languageModelRegistry: LanguageModelRegistry;

    get host(): string | undefined {
        return this._host ?? process.env.OLLAMA_HOST;
    }

    // Triggered from frontend. In case you want to use the models on the backend
    // without a frontend then call this yourself
    protected calculateStatus(host: string | undefined): LanguageModelStatus {
        return host ? { status: 'ready' } : { status: 'unavailable', message: 'No Ollama host set' };
    }

    async createOrUpdateLanguageModels(...models: OllamaModelDescription[]): Promise<void> {
        for (const modelDescription of models) {
            const existingModel = await this.languageModelRegistry.getLanguageModel(modelDescription.id);
            const hostProvider = () => this.host;

            const host = hostProvider();
            const normalizedHost = host && !host.includes('://') ? `http://${host}` : host;
            const proxyUrl = getProxyUrl(normalizedHost ?? 'http://localhost:11434', this._proxyUrl);

            if (existingModel) {
                if (!(existingModel instanceof OllamaModel)) {
                    console.warn(`Ollama: model ${modelDescription.id} is not an Ollama model`);
                    continue;
                }
                const status = this.calculateStatus(host);
                await this.languageModelRegistry.patchLanguageModel<OllamaModel>(modelDescription.id, {
                    proxy: proxyUrl,
                    status,
                    reasoningSupport: modelDescription.reasoningSupport
                });
            } else {
                const status = this.calculateStatus(host);
                this.languageModelRegistry.addLanguageModels([
                    new OllamaModel(
                        modelDescription.id,
                        modelDescription.model,
                        status,
                        hostProvider,
                        proxyUrl,
                        modelDescription.reasoningSupport
                    )
                ]);
            }
        }
    }

    removeLanguageModels(...modelIds: string[]): void {
        this.languageModelRegistry.removeLanguageModels(modelIds.map(id => `ollama/${id}`));
    }

    setProxyUrl(proxyUrl: string | undefined): void {
        if (proxyUrl) {
            this._proxyUrl = proxyUrl;
        } else {
            this._proxyUrl = undefined;
        }
    }

    async setHost(host: string | undefined): Promise<void> {
        this._host = host || undefined;
        const models = await this.languageModelRegistry.getLanguageModels();
        const ollamaModels = models.filter(model => model instanceof OllamaModel) as OllamaModel[];
        const status = this.calculateStatus(this.host);
        for (const model of ollamaModels) {
            model.status = status;
        }
    }
}
