import Lock from "@common/lock";
import SSE from "@common/sse";
import { throttle } from "@common/utils";
import delay from "delay";
import { Huggingface } from "./huggingface";
import { approximateTokens, normalizeModel } from "./model";

interface IBaseConnection {
    type: 'kobold' | 'horde';
    instruct: string;
    url?: string;
    apiKey?: string;
    model?: string;
}

interface IKoboldConnection extends IBaseConnection {
    type: 'kobold';
    url: string;
}

interface IHordeConnection extends IBaseConnection {
    type: 'horde';
    apiKey?: string;
    model: string;
}

export type IConnection = IKoboldConnection | IHordeConnection;

interface IHordeWorker {
    id: string;
    models: string[];
    flagged: boolean;
    online: boolean;
    maintenance_mode: boolean;
    max_context_length: number;
    max_length: number;
    performance: string;
}

export interface IHordeModel {
    name: string;
    hordeNames: string[];
    maxLength: number;
    maxContext: number;
    workers: string[];
}

interface IHordeResult {
    faulted: boolean;
    done: boolean;
    finished: number;
    kudos: number;
    generations?: {
        text: string;
    }[];
}

const DEFAULT_GENERATION_SETTINGS = {
    temperature: 0.8,
    min_p: 0.1,
    rep_pen: 1.08,
    rep_pen_range: -1,
    rep_pen_slope: 0.7,
    top_k: 100,
    top_p: 0.92,
    banned_tokens: [] as string[],
    max_length: 300,
    trim_stop: true,
    stop_sequence: ['[INST]', '[/INST]', '</s>', '<|'],
    dry_allowed_length: 5,
    dry_multiplier: 0.8,
    dry_base: 1,
    dry_sequence_breakers: ["\n", ":", "\"", "*"],
    dry_penalty_last_n: 0
}

const MIN_PERFORMANCE = 2.0;
const MIN_WORKER_CONTEXT = 8192;
const MAX_HORDE_LENGTH = 512;
const MAX_HORDE_CONTEXT = 32000;
export const HORDE_ANON_KEY = '0000000000';

export type IGenerationSettings = Partial<typeof DEFAULT_GENERATION_SETTINGS>;

export namespace Connection {
    const AIHORDE = 'https://aihorde.net';

    let abortController = new AbortController();

    export interface TextChunk {
        text: string;
        cost: number;
    }

    async function* generateKobold(url: string, prompt: string, extraSettings: IGenerationSettings = {}): AsyncGenerator<TextChunk> {
        const sse = new SSE(`${url}/api/extra/generate/stream`, {
            payload: JSON.stringify({
                ...DEFAULT_GENERATION_SETTINGS,
                ...extraSettings,
                prompt,
            }),
        });

        const messages: string[] = [];
        const messageLock = new Lock();
        let end = false;

        sse.addEventListener('message', (e) => {
            if (e.data) {
                {
                    const { token, finish_reason } = JSON.parse(e.data);
                    messages.push(token);

                    if (finish_reason && finish_reason !== 'null') {
                        end = true;
                    }
                }
            }
            messageLock.release();
        });

        const handleEnd = () => {
            end = true;
            messageLock.release();
        };

        abortController.signal.addEventListener('abort', handleEnd);
        sse.addEventListener('error', handleEnd);
        sse.addEventListener('abort', handleEnd);
        sse.addEventListener('readystatechange', (e) => {
            if (e.readyState === SSE.CLOSED) handleEnd();
        });


        while (!end || messages.length) {
            while (messages.length > 0) {
                const text = messages.shift();
                if (text != null) {
                    try {
                        yield { text, cost: 0 };
                    } catch { }
                }
            }
            if (!end) {
                await messageLock.wait();
            }
        }

        sse.close();
    }

