using System;
using System.Collections.Generic;
using System.IO;
using System.Linq;
using System.Net.Http;
using System.Net.WebSockets;
using System.Text;
using System.Threading.Tasks;
using Newtonsoft.Json.Linq;
using SwarmUI.Accounts;
using SwarmUI.Utils;
using SwarmUI.WebAPI;

namespace Mrleo1nid.SwarmAssistent;

/// <summary>Ollama transport: model listing, /api/chat calls and the two chat API endpoints.</summary>
public partial class SwarmAssistentExtension
{
    const int DefaultNumCtxFallback = 16384;

    public async Task<JObject> AssistentListModels(Session session, string baseUrl)
    {
        string root = NormalizeBaseUrl(baseUrl);
        try
        {
            using HttpResponseMessage resp = await HttpClient.GetAsync($"{root}/api/tags");
            string body = await resp.Content.ReadAsStringAsync();
            if (!resp.IsSuccessStatusCode)
            {
                return new JObject { ["error"] = $"Ollama /api/tags HTTP {(int)resp.StatusCode}: {Clip(body, 400)}" };
            }
            JObject parsed = JObject.Parse(body);
            JArray all = [];
            foreach (JToken m in parsed["models"] as JArray ?? [])
            {
                string name = m["name"]?.ToString() ?? m["model"]?.ToString() ?? "";
                if (!string.IsNullOrWhiteSpace(name))
                {
                    all.Add(name);
                }
            }
            JObject roles = Config?.LoadOllamaRoles() ?? new JObject();
            HashSet<string> chatSet = new(
                (roles["chat"] as JArray)?.Select(t => t?.ToString()).Where(s => !string.IsNullOrWhiteSpace(s)) ?? [],
                StringComparer.OrdinalIgnoreCase);
            HashSet<string> memSet = new(
                (roles["memory"] as JArray)?.Select(t => t?.ToString()).Where(s => !string.IsNullOrWhiteSpace(s)) ?? [],
                StringComparer.OrdinalIgnoreCase);
            // Heuristic fallbacks when sidecar missing
            if (chatSet.Count == 0 && memSet.Count == 0)
            {
                foreach (JToken t in all)
                {
                    string n = t.ToString();
                    if (LooksLikeEmbedModel(n))
                    {
                        memSet.Add(n);
                    }
                    else
                    {
                        chatSet.Add(n);
                    }
                }
            }
            else
            {
                // Keep only tags that exist; anything unlabeled goes to chat if not memory
                foreach (JToken t in all)
                {
                    string n = t.ToString();
                    if (memSet.Contains(n) || LooksLikeEmbedModel(n))
                    {
                        memSet.Add(n);
                        chatSet.Remove(n);
                    }
                    else if (chatSet.Count == 0 || chatSet.Contains(n))
                    {
                        chatSet.Add(n);
                    }
                    else if (!memSet.Contains(n))
                    {
                        chatSet.Add(n);
                    }
                }
            }
            JArray models = new(all.Select(t => t.ToString()).Where(n => chatSet.Contains(n) && !memSet.Contains(n) && !LooksLikeEmbedModel(n)));
            JArray memoryModels = new(all.Select(t => t.ToString()).Where(n => memSet.Contains(n) || LooksLikeEmbedModel(n)).Distinct(StringComparer.OrdinalIgnoreCase).ToList());
            if (memoryModels.Count == 0)
            {
                string fallback = Config?.LoadAssistant(Config.DefaultPersonaId())["embed_model"]?.ToString() ?? "nomic-embed-text";
                if (all.Any(t => string.Equals(t.ToString(), fallback, StringComparison.OrdinalIgnoreCase)
                    || t.ToString().StartsWith(fallback.Split(':')[0], StringComparison.OrdinalIgnoreCase)))
                {
                    memoryModels.Add(all.Select(t => t.ToString()).First(n =>
                        string.Equals(n, fallback, StringComparison.OrdinalIgnoreCase)
                        || n.StartsWith(fallback.Split(':')[0], StringComparison.OrdinalIgnoreCase)));
                }
            }
            return new JObject
            {
                ["success"] = true,
                ["base_url"] = root,
                ["models"] = models,
                ["memory_models"] = memoryModels,
            };
        }
        catch (Exception ex)
        {
            return new JObject { ["error"] = $"Ollama unreachable at {root}: {ex.Message}" };
        }
    }

