ChatService.cs

using System.Runtime.CompilerServices;
using System.Text;
using System.Text.Json;
using System.Text.Json.Serialization;
using KaiChat.Models;

namespace KaiChat.Services;

public class ChatService : IChatService
{
    private readonly HttpClient _http;
    private readonly KaiChatConfig _config;
    private readonly ILogger<ChatService> _logger;

    public ChatService(HttpClient http, KaiChatConfig config, ILogger<ChatService> logger)
    {
        _http = http;
        _config = config;
        _logger = logger;
    }

    public async Task<string> ChatAsync(
        List<ChatMessage> messages,
        CancellationToken ct = default,
        int? maxTokens = null,
        bool enableTools = false,
        bool requireToolCall = false,
        bool enableThinking = true,
        bool enableAutonomyTools = false)
    {
        var (json, httpRequest) = BuildRequest(messages, stream: false, maxTokens, enableTools, requireToolCall, enableThinking, enableAutonomyTools);

        var response = await _http.SendAsync(httpRequest, ct);
        response.EnsureSuccessStatusCode();

        var body = await response.Content.ReadAsStringAsync(ct);
        using var doc = JsonDocument.Parse(body);
        return doc.RootElement
            .GetProperty("choices")[0]
            .GetProperty("message")
            .GetProperty("content")
            .GetString() ?? "";
    }

    public async IAsyncEnumerable<string> StreamChatAsync(
        List<ChatMessage> messages,
        [EnumeratorCancellation] CancellationToken ct = default,
        int? maxTokens = null,
        bool enableTools = false,
        bool requireToolCall = false,
        bool enableThinking = true,
        bool enableAutonomyTools = false)
    {
        var (json, httpRequest) = BuildRequest(messages, stream: true, maxTokens, enableTools, requireToolCall, enableThinking, enableAutonomyTools);

        _logger.LogInformation("Streaming chat to {Endpoint} with model {Model}", _config.ApiEndpoint, _config.Model);

        var response = await _http.SendAsync(
            httpRequest,
            HttpCompletionOption.ResponseHeadersRead,
            ct);

        response.EnsureSuccessStatusCode();

        var stream = await response.Content.ReadAsStreamAsync(ct);
        using var reader = new StreamReader(stream);
        var toolCallDeltas = new SortedDictionary<int, NativeToolCallDelta>();

        async IAsyncEnumerable<string> FlushToolCalls()
        {
            foreach (var call in toolCallDeltas.Values)
            {
                if (string.IsNullOrWhiteSpace(call.Name)) continue;
                var payload = JsonSerializer.Serialize(new NativeToolCall(
                    string.IsNullOrWhiteSpace(call.Id) ? Guid.NewGuid().ToString("N")[..12] : call.Id,
                    call.Name,
                    call.Arguments.ToString()));
                yield return $"T|{payload}";
                await Task.CompletedTask;
            }
        }

        while (!ct.IsCancellationRequested)
        {
            var line = await reader.ReadLineAsync(ct);
            if (line is null)
            {
                await foreach (var toolCall in FlushToolCalls())
                    yield return toolCall;
                yield break;
            }
            if (line.Length == 0) continue;
            if (!line.StartsWith("data: ")) continue;

            var data = line[6..];
            if (data == "[DONE]")
            {
                await foreach (var toolCall in FlushToolCalls())
                    yield return toolCall;
                yield break;
            }

            string? token = null;
            string? reasoningToken = null;
            int? promptTokens = null;
            int? completionTokens = null;

            try
            {
                using var doc = JsonDocument.Parse(data);
                var root = doc.RootElement;

                // Check for usage object (sent by DeepSeek in the final chunk)
                if (root.TryGetProperty("usage", out var usage))
                {
                    if (usage.TryGetProperty("prompt_tokens", out var pt))
                        promptTokens = pt.GetInt32();
                    if (usage.TryGetProperty("completion_tokens", out var ctVal))
                        completionTokens = ctVal.GetInt32();
                }

                var choice = root.GetProperty("choices")[0];
                var delta = choice.GetProperty("delta");

                if (delta.TryGetProperty("reasoning_content", out var reasoning))
                    reasoningToken = reasoning.GetString();

                if (delta.TryGetProperty("content", out var content))
                    token = content.GetString();

                if (delta.TryGetProperty("tool_calls", out var toolCalls) && toolCalls.ValueKind == JsonValueKind.Array)
                    AccumulateToolCallDeltas(toolCallDeltas, toolCalls);
            }
            catch (Exception ex)
            {
                _logger.LogWarning("Failed to parse streaming chunk: {Ex}", ex.Message);
            }

            // Yield usage info if present
            if (promptTokens.HasValue || completionTokens.HasValue)
                yield return $"U|{promptTokens ?? 0}|{completionTokens ?? 0}";

            if (!string.IsNullOrEmpty(reasoningToken))
                yield return $"R|{reasoningToken}";

            if (token is not null)
                yield return $"C|{token}";
        }
    }

