using BLAIzor.Models;
using System.Text.Json;
using System.Text;

namespace BLAIzor.Services
{
    public class CerebrasAPIService
    {       
        private readonly IConfiguration _configuration;
        private readonly HttpClient _httpClient;
        private static Action<string, string>? _callback;
        private static Action<string>? _onComplete;
        private static Action<string, string>? _onError;

        private const string CerebrasEndpoint = "https://api.cerebras.ai/v1/chat/completions";

        public CerebrasAPIService(IConfiguration configuration, HttpClient httpClient)
        {
            _configuration = configuration;
            _httpClient = httpClient;
            _httpClient.DefaultRequestHeaders.Add("Authorization", $"Bearer {GetApiKey()}");
        }

        private string GetApiKey() =>
            _configuration?.GetSection("Cerebras")?.GetValue<string>("ApiKey") ?? string.Empty;

        private string GetModelName() =>
            _configuration?.GetSection("Cerebras")?.GetValue<string>("Model") ?? string.Empty;

        public void RegisterCallback(Action<string, string> callback, Action<string> onCompleteCallback, Action<string, string> onErrorCallback)
        {
            _callback = callback;
            _onComplete = onCompleteCallback;
            _onError = onErrorCallback;
        }


        public async Task<string> GetSimpleCerebrasResponse(string sessionId, string systemMessage, string userMessage, string? assistantMessage = null, int mode = -1)
        {            

            string modelName = GetModelName();
            //if (mode >= 0)
            //{
            //    if (mode == 0)
            //    {
            //        modelName = "llama-3.3-70b";
            //    }
            //}
            //else
            //{
            //    modelName = "llama-3.3-70b";
            //}

            var requestBody = new ChatGPTRequest
            {
                Model = modelName,
                Temperature = 0.2,
                Messages = assistantMessage == null || assistantMessage == string.Empty
                    ? new[]
                    {
                        new Message { Role = "system", Content = systemMessage },
                        new Message { Role = "user", Content = userMessage }
                    }
                    : new[]
                    {
                        new Message { Role = "system", Content = systemMessage },
                        new Message { Role = "assistant", Content = assistantMessage },
                        new Message { Role = "user", Content = userMessage }
                    },
                Stream = false
            };

            var requestJson = JsonSerializer.Serialize(requestBody, new JsonSerializerOptions
            {
                PropertyNamingPolicy = JsonNamingPolicy.CamelCase
            });

            var requestContent = new StringContent(requestJson, Encoding.UTF8, "application/json");

            using var response = await _httpClient.PostAsync(CerebrasEndpoint, requestContent);
            response.EnsureSuccessStatusCode();

            using var responseStream = await response.Content.ReadAsStreamAsync();
            using var document = await JsonDocument.ParseAsync(responseStream);
            var inputTokens = document.RootElement.GetProperty("usage").GetProperty("prompt_tokens").GetInt32();
            var outputTokens = document.RootElement.GetProperty("usage").GetProperty("completion_tokens").GetInt32();
            var sum = inputTokens + outputTokens;
            Console.WriteLine($"USAGE STATS - Tokens: {inputTokens.ToString()} + {outputTokens.ToString()} = {sum.ToString()}");

            return document.RootElement
                           .GetProperty("choices")[0]
                           .GetProperty("message")
                           .GetProperty("content")
                           .GetString() ?? "No response";
        }