    static bool LooksLikeEmbedModel(string name)
    {
        string n = (name ?? "").ToLowerInvariant();
        return n.Contains("embed") || n.Contains("nomic") || n.Contains("bge-") || n.Contains("minilm") || n.Contains("e5-");
    }

    async Task<(string reply, JObject raw)> CallOllamaChat(
        string root,
        string modelName,
        List<JObject> ollamaMessages,
        bool stream,
        Func<string, Task> onDelta,
        string personaId = null)
    {
        int numCtx = Config.LoadAssistant(AssistentConfig.SafeId(personaId) ?? Config.DefaultPersonaId())["num_ctx"]?.Value<int?>()
            ?? DefaultNumCtxFallback;
        JObject payload = new()
        {
            ["model"] = modelName,
            ["stream"] = stream,
            ["messages"] = new JArray(ollamaMessages),
            ["options"] = new JObject
            {
                ["num_ctx"] = numCtx,
            },
            ["keep_alive"] = "15m",
        };
        using StringContent content = new(payload.ToString(Newtonsoft.Json.Formatting.None), Encoding.UTF8, "application/json");
        using HttpRequestMessage req = new(HttpMethod.Post, $"{root}/api/chat") { Content = content };
        using HttpResponseMessage resp = await HttpClient.SendAsync(req, stream
            ? HttpCompletionOption.ResponseHeadersRead
            : HttpCompletionOption.ResponseContentRead);
        if (!resp.IsSuccessStatusCode)
        {
            string errBody = await resp.Content.ReadAsStringAsync();
            throw new Exception($"Ollama /api/chat HTTP {(int)resp.StatusCode}: {Clip(errBody, 800)}");
        }
        if (!stream)
        {
            string body = await resp.Content.ReadAsStringAsync();
            JObject parsed = JObject.Parse(body);
            string reply = parsed["message"]?["content"]?.ToString() ?? parsed["response"]?.ToString() ?? "";
            return (reply, parsed);
        }
        StringBuilder full = new();
        await using Stream streamBody = await resp.Content.ReadAsStreamAsync();
        using StreamReader reader = new(streamBody, Encoding.UTF8);
        JObject last = null;
        while (true)
        {
            string line = await reader.ReadLineAsync();
            if (line is null)
            {
                break;
            }
            if (string.IsNullOrWhiteSpace(line))
            {
                continue;
            }
            JObject chunk = JObject.Parse(line);
            last = chunk;
            string delta = chunk["message"]?["content"]?.ToString() ?? "";
            if (!string.IsNullOrEmpty(delta))
            {
                full.Append(delta);
                if (onDelta is not null)
                {
                    await onDelta(delta);
                }
            }
            if (chunk["done"]?.Value<bool>() == true)
            {
                break;
            }
        }
        return (full.ToString(), last ?? new JObject());
    }

    /// <summary>SwarmUI hands the whole request body over as the JObject param, so every field is read flat off it.</summary>
    static void ExtractChatPayload(JObject raw, ref string baseUrl, ref string model, ref string pack, ref bool includeBase, out JArray userMessages, out string contextJson, out string persona, out JArray skills)
    {
        JObject whole = raw ?? [];
        if (string.IsNullOrWhiteSpace(baseUrl))
        {
            baseUrl = whole["base_url"]?.ToString() ?? whole["baseUrl"]?.ToString();
        }
        if (string.IsNullOrWhiteSpace(model))
        {
            model = whole["model"]?.ToString();
        }
        if (string.IsNullOrWhiteSpace(pack))
        {
            pack = whole["pack"]?.ToString();
        }
        if (whole["includeBase"] is not null)
        {
            includeBase = whole.Value<bool?>("includeBase") ?? includeBase;
        }
        userMessages = whole["messages"] as JArray;
        contextJson = whole["context_json"]?.ToString();
        persona = whole["persona"]?.ToString() ?? "neutral";
        skills = whole["skills"] as JArray;
    }

    /// <summary>Shared validation for both chat endpoints. Returns an error message, or null when the request is usable.</summary>
    static string ValidateChatRequest(string modelName, JArray userMessages)
    {
        if (string.IsNullOrWhiteSpace(modelName))
        {
            return "model is required";
        }
        if (userMessages is null || userMessages.Count == 0)
        {
            return "messages required";
        }
        return null;
    }