    private (string Json, HttpRequestMessage Request) BuildRequest(
        List<ChatMessage> messages, bool stream, int? maxTokens, bool enableTools, bool requireToolCall, bool enableThinking, bool enableAutonomyTools)
    {
        var requestBody = new Dictionary<string, object?>
        {
            ["model"] = _config.Model,
            ["messages"] = messages.Select(ToApiMessage).ToList(),
            ["stream"] = stream,
            ["max_tokens"] = maxTokens ?? _config.GetDefaultCompletionTokenLimit()
        };

        if (enableTools)
        {
            requestBody["tools"] = BuildToolDefinitions(enableAutonomyTools);
            requestBody["tool_choice"] = "auto";
        }

        if (!enableThinking && SupportsThinkingSwitch())
            requestBody["thinking"] = new { type = "disabled" };

        var json = JsonSerializer.Serialize(requestBody, new JsonSerializerOptions
        {
            DefaultIgnoreCondition = JsonIgnoreCondition.WhenWritingNull
        });

        var httpRequest = new HttpRequestMessage(HttpMethod.Post, _config.ApiEndpoint)
        {
            Content = new StringContent(json, Encoding.UTF8, "application/json")
        };
        httpRequest.Headers.Authorization =
            new System.Net.Http.Headers.AuthenticationHeaderValue("Bearer", _config.ApiKey);

        return (json, httpRequest);
    }

    private bool SupportsThinkingSwitch()
    {
        return _config.Model.StartsWith("deepseek-v4", StringComparison.OrdinalIgnoreCase);
    }

    private static object ToApiMessage(ChatMessage message)
    {
        var role = (message.Role ?? "").Trim().ToLowerInvariant();

        if (role == "tool")
        {
            if (!string.IsNullOrWhiteSpace(message.Tool?.ToolCallId))
            {
                return new Dictionary<string, object?>
                {
                    ["role"] = "tool",
                    ["content"] = message.Content,
                    ["tool_call_id"] = message.Tool.ToolCallId
                };
            }

            return new Dictionary<string, object?>
            {
                ["role"] = "system",
                ["content"] = message.Content
            };
        }

        var payload = new Dictionary<string, object?>
        {
            ["role"] = role,
            ["content"] = message.Content
        };

        var toolCalls = BuildApiToolCalls(message.ToolCalls);
        if (role == "assistant" && toolCalls.Count > 0)
        {
            payload["reasoning_content"] = message.Thinking ?? "";
            payload["tool_calls"] = toolCalls;
        }

        return payload;
    }

    private static List<Dictionary<string, object?>> BuildApiToolCalls(IEnumerable<StoredToolCall> toolCalls)
    {
        return toolCalls
            .Where(call => !string.IsNullOrWhiteSpace(call.Id)
                && !string.IsNullOrWhiteSpace(call.Name)
                && !string.IsNullOrWhiteSpace(call.Arguments))
            .Select(call => new Dictionary<string, object?>
            {
                ["id"] = call.Id,
                ["type"] = "function",
                ["function"] = new Dictionary<string, object?>
                {
                    ["name"] = call.Name,
                    ["arguments"] = call.Arguments
                }
            })
            .ToList();
    }

