path: root/Software/Visual_Studio_22/Tango.Portal.Chat.Web/Services/LlmClient.cs

using System.Net.Http.Headers;
using System.Text;
using System.Text.Json;
using System.Text.Json.Nodes;
using Microsoft.Extensions.Options;
using Tango.Portal.Chat.Web.Models;
using System.Text.RegularExpressions;

namespace Tango.Portal.Chat.Web.Services
{
    public sealed class LlmClient
    {
        private static int MAX_HISTORY = 10;
        private readonly HttpClient _http;
        private readonly LlmOptions _opt;

        public sealed class AssistantRunResult
        {
            public string Answer { get; set; } = "";
            public string ThreadId { get; set; } = "";
        }

        public enum AssistantType
        {
            Data,
            Docs
        }

        public LlmClient(HttpClient http, IOptions<LlmOptions> opt)
        {
            _http = http;
            _opt = opt.Value;
        }

        public async Task<ProposeKqlResult> ProposeKqlAsync(String plannerPrompt, String plotySample,
    string question, string schemaJson, IEnumerable<ChatMessage>? history, CancellationToken ct = default)
        {
            var plan = _opt.Provider switch
            {
                LlmProvider.Claude => await ProposeKqlWithClaudeAsync(plannerPrompt, plotySample, question, schemaJson, history, ct),
                LlmProvider.OpenAI => await ProposeKqlWithOpenAIAsync(plannerPrompt, plotySample, question, schemaJson, history, ct),
                _ => await ProposeKqlWithOpenAIAsync(plannerPrompt, plotySample, question, schemaJson, history, ct) // Default to OpenAI
            };

            plan.Provider = _opt.Provider;
            return plan;
        }

        private async Task<ProposeKqlResult> ProposeKqlWithOpenAIAsync(String plannerPrompt, String plotySample,
    string question, string schemaJson, IEnumerable<ChatMessage>? history, CancellationToken ct)
        {
            var messages = new List<object> { new { role = "system", content = plannerPrompt } };

            if (history != null)
            {
                history = history.DistinctBy(x => x.Content);

                foreach (var m in history.TakeLast(MAX_HISTORY))
                    messages.Add(new { role = m.Role, content = CapString(m.Content, 1000), usedKql = m.UsedKql });
            }

            var schemaBlock = $"SCHEMA:\n{schemaJson}";
            messages.Add(new { role = "user", content = $"Question: {question}\n\n{schemaBlock}\n\n{plotySample}" });

            var payload = new
            {
                model = _opt.Model,
                temperature = _opt.Temperature,
                response_format = new { type = "json_object" },
                messages = messages
            };

            using var req = new HttpRequestMessage(HttpMethod.Post, _opt.Endpoint);
            req.Content = new StringContent(JsonSerializer.Serialize(payload), Encoding.UTF8, "application/json");

            req.Headers.Authorization = new AuthenticationHeaderValue("Bearer", _opt.ApiKey);

            using var resp = await _http.SendAsync(req, ct);
            resp.EnsureSuccessStatusCode();
            var body = await resp.Content.ReadAsStringAsync(ct);

            var root = JsonNode.Parse(body)!.AsObject();
            var content = root["choices"]![0]!["message"]!["content"]?.ToString() ?? "{}";
            content = StripCodeFences(content);

            var opts = new JsonSerializerOptions { PropertyNameCaseInsensitive = true };
            opts.Converters.Add(new FlexibleStringListConverter());
            var result = JsonSerializer.Deserialize<ProposeKqlResult>(content, opts)
                         ?? new ProposeKqlResult { Kql = "", Parameters = new() };
            return result;
        }

