#nullable enable
using System;
using System.Threading;
using Cysharp.Threading.Tasks;
using Mochineko.ChatGPT_API.Memories;
using UnityEngine;

namespace Mochineko.ChatGPT_API.Samples
{
    /// <summary>
    /// A sample component to complete chat by ChatGPT API on Unity.
    /// </summary>
    public sealed class ChatCompletionSample : MonoBehaviour
    {
        /// <summary>
        /// API key generated by OpenAPI.
        /// </summary>
        [SerializeField] private string apiKey = string.Empty;

        /// <summary>
        /// System message to instruct assistant.
        /// </summary>
        [SerializeField, TextArea] private string systemMessage = string.Empty;

        /// <summary>
        /// Message sent to ChatGPT API.
        /// </summary>
        [SerializeField, TextArea] private string message = string.Empty;

        /// <summary>
        /// Max number of chat memory of queue.
        /// </summary>
        [SerializeField] private int maxMemoryCount = 20;

        private ChatCompletionAPIConnection? connection;
        private IChatMemory? memory;

        private void Start()
        {
            // API Key must be set.
            if (string.IsNullOrEmpty(apiKey))
            {
                Debug.LogError("OpenAI API key must be set.");
                return;
            }

            memory = new FiniteQueueChatMemory(maxMemoryCount);

            // Create instance of ChatGPTConnection with specifying chat model.
            connection = new ChatCompletionAPIConnection(
                apiKey,
                memory,
                systemMessage);
        }

        [ContextMenu(nameof(SendChat))]
        public void SendChat()
        {
            SendChatAsync(this.GetCancellationTokenOnDestroy()).Forget();
        }
        
        [ContextMenu(nameof(ClearChatMemory))]
        public void ClearChatMemory()
        {
            memory?.ClearAllMessages();
        }
        
        private async UniTask SendChatAsync(CancellationToken cancellationToken)
        {
            // Validations
            if (connection == null)
            {
                Debug.LogError($"[ChatGPT_API.Samples] Connection is null.");
                return;
            }

            if (string.IsNullOrEmpty(message))
            {
                Debug.LogError($"[ChatGPT_API.Samples] Chat content is empty.");
                return;
            }

            ChatCompletionResponseBody response;
            try
            {
                await UniTask.SwitchToThreadPool();
                
                // Create message by ChatGPT chat completion API.
                response = await connection.CompleteChatAsync(
                        message,
                        cancellationToken);
            }
            catch (Exception e)
            {
                // Exceptions should be caught.
                Debug.LogException(e);
                return;
            }

            await UniTask.SwitchToMainThread(cancellationToken);

            // Log chat completion result.
            Debug.Log($"[ChatGPT_API.Samples] Result:\n{response.ResultMessage}");
        }
    }
}

using TiktokenSharp;

private int CalculateTokenLength()
{
    string text = "A text that you want to calculate token length.";

    // Specify model name.
    Tiktoken tiktoken = TikToken.EncodingForModel("gpt-3.5-turbo");
    
    // Encoding is tokenizing.
    int[] tokens = tikToken.Encode(text);
    
    return tokens.Length;
}

var builder = new StringBuilder();
 // Receive enumerable from ChatGPT chat completion API.
var enumerable = await connection.CompleteChatAsStreamAsync(
    message,
    cancellationToken);

await foreach (var chunk in enumerable.WithCancellation(cancellationToken))
{
    // First chunk has only "role" element.
    if (chunk.Choices[0].Delta.Content is null)
    {
        Debug.Log($"[ChatGPT_API.Samples] Role:{chunk.Choices[0].Delta.Role}.");
        continue;
    }
    
    var delta = chunk.Choices[0].Delta.Content;
    builder.Append(delta);
    Debug.Log($"[ChatGPT_API.Samples] Delta:{delta}, Current:{builder}");
}

// Log chat completion result.
Debug.Log($"[ChatGPT_API.Samples] Completed: \n{builder}");