import SSE from '@common/sse';
import type { TObject } from '@common/typebox';

namespace LLM {
    export interface Connection {
        url: string;
        apiKey: string;
    }

    export interface ToolCall {
        id: string;
        type: 'function';
        function: {
            name: string;
            arguments: string | Record<string, any>; // JSON string
        };
    }

    interface ChatMessageUser {
        role: 'user';
        content: string;
    }

    interface ChatMessageAssistant {
        role: 'assistant';
        content: string;
        reasoning_content?: string;
        tool_calls?: ToolCall[];
    }

    interface ChatMessageSystem {
        role: 'system';
        content: string;
    }

    interface ChatMessageTool {
        role: 'tool';
        content: string;
        tool_call_id?: string;
    }

    export type ChatMessage = ChatMessageUser | ChatMessageAssistant | ChatMessageSystem | ChatMessageTool;

    export interface Tool {
        type: 'function';
        function: {
            name: string;
            description?: string;
            parameters: TObject;
        };
    }

    export interface ChatCompletionRequest {
        model: string;
        messages: ChatMessage[];
        tools?: Tool[];
        tool_choice?: 'none' | 'auto' | 'required' | { type: 'function'; function: { name: string } };
        parallel_tool_calls?: boolean;
        temperature?: number;
        max_tokens?: number;
        max_completion_tokens?: number;
        stop?: string | string[];
        banned_tokens?: string[];
        top_p?: number;
        top_k?: number;
        min_p?: number;
        frequency_penalty?: number;
        repetition_penalty?: number;
        reasoning?: {
            effort?: 'xhigh' | 'high' | 'medium' | 'low' | 'minimal' | 'none';
            exclude?: boolean;
            max_tokens?: number;
        };
        add_generation_prompt?: boolean;
        remove_last_eos?: boolean;
    }

    export interface ChatCompletionChoice {
        index: number;
        message: ChatMessage;
        finish_reason: 'stop' | 'tool_calls';
    }

    export interface ChatCompletionResponse {
        id: string;
        object: 'chat.completion';
        created: number;
        model: string;
        choices: ChatCompletionChoice[];
        usage: {
            prompt_tokens: number;
            completion_tokens: number;
            total_tokens: number;
            spent_kudos?: number;
        };
    }

    export interface ChatCompletionChunkChoice {
        index: number;
        delta: {
            role?: string;
            content?: string;
            reasoning_content?: string;
            tool_calls?: ToolCall[];
        };
        finish_reason: 'stop' | 'tool_calls' | null;
    }

    export interface ChatCompletionChunk {
        id: string;
        object: 'chat.completion.chunk';
        created: number;
        model: string;
        choices: ChatCompletionChunkChoice[];
    }

    export interface ChatCompletionError {
        error: string;
    }

    export interface ImageGenerationSettings {
        width?: number;
        height?: number;
        negative_prompt?: string;
    }

    export interface ImageGenerationRequest {
        model: string;
        prompt: string;
        n?: number;
        size?: string;
        quality?: 'standard' | 'hd';
        output_format?: 'jpeg' | 'webp' | 'png';
        image_settings?: ImageGenerationSettings;
    }

    export interface ImageGenerationResponse {
        created: number;
        data: ({ b64_json: string })[];
    }

    export interface ImageGenerationError {
        error: string;
    }

    type Modality = 'text' | 'image';

    interface BaseModelInfo {
        id: string;
        object: 'model';
        created: number;
        owned_by: string;
        supported_parameters: string[];
        architecture?: {
            input_modalities: Modality[];
            output_modalities: Modality[];
        };
    }

    export interface ModelInfoText extends BaseModelInfo {
        context_length: number;
        top_provider: {
            context_length: number;
            max_completion_tokens: number;
            is_moderated: boolean;
        };
    }

    export interface ModelInfoImage extends BaseModelInfo {
    }

    export type ModelInfo = ModelInfoText | ModelInfoImage;

    const isTextModel = (model: ModelInfo): model is ModelInfoText => ('context_length' in model);
    const isImageModel = (model: ModelInfo): model is ModelInfoImage => Boolean(
        !isTextModel(model) &&
        model.architecture &&
        (model.architecture.output_modalities).includes('image')
    );

    export interface ModelsResponse<T extends ModelInfo = ModelInfo> {
        object: 'list';
        data: T[];
    }

    interface CountTokensRequestString {
        model: string;
        input: string;
    }

    interface CountTokensRequestMessages {
        model: string;
        input: LLM.ChatMessage[];
        tools?: LLM.Tool[];
        add_generation_prompt?: boolean;
        reasoning?: {
            effort?: string;
            exclude?: boolean;
            max_tokens?: number;
        };
    }

