import type { Entity } from '@triplit/client';
import type schema from '#triplit/schema';

type Provider = Entity<typeof schema, 'providers'>;
type Model = Entity<typeof schema, 'models'>;

export interface ModelWithProvider extends Model {
    provider: Provider;
}

export interface ProviderWithModels extends Provider {
    models: Model[];
}

export const useModels = () => {
    const nuxtApp = useNuxtApp();
    const triplit = useTriplitClient();

    if (!nuxtApp._modelsState) {
        nuxtApp._modelsState = {
            providers: shallowRef([]),
            isReady: ref(false)
        };
    }


    const state = nuxtApp._modelsState as {
        providers: Ref<ProviderWithModels[]>;
        initPromise: Promise<void> | null;
    };

    const init = () => {
        if (state.initPromise) return;

        state.initPromise = (async () => {
            const query = triplit.query('providers').Include('models');

            const { results } = await useQuery('providers', triplit, query)
            watch(results, (newProviders) => {
                state.providers.value = newProviders as unknown as ProviderWithModels[] || [];
            }, { immediate: true, flush: 'sync' });
        })();
    };

    const allModels = computed<ModelWithProvider[]>(() => {
        const result: ModelWithProvider[] = [];
        for (const provider of state.providers.value) {
            if (!provider.enabled) continue;
            for (const model of provider.models || []) {
                if (!model.enabled) continue;
                result.push({ ...model, provider });
            }
        }
        return result;
    });

    const getModel = (id: string): ModelWithProvider | undefined => {
        return allModels.value.find((model) => model.id === id);
    };

    const getProvider = (id: string): ProviderWithModels | undefined => {
        return state.providers.value.find((provider) => provider.id === id);
    };

    const getFirstAvailableModel = () => {
        return allModels.value[0] ?? null;
    };

    return {
        init,
        providers: state.providers,
        allModels,
        getModel,
        getProvider,
        getFirstAvailableModel
    };
};