"use strict";

import {__} from "@wordpress/i18n";
import OpenAI from "openai";

export default class Grok {
    /**
     * API client credentials
     * @type {{baseURL: string|null, apiKey: string|null, model: string|null}}
     */
    #credentials = {
        baseURL: null,
        apiKey: null,
        model: null,
    };

    /**
     * API client
     *
     * @type {OpenAI|null}
     */
    #client = null;

    /**
     * Construct
     *
     * @param {string} baseUrl API base URL
     * @param {string} apiKey API key
     * @param {string} model AI client model
     */
    constructor({baseUrl = "https://api.x.ai/v1", apiKey, model}) {
        /**
         * Init properties
         */
        this.#credentials.baseURL = baseUrl;
        this.#credentials.apiKey = apiKey;
        this.#credentials.model = model;
        // Init API client
        this.#initApiClient();
    }

    /**
     * Initialize the API client
     */
    #initApiClient() {
        // Set configs to init the AI
        const config = {
            apiKey: this.#credentials.apiKey,
            baseURL: this.#credentials.baseURL,
        };
        // Init API client
        this.#client = new OpenAI({
            ...config,
            dangerouslyAllowBrowser: true
        });
    }

    /**
     * Is the API client initialized successfully
     *
     * @returns {boolean}
     */
    isApiClientInitialized() {
        return !!this.#client;
    }

    /**
     * Generate content
     *
     * @param {string} systemInfo AI description
     * @param {string} prompt Request
     * @returns {{status: boolean,data: string}}
     */
    async generateContent({systemInfo = "You are a helpful assistant.", prompt}) {
        try {
            const result = await this.chatCompletion({systemInfo, content: prompt});
            return {
                status: result.status,
                data: result.status ? result.data?.message.content : result.data
            };
        } catch (e) {
            return {
                status: false,
                data: e.message
            };
        }
    }

    /**
     * The chat completion endpoint request
     *
     * @param {string} systemInfo System information
     * @param {string} content Content
     * @returns {Promise<{status: boolean, data: any}>}
     */
    async chatCompletion({systemInfo = "You are a helpful assistant.", content}) {
        try {
            if (!this.isApiClientInitialized()) {
                return {
                    status: false,
                    data: __("The Grok API client isn't initialized.", 'limb-chatbot')
                };
            }
            const completion = await this.#client.chat.completions.create({
                messages: [
                    {role: "system", content: systemInfo},
                    {role: "user", content},
                ],
                model: this.#credentials.model,
            });

            // TODO: track usage
            // {
            //     "id": "chatcmpl-123",
            //     "object": "chat.completion",
            //     "created": 1677652288,
            //     "model": "grok-beta",
            //     "choices": [{
            //         "index": 0,
            //         "message": {
            //             "role": "assistant",
            //             "content": "Hello!"
            //         },
            //         "finish_reason": "stop"
            //     }],
            //     "usage": {
            //         "prompt_tokens": 9,
            //         "completion_tokens": 12,
            //         "total_tokens": 21
            //     }
            // }

            return {
                status: true,
                data: completion.choices[0]
            };
        } catch (e) {
            return {
                status: false,
                data: e.message
            };
        }
    }

    /**
     * Chat with conversation history
     *
     * @param {Array} messages Array of messages in format [{role: "user", content: "text"}, {role: "assistant", content: "text"}]
     * @param {string} systemInfo System information (optional)
     * @return {Promise<{status: boolean, data: string}>}
     */
    async chat({messages, systemInfo = null}) {
        try {
            if (!this.isApiClientInitialized()) {
                return {
                    status: false,
                    data: __("The Grok API client isn't initialized.", 'limb-chatbot')
                };
            }

            // Prepend system message if provided
            const messagesWithSystem = systemInfo
                ? [{role: "system", content: systemInfo}, ...messages]
                : messages;

            const completion = await this.#client.chat.completions.create({
                messages: messagesWithSystem,
                model: this.#credentials.model,
            });

            // TODO: track usage
            // completion.usage
            // { prompt_tokens: 25, completion_tokens: 25, total_tokens: 50 }

            return {
                status: true,
                data: completion.choices[0]?.message?.content || ""
            };
        } catch (e) {
            return {
                status: false,
                data: e.message
            };
        }
    }

    /**
     * Stream chat completion
     *
     * @param {Array} messages Array of messages
     * @param {string} systemInfo System information (optional)
     * @param {Function} onChunk Callback function for each chunk
     * @return {Promise<{status: boolean, data: string}>}
     */
    async streamChat({messages, systemInfo = null, onChunk}) {
        try {
            if (!this.isApiClientInitialized()) {
                return {
                    status: false,
                    data: __("The Grok API client isn't initialized.", 'limb-chatbot')
                };
            }

            // Prepend system message if provided
            const messagesWithSystem = systemInfo
                ? [{role: "system", content: systemInfo}, ...messages]
                : messages;

            const stream = await this.#client.chat.completions.create({
                messages: messagesWithSystem,
                model: this.#credentials.model,
                stream: true,
            });

            let fullText = "";

            for await (const chunk of stream) {
                const content = chunk.choices[0]?.delta?.content || "";
                if (content) {
                    fullText += content;
                    if (onChunk && typeof onChunk === 'function') {
                        onChunk(content);
                    }
                }
            }

            return {
                status: true,
                data: fullText
            };
        } catch (e) {
            return {
                status: false,
                data: e.message
            };
        }
    }

    /**
     * The embedding endpoint request
     *
     * @param {string} input Input text
     * @param {string} model Embedding model (optional, defaults to chat model)
     * @returns {Promise<{status: boolean, data: any}>}
     */
    async embedding({input, model = null}) {
        try {
            if (!this.isApiClientInitialized()) {
                return {
                    status: false,
                    data: __("The Grok API client isn't initialized.", 'limb-chatbot')
                };
            }

            const embedding = await this.#client.embeddings.create({
                model: model || this.#credentials.model,
                input
            });

            return {
                status: true,
                data: embedding
            };
        } catch (e) {
            return {
                status: false,
                data: e.message
            };
        }
    }
}