    private static object[] BuildToolDefinitions(bool enableAutonomyTools)
    {
        var toolAreas = enableAutonomyTools
            ? "files, search, memory, Nightscout, the local KaiTimer dashboard/wheel/milestones/counters, autonomy production guidance, instruction rewrites, and explicit git commits"
            : "files, search, memory, Nightscout, the local KaiTimer dashboard/wheel/milestones/counters, instruction rewrites, and explicit git commits";
        var autonomyDescription = enableAutonomyTools
            ? " For autonomy production guidance, use GET /api/autonomy/status, POST /api/autonomy/guidance/propose, POST /api/autonomy/guidance/approve, or POST /api/autonomy/guidance/reject."
            : "";
        var pathDescription = enableAutonomyTools
            ? "KaiChat tool endpoint path, for example /api/kaitimer/dashboard, /api/kaitimer/wheel/spin, /api/autonomy/status, /api/autonomy/guidance/propose, /api/files/read, /api/memory/ai-manage."
            : "KaiChat tool endpoint path, for example /api/kaitimer/dashboard, /api/kaitimer/wheel/spin, /api/files/read, /api/memory/ai-manage.";

        return
        [
            new
            {
                type = "function",
                function = new
                {
                    name = "kai_tool_call",
                    description = $"Call one KaiChat server tool. Use this for {toolAreas}. Do not describe tool use in text when this function can be called. Successful write tools such as /api/instructions/rewrite, /api/memory/* writes, /api/bible/update, and KaiTimer write controls auto-commit; do not call /api/git/commit after them unless the tool result says commit failed or Raymond explicitly asks. For KaiTimer, use GET /api/kaitimer/dashboard, POST /api/kaitimer/start|pause|reset, POST /api/kaitimer/wheel/segments, /api/kaitimer/wheel/segments/update, /api/kaitimer/wheel/segments/delete, /api/kaitimer/wheel/spin, /api/kaitimer/milestones, /api/kaitimer/milestones/update, /api/kaitimer/milestones/delete, /api/kaitimer/counters, /api/kaitimer/counters/update, /api/kaitimer/counters/increment, /api/kaitimer/counters/decrement, or /api/kaitimer/counters/delete.{autonomyDescription}",
                    parameters = new
                    {
                        type = "object",
                        additionalProperties = false,
                            properties = new
                            {
                                method = new
                                {
                                    type = "string",
                                    description = "HTTP method for the KaiChat tool endpoint.",
                                    @enum = new[] { "GET", "POST", "PUT", "PATCH", "DELETE" }
                                },
                            path = new
                            {
                                type = "string",
                                description = pathDescription
                            },
                            body = new
                            {
                                type = "object",
                                description = "JSON request body for the KaiChat tool endpoint. Use {} when the endpoint does not need parameters.",
                                additionalProperties = true
                            }
                        },
                        required = new[] { "method", "path", "body" }
                    }
                }
            }
        ];
    }

    private static void AccumulateToolCallDeltas(
        SortedDictionary<int, NativeToolCallDelta> toolCallDeltas,
        JsonElement toolCalls)
    {
        foreach (var toolCall in toolCalls.EnumerateArray())
        {
            var index = toolCall.TryGetProperty("index", out var indexProp) && indexProp.TryGetInt32(out var parsedIndex)
                ? parsedIndex
                : toolCallDeltas.Count;

            if (!toolCallDeltas.TryGetValue(index, out var delta))
            {
                delta = new NativeToolCallDelta();
                toolCallDeltas[index] = delta;
            }

            if (toolCall.TryGetProperty("id", out var idProp))
                delta.Id = Join(delta.Id, idProp.GetString());

            if (!toolCall.TryGetProperty("function", out var function) || function.ValueKind != JsonValueKind.Object)
                continue;

            if (function.TryGetProperty("name", out var nameProp))
                delta.Name = Join(delta.Name, nameProp.GetString());

            if (function.TryGetProperty("arguments", out var argsProp))
                delta.Arguments.Append(argsProp.GetString());
        }
    }

    private static string Join(string current, string? next)
    {
        if (string.IsNullOrEmpty(next)) return current;
        return current + next;
    }

    private sealed class NativeToolCallDelta
    {
        public string Id { get; set; } = "";
        public string Name { get; set; } = "";
        public StringBuilder Arguments { get; } = new();
    }

    private sealed record NativeToolCall(string Id, string Name, string Arguments);
}