    /// <summary>Proxy to Ollama /api/chat (non-stream), with optional Civitai search hop.</summary>
    public async Task<JObject> AssistentChat(Session session, string baseUrl, string model, string pack, bool includeBase, JObject raw)
    {
        ExtractChatPayload(raw, ref baseUrl, ref model, ref pack, ref includeBase, out JArray userMessages, out string contextJson, out string persona, out JArray skills);
        string root = NormalizeBaseUrl(baseUrl);
        string modelName = (model ?? "").Trim();
        string invalid = ValidateChatRequest(modelName, userMessages);
        if (invalid is not null)
        {
            return new JObject { ["error"] = invalid };
        }
        string packName = (pack ?? "write_prompt").Trim();
        string embedModel = raw?["embed_model"]?.ToString() ?? Config.LoadSettings()["embed_model"]?.ToString();
        try
        {
            (string reply, JObject parsed, JArray civitai, int systemChars) = await RunChatWithHops(
                session, root, modelName, packName, includeBase, contextJson, userMessages, personaId: persona, skillIds: skills, embedModel: embedModel);
            return new JObject
            {
                ["success"] = true,
                ["reply"] = reply,
                ["model"] = modelName,
                ["pack"] = packName,
                ["persona"] = persona,
                ["raw"] = parsed,
                ["civitai_results"] = civitai,
                ["system_chars"] = systemChars,
            };
        }
        catch (Exception ex)
        {
            return new JObject { ["error"] = $"Ollama chat failed: {ex.Message}" };
        }
    }

    /// <summary>WebSocket streaming chat (Ollama stream:true) + Civitai hops.</summary>
    public async Task<JObject> AssistentChatWS(Session session, WebSocket ws, string baseUrl, string model, string pack, bool includeBase, JObject raw)
    {
        ExtractChatPayload(raw, ref baseUrl, ref model, ref pack, ref includeBase, out JArray userMessages, out string contextJson, out string persona, out JArray skills);
        string root = NormalizeBaseUrl(baseUrl);
        string modelName = (model ?? "").Trim();
        string invalid = ValidateChatRequest(modelName, userMessages);
        if (invalid is not null)
        {
            await ws.SendJson(new JObject { ["error"] = invalid }, API.WebsocketTimeout);
            return null;
        }
        string packName = (pack ?? "write_prompt").Trim();
        string embedModel = raw?["embed_model"]?.ToString() ?? Config.LoadSettings()["embed_model"]?.ToString();
        try
        {
            if (ws.State == WebSocketState.Open)
            {
                await ws.SendJson(new JObject
                {
                    ["phase"] = "waiting_ollama",
                    ["notice"] = "Loading model into GPU…",
                }, API.WebsocketTimeout);
            }
            async Task OnDelta(string delta)
            {
                if (ws.State == WebSocketState.Open)
                {
                    await ws.SendJson(new JObject { ["delta"] = delta }, API.WebsocketTimeout);
                }
            }
            async Task OnHopStart(int hop)
            {
                if (ws.State == WebSocketState.Open && hop > 0)
                {
                    await ws.SendJson(new JObject
                    {
                        ["clear_stream"] = true,
                        ["hop"] = hop + 1,
                        ["notice"] = "Civitai search done — refining…",
                    }, API.WebsocketTimeout);
                }
            }
            (string reply, JObject parsed, JArray civitai, int systemChars) = await RunChatWithHops(
                session, root, modelName, packName, includeBase, contextJson, userMessages, OnDelta, OnHopStart, persona, skills, embedModel);
            await ws.SendJson(new JObject
            {
                ["success"] = true,
                ["done"] = true,
                ["reply"] = reply,
                ["model"] = modelName,
                ["pack"] = packName,
                ["persona"] = persona,
                ["raw"] = parsed,
                ["civitai_results"] = civitai,
                ["system_chars"] = systemChars,
            }, API.WebsocketTimeout);
        }
        catch (Exception ex)
        {
            await ws.SendJson(new JObject { ["error"] = $"Ollama chat failed: {ex.Message}" }, API.WebsocketTimeout);
        }
        return null;
    }
}