    export type CountTokensRequest = CountTokensRequestString | CountTokensRequestMessages;

    export interface CountTokensResponse {
        object: 'response.input_tokens';
        input_tokens: number;
    }

    async function request<T>(connection: Connection, path: string, method: string = 'GET', body?: unknown): Promise<T> {
        const headers = {
            'Content-Type': 'application/json',
            'Authorization': `Bearer ${connection.apiKey}`,
        };

        const url = new URL(connection.url);
        url.pathname = path;

        const response = await fetch(url, {
            method,
            headers,
            body: body ? JSON.stringify(body) : undefined,
        });
        if (!response.ok) {
            let text = '';
            try {
                text = await response.text();
            } catch { }
            throw new Error(`HTTP error! status: ${response.status}, text: ${text}`);
        }
        return response.json();
    }

    async function* streamRequest<T>(connection: Connection, path: string, method: string = 'GET', body?: unknown): AsyncGenerator<T> {
        const headers = {
            'Content-Type': 'application/json',
            'Authorization': `Bearer ${connection.apiKey}`,
        };

        const url = new URL(connection.url);
        url.pathname = path;

        using sse = new SSE(url.toString(), {
            headers,
            method,
            payload: body ? JSON.stringify(body) : undefined,
        });

        const readable = new ReadableStream<string>({
            start: async (controller) => {
                sse.addEventListener('message', (e) => {
                    if (isMessageEvent(e)) {
                        if (e.data === '[DONE]') {
                            controller.close();
                        } else {
                            controller.enqueue(e.data);
                        }
                    }
                });

                let closed = false;
                const handleEnd = (_e?: unknown) => {
                    if (closed) return;
                    closed = true;
                    controller.close();
                };

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

        const reader = readable.getReader();

        while (true) {
            const { value, done } = await reader.read();
            if (done) {
                break;
            }
            try {
                yield JSON.parse(value);
            } catch {
                break;
            }
        }

        await reader.closed;
        sse.close();
    }

    function isMessageEvent(e: unknown): e is { data: string } {
        return e != null && typeof e === 'object' && 'data' in e && typeof e.data === 'string';
    }

    export async function getTextModels(connection: Connection): Promise<ModelsResponse<ModelInfoText>> {
        const response = await request<ModelsResponse>(connection, '/v1/models');

        response.data = response.data.filter(isTextModel);
        return response as ModelsResponse<ModelInfoText>;
    }

    export async function getImageModels(connection: Connection): Promise<ModelsResponse<ModelInfoImage>> {
        const response = await request<ModelsResponse>(connection, '/v1/models');

        response.data = response.data.filter(isImageModel);
        return response as ModelsResponse<ModelInfoImage>;
    }

    export async function countTokens(connection: Connection, body: CountTokensRequest) {
        return request<CountTokensResponse>(connection, '/v1/responses/input_tokens', 'POST', body);
    }

    export async function* generateStream(connection: Connection, config: ChatCompletionRequest) {
        yield* streamRequest<ChatCompletionChunk | ChatCompletionError>(connection, '/v1/chat/completions', 'POST', {
            ...config,
            stream: true,
        });
    }

    export async function generate(connection: Connection, config: ChatCompletionRequest) {
        return request<ChatCompletionResponse | ChatCompletionError>(connection, '/v1/chat/completions', 'POST', {
            ...config,
            stream: false,
        });
    }

    export async function generateImage(connection: Connection, config: ImageGenerationRequest) {
        return request<ImageGenerationResponse | ImageGenerationError>(connection, '/v1/images/generations', 'POST', config);
    }

    const SUMMARIZATION_PROMPT = `Summarize the following text concisely while preserving key information and meaning. {level}

Text:
{text}

Provide a clear and coherent summary:`;

    export type SummarizationLevel = 'sentence' | 'paragraph' | 'arbitrary';
    const LEVEL_INSTRUCTIONS: Record<SummarizationLevel, string> = {
        sentence: 'Summarize in exactly one sentence.',
        paragraph: 'Summarize in exactly one paragraph (2-4 sentences).',
        arbitrary: 'Summarize in a way you think is appropriate for the text length and complexity.',
    };

    export async function summarize(connection: Connection, model: string, text: string, level: SummarizationLevel = 'arbitrary'): Promise<string> {

        const prompt = SUMMARIZATION_PROMPT
            .replace('{text}', text)
            .replace('{level}', LEVEL_INSTRUCTIONS[level]);

        const response = await generate(connection, {
            model,
            messages: [{
                role: 'user',
                content: prompt,
            }],
            temperature: 0.3,
            max_tokens: 500,
        });

        if ('error' in response) {
            throw new Error(response.error);
        }

        return response.choices[0]?.message.content ?? '';
    }
}

export default LLM;