        public async Task<ProposeKqlResult> ProposeKqlWithClaudeAsync(String plannerPrompt, String plotySample,
    string question, string schemaJson, IEnumerable<ChatMessage>? history, CancellationToken ct)
        {
            var messages = new List<object>();

            if (history != null)
            {
                history = history.DistinctBy(x => x.Content);

                foreach (var m in history.TakeLast(MAX_HISTORY))
                {
                    messages.Add(new
                    {
                        role = m.Role == "assistant" ? "assistant" : "user",
                        content = CapString(m.Content, 1000)
                    });
                }
            }

            var schemaBlock = $"SCHEMA:\n{schemaJson}";
            var userMessage = $"Question: {question}\n\n{schemaBlock}\n\n{plotySample}\n\nPlease respond with valid JSON only.";
            messages.Add(new { role = "user", content = userMessage });

            var payload = new
            {
                model = !string.IsNullOrEmpty(_opt.ClaudeModel) ? _opt.ClaudeModel : "claude-3-5-sonnet-20241022",
                max_tokens = _opt.MaxTokens,
                temperature = _opt.Temperature,
                system = plannerPrompt,
                messages = messages
            };

            var endpoint = !string.IsNullOrEmpty(_opt.ClaudeEndpoint) ? _opt.ClaudeEndpoint : "https://api.anthropic.com/v1/messages";
            var apiKey = !string.IsNullOrEmpty(_opt.ClaudeApiKey) ? _opt.ClaudeApiKey : _opt.ApiKey;

            using var req = new HttpRequestMessage(HttpMethod.Post, endpoint);
            req.Content = new StringContent(JsonSerializer.Serialize(payload), Encoding.UTF8, "application/json");
            req.Headers.Add("x-api-key", apiKey);
            req.Headers.Add("anthropic-version", "2023-06-01");

            using var resp = await _http.SendAsync(req, ct);
            resp.EnsureSuccessStatusCode();
            var body = await resp.Content.ReadAsStringAsync(ct);

            var root = JsonNode.Parse(body)!.AsObject();
            var contentArray = root["content"]?.AsArray();
            var content = contentArray?[0]?["text"]?.ToString() ?? "{}";
            content = StripCodeFences(content);

            var opts = new JsonSerializerOptions { PropertyNameCaseInsensitive = true };
            opts.Converters.Add(new FlexibleStringListConverter());
            var result = JsonSerializer.Deserialize<ProposeKqlResult>(content, opts)
                         ?? new ProposeKqlResult { Kql = "", Parameters = new() };
            return result;
        }

        private static string StripCodeFences(string s)
        {
            if (string.IsNullOrWhiteSpace(s)) return s ?? string.Empty;
            var t = s.Trim();

            if (t.StartsWith("```"))
            {
                // Remove first line (the ```json or ``` block)
                var firstNl = t.IndexOf('\n');
                if (firstNl >= 0 && firstNl + 1 < t.Length)
                {
                    var inner = t.Substring(firstNl + 1);

                    // Remove trailing fence
                    var fence = inner.LastIndexOf("```", StringComparison.Ordinal);
                    if (fence >= 0)
                        inner = inner.Substring(0, fence);

                    return inner.Trim();
                }
            }
            return t;
        }


        public async Task<string> AnswerFromFactsAsync(string question, string factsJson, string kqlForDisplay, CancellationToken ct = default)
        {
            var system = string.Join("\n", new[] {
                "You are a precise analyst.",
                "Answer ONLY from the provided ADX facts. If insufficient, say so.",
                "Be explicit about the time range and columns used.",
                "Ink quantities are stored as nanoliters. Make them humanly readable by converting them to milliliters or liters depending on what makes more sense."
            });

            var user = $"Question: {question}\n\nFacts(JSON):\n{factsJson}\n\nKQL used:\n{kqlForDisplay}";

            var payload = new
            {
                model = _opt.Model,
                temperature = 0.2,
                messages = new object[] {
                    new { role = "system", content = system },
                    new { role = "user", content = user }
                }
            };

            using var req = new HttpRequestMessage(HttpMethod.Post, _opt.Endpoint);
            var json = JsonSerializer.Serialize(payload);
            req.Content = new StringContent(json, Encoding.UTF8, "application/json");

            req.Headers.Authorization = new AuthenticationHeaderValue("Bearer", _opt.ApiKey);

            using var resp = await _http.SendAsync(req, ct);
            resp.EnsureSuccessStatusCode();
            var body = await resp.Content.ReadAsStringAsync(ct);

            var root = JsonNode.Parse(body)!.AsObject();
            var content = root["choices"]![0]!["message"]!["content"]!.ToString();
            return content;
        }

        // Add once in your class
        private void AddOpenAIHeaders(HttpRequestMessage req)
        {
            // MUST be a standard OpenAI key (sk-...), not an Azure key
            req.Headers.Authorization =
                new System.Net.Http.Headers.AuthenticationHeaderValue("Bearer", _opt.ApiKey);

            // REQUIRED for Assistants v2
            req.Headers.Add("OpenAI-Beta", "assistants=v2");

            // If your org enforces projects or you want to scope, uncomment:
            // req.Headers.Add("OpenAI-Organization", "<org_xxx>");
            // req.Headers.Add("OpenAI-Project", "<proj_xxx>");
        }

