import { createCerebras, type CerebrasProvider } from "@ai-sdk/cerebras";
import { createGoogleGenerativeAI, type GoogleGenerativeAIProvider } from "@ai-sdk/google";
import { createOpenRouter, type OpenRouterProvider } from "@openrouter/ai-sdk-provider";
import { createOllama, type OllamaProvider } from "ai-sdk-ollama";
import { createOpenAICompatible, type OpenAICompatibleProvider } from '@ai-sdk/openai-compatible';
import { createCohere, type CohereProvider } from '@ai-sdk/cohere';
import { type Entity } from "@triplit/client";
import { schema } from "~~/triplit/schema";
import { type StreamTextTransform } from "ai";
import { transformCerebrasReasoningStream } from "./cerebras";
import { providerBaseUrls } from "~/types/model";
import { type Result, Err, Ok } from "~~/types/result";
// import { createLongcatTransformer } from "./longcat";

export type ModelGateway = OpenRouterProvider | OllamaProvider | CerebrasProvider | GoogleGenerativeAIProvider | OpenAICompatibleProvider | CohereProvider;

export interface Gateway {
    gateway: ModelGateway;
    streamTransformer: StreamTextTransform<{}> | StreamTextTransform<{}>[] | undefined;
    textTransformer: ((text: string) => string) | ((text: string) => string)[] | undefined;
}

export interface Provider {
    gateway: Gateway | null;
    endpoint: {
        baseURL: string;
        modelsEndpoint: string | null;
        headers: Record<string, string>;
    }
}

export enum GatewayFetchError {
    NoProviderApiKey = 0,
    NoProviderBaseUrl,
}

export async function getProviderDetails(provider: Entity<typeof schema, 'providers'>, providerApiKey?: string, model?: Entity<typeof schema, 'models'>): Promise<Result<Provider, GatewayFetchError>> {
    let gateway: Gateway = {} as Gateway;

    let baseURL = undefined;
    let modelsEndpoint = undefined;
    let headers: Record<string, string> = {};

    if (provider.config.apiProxyUrl && provider.config.apiProxyUrl.trim() !== '') {
        baseURL = provider.config.apiProxyUrl;
    } else {
        baseURL = providerBaseUrls[provider.type];
    }
    baseURL = baseURL.replace(/\/$/, '');

    switch (provider.type) {
        case 'openrouter': {
            if (providerApiKey === undefined) {
                return Err(GatewayFetchError.NoProviderApiKey);
            }

            gateway.gateway = createOpenRouter({
                apiKey: providerApiKey,
                headers: {
                    'HTTP-Referer': 'https://localhost:3000',
                    'X-Title': 'Veridian',
                },
            });
            headers['Authorization'] = `Bearer ${providerApiKey}`
            modelsEndpoint = `/models`;
        } break;
        case 'ollama': {
            if (baseURL === undefined) {
                return Err(GatewayFetchError.NoProviderBaseUrl);
            }

            if (providerApiKey !== undefined) {
                headers['Authorization'] = `Bearer ${providerApiKey}`
            }
            modelsEndpoint = `/api/tags`;

            if (model !== undefined) {
                const innerGateway = createOllama({
                    apiKey: providerApiKey,
                    baseURL,
                })

                gateway.gateway = ((modelId: string) => innerGateway(modelId, { think: [...model.attributes.capabilities].includes('reasoning') })) as OllamaProvider;
            }
        } break;
        case 'cerebras': {
            if (providerApiKey === undefined) {
                return Err(GatewayFetchError.NoProviderApiKey);
            }

            gateway.gateway = createCerebras({
                apiKey: providerApiKey,
                baseURL,
            })

            gateway.streamTransformer = transformCerebrasReasoningStream() as StreamTextTransform<{}>;
            gateway.textTransformer = (text: string) => {
                return text.split('</think>').at(-1)!.trim()
            };
            headers['Authorization'] = `Bearer ${providerApiKey}`
            modelsEndpoint = `/models`;
        } break;
        case 'google': {
            if (providerApiKey === undefined) {
                return Err(GatewayFetchError.NoProviderApiKey);
            }

            gateway.gateway = createGoogleGenerativeAI({
                apiKey: providerApiKey,
                baseURL,
            });
            headers['x-goog-api-key'] = `${providerApiKey}`
            modelsEndpoint = `/models`;
        } break;
        case 'longcat': {
            if (providerApiKey === undefined) {
                return Err(GatewayFetchError.NoProviderApiKey);
            }

            gateway.gateway = createOpenAICompatible({
                name: 'LongCat',
                apiKey: providerApiKey,
                baseURL: baseURL ?? providerBaseUrls[provider.type],
                includeUsage: true,
            })
            headers['Authorization'] = `Bearer ${providerApiKey}`
            modelsEndpoint = null;
            // streamTransformer = createLongcatTransformer() as StreamTextTransform<{}>;
        } break;
        case 'cohere': {
            if (providerApiKey === undefined) {
                return Err(GatewayFetchError.NoProviderApiKey);
            }

            gateway.gateway = createCohere({
                apiKey: providerApiKey,
                baseURL,
            });
            headers['Authorization'] = `Bearer ${providerApiKey}`
            modelsEndpoint = `/models`;
        } break;
    }

    return Ok({
        gateway,
        endpoint: {
            baseURL,
            modelsEndpoint,
            headers,
        }
    });
}