        public async Task<string> GetSimpleCerebrasResponseNoSession(string systemMessage, string userMessage, string? assistantMessage = null, int mode = -1)
        {
            string modelName = GetModelName();
            //if (mode >= 0)
            //{
            //    if (mode == 0)
            //    {
            //        modelName = "llama-3.3-70b";
            //    }
            //}
            //else
            //{
            //    modelName = "llama-3.3-70b";
            //}

            var requestBody = new ChatGPTRequest
            {
                Model = modelName,
                Temperature = 0.2,
                Messages = assistantMessage == null || assistantMessage == string.Empty
                    ? new[]
                    {
                        new Message { Role = "system", Content = systemMessage },
                        new Message { Role = "user", Content = userMessage }
                    }
                    : new[]
                    {
                        new Message { Role = "system", Content = systemMessage },
                        new Message { Role = "assistant", Content = assistantMessage },
                        new Message { Role = "user", Content = userMessage }
                    },
                Stream = false
            };

            var requestJson = JsonSerializer.Serialize(requestBody, new JsonSerializerOptions
            {
                PropertyNamingPolicy = JsonNamingPolicy.CamelCase
            });

            var requestContent = new StringContent(requestJson, Encoding.UTF8, "application/json");

            using var response = await _httpClient.PostAsync(CerebrasEndpoint, requestContent);
            response.EnsureSuccessStatusCode();

            using var responseStream = await response.Content.ReadAsStreamAsync();
            using var document = await JsonDocument.ParseAsync(responseStream);

            return document.RootElement
                           .GetProperty("choices")[0]
                           .GetProperty("message")
                           .GetProperty("content")
                           .GetString() ?? "No response";
        }

        public async Task<string> GetCerebrasStreamedResponse(string sessionId, string systemMessage, string userMessage, string? assistantMessage = null, int mode = -1)
        {
            string modelName = GetModelName();
            //if (mode >= 0)
            //{
            //    if (mode == 0)
            //    {
            //        modelName = "llama-3.3-70b";
            //    }
            //}
            //else
            //{
            //    modelName = "llama-3.3-70b";
            //}

            var requestBody = new ChatGPTRequest
            {
                Model = modelName,
                Temperature = 0.2,
                Messages = assistantMessage == null
                    ? new[]
                    {
                        new Message { Role = "system", Content = systemMessage },
                        new Message { Role = "user", Content = userMessage }
                    }
                    : new[]
                    {
                        new Message { Role = "system", Content = systemMessage },
                        new Message { Role = "assistant", Content = assistantMessage },
                        new Message { Role = "user", Content = userMessage }
                    },
                Stream = true
            };

            var requestJson = JsonSerializer.Serialize(requestBody, new JsonSerializerOptions
            {
                PropertyNamingPolicy = JsonNamingPolicy.CamelCase
            });

            var requestContent = new StringContent(requestJson, Encoding.UTF8, "application/json");
            using var httpRequest = new HttpRequestMessage(HttpMethod.Post, CerebrasEndpoint)
            {
                Content = requestContent
            };

            using var response = await _httpClient.SendAsync(httpRequest, HttpCompletionOption.ResponseHeadersRead);
            response.EnsureSuccessStatusCode();

            var stringBuilder = new StringBuilder();
            using var responseStream = await response.Content.ReadAsStreamAsync();
            using var reader = new StreamReader(responseStream);

            try
            {
                while (!reader.EndOfStream)
                {
                    var line = await reader.ReadLineAsync();

                    if (string.IsNullOrWhiteSpace(line) || !line.StartsWith("data: ")) continue;

                    var jsonResponse = line.Substring(6);

                    // ✅ Detect explicit end of stream
                    if (jsonResponse == "[DONE]")
                    {
                        _onComplete?.Invoke(sessionId); // Optional: notify stream end
                        break;
                    }

                    try
                    {
                        using var jsonDoc = JsonDocument.Parse(jsonResponse);

                        if (jsonDoc.RootElement.TryGetProperty("choices", out var choices) &&
                            choices[0].TryGetProperty("delta", out var delta) &&
                            delta.TryGetProperty("content", out var contentElement))
                        {
                            var content = contentElement.GetString();

                            if (!string.IsNullOrEmpty(content))
                            {
                                stringBuilder.Append(content);
                                _callback?.Invoke(sessionId, stringBuilder.ToString());
                            }
                        }
                    }
                    catch (JsonException ex)
                    {
                        _onError?.Invoke(sessionId, $"Malformed JSON: {ex.Message}");
                        break; // Optionally stop stream
                    }
                }

                // ✅ Check for unexpected end (in case no [DONE])
                if (reader.EndOfStream && !stringBuilder.ToString().EndsWith("[DONE]"))
                {
                    _onError?.Invoke(sessionId, "Unexpected end of stream");
                }
            }
            catch (Exception ex)
            {
                _onError?.Invoke(sessionId, $"Exception: {ex.Message}");
            }


            return stringBuilder.ToString();
        }
    }
}