        private static async Task<string> ReadBodyOrThrowAsync(HttpResponseMessage res, CancellationToken ct)
        {
            var body = await res.Content.ReadAsStringAsync(ct);
            if (!res.IsSuccessStatusCode)
                throw new HttpRequestException($"{(int)res.StatusCode} {res.ReasonPhrase}: {body}");
            return body;
        }


        public async Task<string> AnswerWithAssistantAsync(
            AssistantType assistant,
            string question,
            string factsJson,
            string kql,
            CancellationToken ct = default)
        {
            // 1) Create a thread (empty is fine)
            using var tReq = new HttpRequestMessage(HttpMethod.Post, "https://api.openai.com/v1/threads")
            {
                Content = new StringContent("{}", System.Text.Encoding.UTF8, "application/json")
            };
            AddOpenAIHeaders(tReq);
            var tBody = await ReadBodyOrThrowAsync(await _http.SendAsync(tReq, ct), ct);
            var threadId = System.Text.Json.JsonDocument.Parse(tBody).RootElement.GetProperty("id").GetString();

            // 2) Add a single user message with the three sections

            String messageText = String.Empty;

            if (assistant == AssistantType.Data)
            {
                messageText = $"Question:\n{question}\n\nFacts(JSON):\n{factsJson}\n\nKQL used:\n{kql}";
            }
            else
            {
                messageText = $"Question:\n{question}";
            }

            using var mReq = new HttpRequestMessage(HttpMethod.Post, $"https://api.openai.com/v1/threads/{threadId}/messages")
            {
                // v2 expects "content" to be an array of content parts
                Content = JsonContent.Create(new
                {
                    role = "user",
                    content = new object[] { new { type = "text", text = messageText } }
                })
            };
            AddOpenAIHeaders(mReq);
            await ReadBodyOrThrowAsync(await _http.SendAsync(mReq, ct), ct);

            // 3) Create a run targeting your Assistant (must have File Search enabled & schema.json attached)
            using var rReq = new HttpRequestMessage(HttpMethod.Post, $"https://api.openai.com/v1/threads/{threadId}/runs")
            {
                Content = JsonContent.Create(new
                {
                    assistant_id = assistant == AssistantType.Data ? _opt.AnswererAssistantId : _opt.DocsAssistantId,
                    // You can override instructions here if you ever need to:
                    // instructions = "..."
                })
            };
            AddOpenAIHeaders(rReq);
            var rBody = await ReadBodyOrThrowAsync(await _http.SendAsync(rReq, ct), ct);
            var runId = System.Text.Json.JsonDocument.Parse(rBody).RootElement.GetProperty("id").GetString();

            // 4) Poll run until completed
            while (true)
            {
                await Task.Delay(600, ct);
                using var gReq = new HttpRequestMessage(HttpMethod.Get, $"https://api.openai.com/v1/threads/{threadId}/runs/{runId}");
                AddOpenAIHeaders(gReq);
                var gBody = await ReadBodyOrThrowAsync(await _http.SendAsync(gReq, ct), ct);
                var root = System.Text.Json.JsonDocument.Parse(gBody).RootElement;
                var status = root.GetProperty("status").GetString();
                if (status == "completed") break;
                if (status == "failed" || status == "cancelled" || status == "expired")
                {
                    var lastError = root.TryGetProperty("last_error", out var le) ? le.ToString() : "unknown";
                    throw new Exception($"Assistant run {status}: {lastError}");
                }
            }

            // 5) Fetch messages and return the latest assistant text
            using var lReq = new HttpRequestMessage(HttpMethod.Get, $"https://api.openai.com/v1/threads/{threadId}/messages");
            AddOpenAIHeaders(lReq);
            var lBody = await ReadBodyOrThrowAsync(await _http.SendAsync(lReq, ct), ct);
            using var doc = System.Text.Json.JsonDocument.Parse(lBody);

            // Messages are returned most-recent-first; take the first assistant message
            foreach (var msg in doc.RootElement.GetProperty("data").EnumerateArray())
            {
                var role = msg.GetProperty("role").GetString();
                if (role == "assistant")
                {
                    foreach (var part in msg.GetProperty("content").EnumerateArray())
                    {
                        if (part.GetProperty("type").GetString() == "text")
                            return part.GetProperty("text").GetProperty("value").GetString() ?? "";
                    }
                }
            }

            return "(no assistant message found)";
        }