    async function* generateHorde(connection: IHordeConnection, prompt: string, extraSettings: IGenerationSettings = {}): AsyncGenerator<TextChunk> {
        if (!connection.model) {
            throw new Error('Horde not connected');
        }

        const models = await getHordeModels();
        const model = models.get(connection.model);
        if (model) {
            let maxLength = Math.min(model.maxLength, DEFAULT_GENERATION_SETTINGS.max_length);
            if (extraSettings.max_length && extraSettings.max_length < maxLength) {
                maxLength = extraSettings.max_length;
            }
            const baseTemperature = extraSettings.temperature ?? DEFAULT_GENERATION_SETTINGS.temperature;
            let currentTemperature = baseTemperature;
            const MAX_TEMPERATURE = 2.0;
            const TEMP_INCREMENT = 0.15;
            const RECOVERY_LENGTH = 16;

            const requestData = {
                prompt,
                params: {
                    ...DEFAULT_GENERATION_SETTINGS,
                    ...extraSettings,
                    n: 1,
                    max_context_length: model.maxContext,
                    max_length: maxLength,
                    rep_pen_range: Math.min(model.maxContext, 4096),
                    temperature: currentTemperature,
                },
                models: model.hordeNames,
                workers: model.workers,
            };
            const bannedTokens = requestData.params.banned_tokens ?? [];
            let recoveryMode = false;

            const { signal } = abortController;

            while (true) {
                const generateResponse = await fetch(`${AIHORDE}/api/v2/generate/text/async`, {
                    method: 'POST',
                    body: JSON.stringify(requestData),
                    headers: {
                        'Content-Type': 'application/json',
                        apikey: connection.apiKey || HORDE_ANON_KEY,
                    },
                    signal,
                });

                if (!generateResponse.ok || generateResponse.status >= 400) {
                    throw new Error(`Error starting generation: ${generateResponse.statusText}: ${await generateResponse.text()}`);
                }

                const { id } = await generateResponse.json() as { id: string };
                const request = async (method = 'GET'): Promise<TextChunk | null> => {
                    const response = await fetch(`${AIHORDE}/api/v2/generate/text/status/${id}`, { method });
                    if (response.ok && response.status < 400) {
                        const result: IHordeResult = await response.json();
                        if (result.generations?.length === 1) {
                            const { text } = result.generations[0];

                            return { text, cost: result.kudos };
                        }
                    } else {
                        throw new Error(await response.text());
                    }

                    return null;
                };

                const deleteRequest = async () => (await request('DELETE')) ?? { text: '', cost: 0 };
                let text: string | null = null;

                while (!text) {
                    try {
                        await delay(2500, { signal });

                        const response = await request();

                        if (response?.text) {
                            text = response.text;
                            let minStopIdx = text.length;
                            for (const sequence of requestData.params.stop_sequence) {
                                const stopIdx = text.indexOf(sequence);
                                if (stopIdx >= 0 && stopIdx < minStopIdx) {
                                    minStopIdx = stopIdx;
                                }
                            }
                            if (minStopIdx < text.length) {
                                text = text.slice(0, minStopIdx);
                            }

                            const locaseText = text.toLowerCase();
                            let unsloppedText = text;
                            let slopDetected = false;
                            let minSlopIdx = text.length;
                            let detectedBan = '';
                            for (const ban of bannedTokens) {
                                const slopIdx = locaseText.indexOf(ban.toLowerCase());
                                if (slopIdx >= 0 && slopIdx < minSlopIdx) {
                                    minSlopIdx = slopIdx;
                                    detectedBan = ban;
                                    slopDetected = true;
                                }
                            }
                            if (slopDetected) {
                                console.log(`[horde] slop '${detectedBan}' detected at ${minSlopIdx}`);
                                unsloppedText = unsloppedText.slice(0, minSlopIdx).trimEnd();
                            }

                            yield { text: unsloppedText, cost: response.cost };

                            requestData.prompt += unsloppedText;

                            if (slopDetected) {
                                recoveryMode = true;
                                requestData.params.max_length = RECOVERY_LENGTH;
                                currentTemperature = Math.min(MAX_TEMPERATURE, currentTemperature + TEMP_INCREMENT);
                                requestData.params.temperature = currentTemperature;
                                requestData.params.top_p = Math.min(0.98, 0.92 + (currentTemperature - baseTemperature) * 0.02);
                            } else if (recoveryMode) {
                                recoveryMode = false;
                                requestData.params.max_length = maxLength;
                                requestData.params.temperature = baseTemperature;
                                requestData.params.top_p = 0.92;
                                currentTemperature = baseTemperature;
                            } else {
                                return; // we are finished
                            }
                        }
                    } catch (e) {
                        if (!signal.aborted) {
                            console.error('Error in horde generation:', e);
                        }
                        return yield deleteRequest();
                    }
                }
            }
        }

        throw new Error(`Model ${connection.model} is offline`);
    }

