using System.ClientModel;
using System.Net.Http.Json;
using System.Text.Json.Serialization;
using OpenAI;
using OpenAI.Chat;

namespace TelegramBot;

public class OpenAiAgent {
    private ApiKeyCredential apikey;
    private OpenAIClient oaiClient;
    private ChatClient chatClient;
    private HttpClient httpClient;
    
    private Dictionary<(long, Actor), List<(ChatMessage, int)>> oaiChats = new();
    
    public int ContextSize { get; set; } = 4096;
    
    public OpenAiAgent(string baseUrl, string apiKey, string model) {
        OpenAIClientOptions options = new OpenAIClientOptions() {
            Endpoint = new(baseUrl),
            NetworkTimeout = new(0, 0, 15)
        };

        apikey = new(apiKey);
        oaiClient = new(apikey, options);
        chatClient = oaiClient.GetChatClient(model);

        httpClient = new() {
            BaseAddress = new(baseUrl),
            Timeout = new(0, 0, 15)
        };
        
        Console.WriteLine(
            $"""
                Base URL: {baseUrl}
                Model: {model}
                API Key: {apiKey}
             """);
    }

    public int GetTokenLenght(string message) {
        TokenizeRequest request = new(message);
        var req = httpClient.PostAsync("tokenize", JsonContent.Create(request));
        req.Wait();
        var response = req.Result.Content.ReadFromJsonAsync<TokenizeResponseBase>();
        response.Wait();
        return response.Result.Tokens.Length;
    }

    public void ChatHistoryAppend(Actor actor, long chatId, string message) {
        //get the chat from the dictionary
        var chat = GetChatHistory((chatId, actor));
        //get the token lenght of the message
        var tokenLenght = GetTokenLenght(message);
        //create a new chat message
        switch (actor) {
            case Actor.User:
                chat.Add((new UserChatMessage(message), tokenLenght));
                break;
            case Actor.Krolik:
                chat.Add((new AssistantChatMessage(message), tokenLenght));
                break;
            case Actor.Nemesis:
                chat.Add((new AssistantChatMessage(message), tokenLenght));
                break;
            case Actor.System:
                chat.Add((new SystemChatMessage(message), tokenLenght));
                break;
        }
    }

    public List<(ChatMessage, int)> GetChatHistory((long, Actor) key) => GetChatHistory(key.Item1, key.Item2);
    
    public List<(ChatMessage, int)> GetChatHistory(long chatId, Actor actor) {
        oaiChats.TryGetValue((chatId, actor), out var chat);
        if(chat == null) {
            AddChatToDictionary(chatId, actor);
            chat = oaiChats[(chatId, actor)];
        }
        return chat!;
    }

    public string GetChatResponse(long chatId, Agent agent) {
        int currentContextSize = agent.SystemPromptLength;
        List<ChatMessage> chatHistory = new();
        chatHistory.Add(new SystemChatMessage(agent.SystemPrompt));
        
        //Fetch the chat history from the dictionary trimming to the context size
        var history = GetChatHistory(chatId, agent.Actor).ToList();
        history.Reverse();
        //Add the chat history to the list until the context size is reached
        foreach (var (message, tokenLenght) in history) {
            if (currentContextSize + tokenLenght > ContextSize) break;
            chatHistory.Add(message);
            currentContextSize += tokenLenght;
        }
        //Reverse the chat history to get the correct order
        chatHistory.Reverse(1, chatHistory.Count - 1);
        
        var completion = chatClient.CompleteChat(chatHistory).Value.Content[0].Text;

        //Add the response to the chat history
        ChatHistoryAppend(agent.Actor, chatId, completion);
        
        return completion;
    }

    public void AddChatToDictionary(long id) {
        AddChatToDictionary(id, Actor.Krolik);
        AddChatToDictionary(id, Actor.Nemesis);
    }

    public void AddChatToDictionary(long id, Actor actor) {
        //Check if the chat already exists
        if (oaiChats.ContainsKey((id, actor))) return;
        //Create a new chat object
        var chat = new List<(ChatMessage, int)>();
        //chat.Add(new SystemChatMessage(nemesisPrompt));
        //add the entry to the dictionary
        oaiChats.Add((id, actor), chat);
    }
    
    public void ResetChat(long chatId) {
        //Remove the chat from the dictionary
        oaiChats.Where(x => x.Key.Item1 == chatId).ToList().ForEach(x => oaiChats.Remove(x.Key));
        //Add the chat back to the dictionary
        AddChatToDictionary(chatId);
    }
}

public struct TokenizeRequest(string content) {
    [JsonPropertyName("content")] public string Content { get; set; } = content;
    [JsonPropertyName("add_special")] public bool AddSpecial { get; set; } = false;
    [JsonPropertyName("with_pieces")] public bool WithPieces { get; set; } = false;
}

public struct TokenizeResponseBase(int[] tokens) {
    [JsonPropertyName("tokens")] public int[] Tokens { get; set; } = tokens;
}