        public async Task<AssistantRunResult> AnswerWithAssistantAsync(
            AssistantType assistant,
            string question,
            string factsJson,
            string kql,
            string? threadId,
            CancellationToken ct = default)
        {
            // 1) Use existing thread or create a new one
            if (string.IsNullOrEmpty(threadId))
            {
                using var tReq = new HttpRequestMessage(HttpMethod.Post, "https://api.openai.com/v1/threads")
                { Content = new StringContent("{}", Encoding.UTF8, "application/json") };
                AddOpenAIHeaders(tReq);
                var tBody = await ReadBodyOrThrowAsync(await _http.SendAsync(tReq, ct), ct);
                threadId = System.Text.Json.JsonDocument.Parse(tBody).RootElement.GetProperty("id").GetString();
            }

            // 2) Add the user message (same text you build today)
            var messageText = assistant == AssistantType.Data
                ? $"Question:\n{question}\n\nFacts(JSON):\n{factsJson}\n\nKQL used:\n{kql}"
                : $"Question:\n{question}";

            using (var mReq = new HttpRequestMessage(HttpMethod.Post, $"https://api.openai.com/v1/threads/{threadId}/messages")
            { Content = JsonContent.Create(new { role = "user", content = new object[] { new { type = "text", text = messageText } } }) })
            {
                AddOpenAIHeaders(mReq);
                await ReadBodyOrThrowAsync(await _http.SendAsync(mReq, ct), ct);
            }

            // 3) Run with the correct assistant id (unchanged logic)
            using (var rReq = new HttpRequestMessage(HttpMethod.Post, $"https://api.openai.com/v1/threads/{threadId}/runs")
            { Content = JsonContent.Create(new { assistant_id = assistant == AssistantType.Data ? _opt.AnswererAssistantId : _opt.DocsAssistantId }) })
            {
                AddOpenAIHeaders(rReq);
                var rBody = await ReadBodyOrThrowAsync(await _http.SendAsync(rReq, ct), ct);
                var runId = System.Text.Json.JsonDocument.Parse(rBody).RootElement.GetProperty("id").GetString();

                // Poll until completed (same as your existing loop)
                while (true)
                {
                    await Task.Delay(2000, ct);
                    using var gReq = new HttpRequestMessage(HttpMethod.Get, $"https://api.openai.com/v1/threads/{threadId}/runs/{runId}");
                    AddOpenAIHeaders(gReq);
                    var gBody = await ReadBodyOrThrowAsync(await _http.SendAsync(gReq, ct), ct);
                    var root = System.Text.Json.JsonDocument.Parse(gBody).RootElement;
                    var status = root.GetProperty("status").GetString();
                    if (status == "completed") break;
                    if (status == "failed" || status == "cancelled" || status == "expired")
                    {
                        var lastError = root.TryGetProperty("last_error", out var le) ? le.ToString() : "unknown";
                        throw new Exception($"Assistant run {status}: {lastError}");
                    }
                }
            }

            // 4) Fetch the latest assistant text and return it with the thread id
            using var lReq = new HttpRequestMessage(HttpMethod.Get, $"https://api.openai.com/v1/threads/{threadId}/messages");
            AddOpenAIHeaders(lReq);
            var lBody = await ReadBodyOrThrowAsync(await _http.SendAsync(lReq, ct), ct);
            using var doc = System.Text.Json.JsonDocument.Parse(lBody);
            foreach (var msg in doc.RootElement.GetProperty("data").EnumerateArray())
            {
                if (msg.GetProperty("role").GetString() == "assistant")
                {
                    foreach (var part in msg.GetProperty("content").EnumerateArray())
                        if (part.GetProperty("type").GetString() == "text")
                        {
                            var raw = part.GetProperty("text").GetProperty("value").GetString() ?? "";
                            var cleaned = assistant == AssistantType.Docs ? StripCitations(raw) : raw;
                            return new AssistantRunResult
                            {
                                Answer = cleaned,
                                ThreadId = threadId!
                            };
                        }
                }
            }
            return new AssistantRunResult { Answer = "(no assistant message found)", ThreadId = threadId! };
        }

        private static string StripCitations(string s)
        {
            if (string.IsNullOrWhiteSpace(s)) return s ?? string.Empty;
            // Remove any inline citation markers like:    or  
            return Regex.Replace(s, @"\s*【[^】]*】", string.Empty);
        }

        private static string CapString(string input, int maxLength)
        {
            if (string.IsNullOrEmpty(input) || maxLength <= 0)
                return string.Empty;

            if (input.Length <= maxLength)
                return input;

            return input.Substring(0, maxLength) + "…";
        }
    }
}