    export async function* generate(connection: IConnection, prompt: string, extraSettings: IGenerationSettings = {}) {
        if (connection.type === 'kobold') {
            yield* generateKobold(connection.url, prompt, extraSettings);
        } else if (connection.type === 'horde') {
            yield* generateHorde(connection, prompt, extraSettings);
        }
    }

    export function stopGeneration() {
        abortController.abort();
        abortController = new AbortController(); // refresh
    }

    async function requestHordeModels(): Promise<Map<string, IHordeModel>> {
        try {
            const response = await fetch(`${AIHORDE}/api/v2/workers?type=text`);
            if (response.ok) {
                const workers: IHordeWorker[] = await response.json();
                const goodWorkers = workers.filter(w =>
                    w.online
                    && !w.maintenance_mode
                    && !w.flagged
                    && w.max_context_length >= MIN_WORKER_CONTEXT
                    && parseFloat(w.performance) >= MIN_PERFORMANCE
                );

                const models = new Map<string, IHordeModel>();

                for (const worker of goodWorkers) {
                    for (const modelName of worker.models) {
                        const normName = normalizeModel(modelName);
                        let model = models.get(normName);
                        if (!model) {
                            model = {
                                hordeNames: [],
                                maxContext: MAX_HORDE_CONTEXT,
                                maxLength: MAX_HORDE_LENGTH,
                                name: normName,
                                workers: []
                            }
                        }

                        if (!model.hordeNames.includes(modelName)) {
                            model.hordeNames.push(modelName);
                        }
                        if (!model.workers.includes(worker.id)) {
                            model.workers.push(worker.id);
                        }

                        model.maxContext = Math.min(model.maxContext, worker.max_context_length);
                        model.maxLength = Math.min(model.maxLength, worker.max_length);

                        models.set(normName, model);
                    }
                }

                return models;
            }
        } catch (e) {
            console.error(e);
        }

        return new Map();
    };

    export const getHordeModels = throttle(requestHordeModels, 10000);

    export async function getModelName(connection: IConnection): Promise<string> {
        if (connection.type === 'kobold') {
            try {
                const response = await fetch(`${connection.url}/api/v1/model`);
                if (response.ok) {
                    const { result } = await response.json();
                    return result;
                }
            } catch (e) {
                console.error('Error getting max tokens', e);
            }
        } else if (connection.type === 'horde') {
            return connection.model;
        }

        return '';
    }

    export async function getContextLength(connection: IConnection): Promise<number> {
        if (connection.type === 'kobold') {
            try {
                const response = await fetch(`${connection.url}/api/extra/true_max_context_length`);
                if (response.ok) {
                    const { value } = await response.json();
                    return value;
                }
            } catch (e) {
                console.error('Error getting max tokens', e);
            }
        } else if (connection.type === 'horde' && connection.model) {
            const models = await getHordeModels();
            const model = models.get(connection.model);
            if (model) {
                return model.maxContext;
            }
        }

        return 0;
    }

    export async function countTokens(connection: IConnection, prompt: string) {
        if (connection.type === 'kobold') {
            try {
                const response = await fetch(`${connection.url}/api/extra/tokencount`, {
                    body: JSON.stringify({ prompt }),
                    headers: { 'Content-Type': 'applicarion/json' },
                    method: 'POST',
                });
                if (response.ok) {
                    const { value } = await response.json();
                    return value;
                }
            } catch (e) {
                console.error('Error counting tokens:', e);
            }
        } else if (connection.type === 'horde') {
            const model = await getModelName(connection);
            const tokenizer = await Huggingface.findTokenizer(model);
            if (tokenizer) {
                try {
                    const { ids } = tokenizer.encode(prompt);
                    return ids.length;
                } catch (e) {
                    console.error('Error counting tokens with tokenizer:', e);
                }
            }
        }

        return approximateTokens(prompt);
    }
}
