using System;
using System.Collections.Generic;
using System.IO;
using System.Linq;
using System.Net.Http;
using System.Net.Http.Headers;
using System.Net.WebSockets;
using System.Text;
using System.Text.RegularExpressions;
using System.Threading.Tasks;
using FreneticUtilities.FreneticExtensions;
using Newtonsoft.Json.Linq;
using SwarmUI.Accounts;
using SwarmUI.Core;
using SwarmUI.Text2Image;
using SwarmUI.Utils;
using SwarmUI.WebAPI;

namespace Mrleo1nid.SwarmAssistent;

/// <summary>Krea 2 collaborative assistant: Ollama chat + vision + prompt/LoRA/params patches + Generate/Civitai.</summary>
public class SwarmAssistentExtension : Extension
{
    public static PermInfo PermUse = Permissions.Register(new(
        "swarm_assistent_use",
        "[Swarm Assistent] Use",
        "Allows using the Swarm Assistent chat (Ollama proxy).",
        PermissionDefault.USER,
        Permissions.GroupUser));

    public static HttpClient HttpClient;

    public static readonly string[] PackNames =
    [
        "base_krea2",
        "write_prompt",
        "critique_image",
        "compose_scene",
        "fix_params",
        "inpaint_edit",
        "describe_ref",
    ];

    const int MaxCivitaiHops = 2;
    const int MaxLorasInInventory = 120;
    const int MaxWildcardsInInventory = 80;
    /// <summary>Ollama default num_ctx is 4096; Assistent system+inventory+vision exceeds that.</summary>
    const int DefaultNumCtx = 16384;

    static readonly Regex JsonFenceRe = new(@"```(?:json)?\s*([\s\S]*?)```", RegexOptions.IgnoreCase | RegexOptions.Compiled);

    public override void OnPreInit()
    {
        ScriptFiles.Add("Assets/assistent.js");
        StyleSheetFiles.Add("Assets/assistent.css");
        ExtensionAuthor = "mrleo1nid";
        Description = "Collaborative Krea 2 assistant: Ollama chat, multi-window board, slash commands, img2img/inpaint, Generate loop, Civitai Confirm.";
        License = "MIT";
        Version = "0.5.0";
        Tags = ["tabs", "ui", "llm", "ollama", "krea", "inpaint"];
    }

    public override void OnInit()
    {
        HttpClient ??= new HttpClient { Timeout = TimeSpan.FromMinutes(10) };
        API.RegisterAPICall(AssistentListModels, false, PermUse);
        API.RegisterAPICall(AssistentGetPacks, false, PermUse);
        API.RegisterAPICall(AssistentListInventory, false, PermUse);
        API.RegisterAPICall(AssistentSearchCivitai, false, PermUse);
        API.RegisterAPICall(AssistentChat, true, PermUse);
        API.RegisterAPICall(AssistentChatWS, true, PermUse);
        Logs.Init("Swarm Assistent extension loaded (Ollama proxy + Krea 2 packs + inventory/Civitai)");
    }

    static string Clip(string text, int max)
    {
        if (string.IsNullOrEmpty(text) || text.Length <= max)
        {
            return text ?? "";
        }
        return text[..max] + "…";
    }

    public static string NormalizeBaseUrl(string raw)
    {
        string url = (raw ?? "").Trim();
        if (string.IsNullOrWhiteSpace(url))
        {
            url = "http://127.0.0.1:11434";
        }
        return url.TrimEnd('/');
    }

    public string ReadPackFile(string name)
    {
        string safe = name.Replace('\\', '/').AfterLast('/').Replace("..", "");
        if (!PackNames.Contains(safe))
        {
            return null;
        }
        string path = Path.Combine(FilePath, "Prompts", $"{safe}.md");
        if (!File.Exists(path))
        {
            return null;
        }
        return File.ReadAllText(path, Encoding.UTF8);
    }

    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 models = [];
            foreach (JToken m in parsed["models"] as JArray ?? [])
            {
                models.Add(m["name"]?.ToString() ?? "");
            }
            return new JObject { ["success"] = true, ["base_url"] = root, ["models"] = models };
        }
        catch (Exception ex)
        {
            return new JObject { ["error"] = $"Ollama unreachable at {root}: {ex.Message}" };
        }
    }

    public async Task<JObject> AssistentGetPacks(Session session)
    {
        JObject packs = new();
        foreach (string name in PackNames)
        {
            string text = ReadPackFile(name);
            if (text is not null)
            {
                packs[name] = text;
            }
        }
        return new JObject { ["success"] = true, ["packs"] = packs, ["order"] = new JArray(PackNames) };
    }

    /// <summary>Server-side LoRA / checkpoint / wildcard inventory (not DOM scrape).</summary>
    public async Task<JObject> AssistentListInventory(Session session)
    {
        await Task.CompletedTask;
        JArray loras = [];
        JArray checkpoints = [];
        JArray wildcards = [];

        if (Program.T2IModelSets.TryGetValue("LoRA", out T2IModelHandler loraHandler))
        {
            foreach (T2IModel model in loraHandler.Models.Values.OrderBy(m => m.Name).Take(MaxLorasInInventory))
            {
                loras.Add(new JObject
                {
                    ["name"] = model.Name,
                    ["title"] = model.Metadata?.Title ?? model.Title ?? model.Name,
                    ["trigger_phrase"] = model.Metadata?.TriggerPhrase,
                    ["architecture"] = model.ModelClass?.ID,
                    ["compat_class"] = model.ModelClass?.CompatClass?.ID,
                    ["hash"] = model.Metadata?.Hash ?? "",
                });
            }
        }

        if (Program.T2IModelSets.TryGetValue("Stable-Diffusion", out T2IModelHandler ckptHandler))
        {
            foreach (T2IModel model in ckptHandler.Models.Values.OrderBy(m => m.Name).Take(60))
            {
                checkpoints.Add(new JObject
                {
                    ["name"] = model.Name,
                    ["title"] = model.Metadata?.Title ?? model.Title ?? model.Name,
                    ["architecture"] = model.ModelClass?.ID,
                    ["compat_class"] = model.ModelClass?.CompatClass?.ID,
                });
            }
        }

        try
        {
            foreach (string name in WildcardsHelper.ListFiles.OrderBy(n => n).Take(MaxWildcardsInInventory))
            {
                wildcards.Add(new JObject { ["name"] = name });
            }
        }
        catch (Exception ex)
        {
            Logs.Debug($"AssistentListInventory wildcards: {ex.Message}");
        }

        bool hasCivitaiKey = !string.IsNullOrWhiteSpace(session.User.GetGenericData("civitai_api", "key"));

        return new JObject
        {
            ["success"] = true,
            ["loras"] = loras,
            ["checkpoints"] = checkpoints,
            ["wildcards"] = wildcards,
            ["has_civitai_key"] = hasCivitaiKey,
        };
    }

    /// <summary>Search Civitai for LoRAs (prefers Krea 2 base). Uses Swarm-stored civitai_api key.</summary>
    public async Task<JObject> AssistentSearchCivitai(Session session, string query, int limit = 8)
    {
        string q = (query ?? "").Trim();
        if (string.IsNullOrWhiteSpace(q))
        {
            return new JObject { ["error"] = "query is required" };
        }
        limit = Math.Clamp(limit, 1, 20);
        string apiKey = session.User.GetGenericData("civitai_api", "key") ?? "";
        HashSet<string> installedNames = CollectInstalledLoraNames();
        HashSet<string> installedHashes = CollectInstalledLoraHashes();

        string[] hosts = ["civitai.red", "civitai.com"];
        Exception lastEx = null;
        foreach (string host in hosts)
        {
            try
            {
                string url = $"https://{host}/api/v1/models?limit={limit}&types=LORA&query={Uri.EscapeDataString(q)}";
                using HttpRequestMessage req = new(HttpMethod.Get, url);
                if (!string.IsNullOrWhiteSpace(apiKey))
                {
                    req.Headers.Authorization = new AuthenticationHeaderValue("Bearer", apiKey.Trim());
                }
                using HttpResponseMessage resp = await HttpClient.SendAsync(req);
                string body = await resp.Content.ReadAsStringAsync();
                if (!resp.IsSuccessStatusCode)
                {
                    lastEx = new Exception($"HTTP {(int)resp.StatusCode}: {Clip(body, 200)}");
                    continue;
                }
                JObject parsed = JObject.Parse(body);
                JArray items = parsed["items"] as JArray ?? [];
                JArray results = [];
                foreach (JToken item in items)
                {
                    if (item is not JObject mo)
                    {
                        continue;
                    }
                    JObject card = BuildCivitaiCard(mo, installedNames, installedHashes);
                    if (card is not null)
                    {
                        results.Add(card);
                    }
                }
                // Prefer Krea-compatible first
                JArray sorted = new(results.OrderByDescending(t => LooksLikeKrea(t["base_model"]?.ToString())).ThenBy(t => t["name"]?.ToString()));
                return new JObject
                {
                    ["success"] = true,
                    ["query"] = q,
                    ["host"] = host,
                    ["results"] = sorted,
                    ["has_civitai_key"] = !string.IsNullOrWhiteSpace(apiKey),
                };
            }
            catch (Exception ex)
            {
                lastEx = ex;
            }
        }
        return new JObject { ["error"] = $"Civitai search failed: {lastEx?.Message ?? "unknown"}" };
    }

    static bool LooksLikeKrea(string text) => !string.IsNullOrEmpty(text) && Regex.IsMatch(text, @"krea", RegexOptions.IgnoreCase);

    static HashSet<string> CollectInstalledLoraNames()
    {
        HashSet<string> names = new(StringComparer.OrdinalIgnoreCase);
        if (!Program.T2IModelSets.TryGetValue("LoRA", out T2IModelHandler handler))
        {
            return names;
        }
        foreach (T2IModel m in handler.Models.Values)
        {
            names.Add(m.Name);
            string leaf = m.Name.Replace('\\', '/').AfterLast('/');
            if (!string.IsNullOrEmpty(leaf))
            {
                names.Add(leaf);
                names.Add(Path.GetFileNameWithoutExtension(leaf));
            }
        }
        return names;
    }

    static HashSet<string> CollectInstalledLoraHashes()
    {
        HashSet<string> hashes = new(StringComparer.OrdinalIgnoreCase);
        if (!Program.T2IModelSets.TryGetValue("LoRA", out T2IModelHandler handler))
        {
            return hashes;
        }
        foreach (T2IModel m in handler.Models.Values)
        {
            string h = m.Metadata?.Hash;
            if (!string.IsNullOrWhiteSpace(h))
            {
                hashes.Add(h.Trim().ToLowerInvariant());
            }
        }
        return hashes;
    }

    static JObject BuildCivitaiCard(JObject model, HashSet<string> installedNames, HashSet<string> installedHashes)
    {
        string name = model["name"]?.ToString() ?? "";
        JArray versions = model["modelVersions"] as JArray;
        JObject ver = versions?.FirstOrDefault() as JObject;
        if (ver is null)
        {
            return null;
        }
        string baseModel = ver["baseModel"]?.ToString() ?? "";
        JArray trained = ver["trainedWords"] as JArray ?? [];
        List<string> triggers = trained.Select(t => t?.ToString()).Where(s => !string.IsNullOrWhiteSpace(s)).Take(8).ToList();
        JObject file = null;
        foreach (JToken f in ver["files"] as JArray ?? [])
        {
            if (f is JObject fo && (fo["primary"]?.Value<bool>() == true || (fo["name"]?.ToString() ?? "").EndsWith(".safetensors", StringComparison.OrdinalIgnoreCase)))
            {
                file = fo;
                break;
            }
        }
        file ??= (ver["files"] as JArray)?.FirstOrDefault() as JObject;
        string downloadUrl = file?["downloadUrl"]?.ToString() ?? ver["downloadUrl"]?.ToString() ?? "";
        string fileName = file?["name"]?.ToString() ?? "";
        string sha = file?["hashes"]?["SHA256"]?.ToString() ?? file?["hashes"]?["AutoV2"]?.ToString() ?? "";
        string saveName = string.IsNullOrWhiteSpace(fileName)
            ? Regex.Replace(name, @"[^\w\-.]+", "_").Trim('_')
            : Path.GetFileNameWithoutExtension(fileName);

        bool already = false;
        if (!string.IsNullOrWhiteSpace(sha) && installedHashes.Contains(sha.Trim().ToLowerInvariant()))
        {
            already = true;
        }
        else if (installedNames.Contains(saveName) || installedNames.Contains(name) || installedNames.Contains(fileName))
        {
            already = true;
        }

        return new JObject
        {
            ["id"] = model["id"],
            ["version_id"] = ver["id"],
            ["name"] = name,
            ["base_model"] = baseModel,
            ["krea_likely"] = LooksLikeKrea(baseModel),
            ["triggers"] = new JArray(triggers),
            ["download_url"] = downloadUrl,
            ["file_name"] = saveName,
            ["sha256"] = sha,
            ["already_installed"] = already,
            ["n_sfw"] = model["nsfw"]?.Value<bool>() ?? false,
        };
    }

    List<JObject> BuildOllamaMessages(string packName, bool includeBase, string contextJson, JArray userMessages, string extraSystem = null)
    {
        List<JObject> ollamaMessages = [];
        StringBuilder system = new();
        if (includeBase)
        {
            string basePack = ReadPackFile("base_krea2");
            if (!string.IsNullOrWhiteSpace(basePack))
            {
                system.AppendLine(basePack);
            }
        }
        if (!string.IsNullOrWhiteSpace(packName) && packName != "base_krea2")
        {
            string situational = ReadPackFile(packName);
            if (!string.IsNullOrWhiteSpace(situational))
            {
                system.AppendLine();
                system.AppendLine($"## Active mode: {packName}");
                system.AppendLine(situational);
            }
        }
        if (!string.IsNullOrWhiteSpace(contextJson))
        {
            system.AppendLine();
            system.AppendLine("## Live SwarmUI context (JSON — trust this over guesses)");
            system.AppendLine("```json");
            system.AppendLine(contextJson);
            system.AppendLine("```");
        }
        if (!string.IsNullOrWhiteSpace(extraSystem))
        {
            system.AppendLine();
            system.AppendLine(extraSystem);
        }
        if (system.Length > 0)
        {
            ollamaMessages.Add(new JObject
            {
                ["role"] = "system",
                ["content"] = system.ToString(),
            });
        }
        foreach (JToken msg in userMessages ?? [])
        {
            if (msg is not JObject mo)
            {
                continue;
            }
            JObject copy = new()
            {
                ["role"] = mo["role"]?.ToString() ?? "user",
                ["content"] = mo["content"]?.ToString() ?? "",
            };
            if (mo["images"] is JArray images && images.Count > 0)
            {
                copy["images"] = images;
            }
            ollamaMessages.Add(copy);
        }
        return ollamaMessages;
    }

    static JObject TryParsePatch(string reply)
    {
        if (string.IsNullOrWhiteSpace(reply))
        {
            return null;
        }
        foreach (Match match in JsonFenceRe.Matches(reply))
        {
            string raw = match.Groups[1].Value.Trim();
            try
            {
                JObject obj = JObject.Parse(raw);
                if (obj is not null && (obj["prompt"] != null || obj["loras"] != null || obj["width"] != null
                    || obj["height"] != null || obj["steps"] != null || obj["cfg"] != null
                    || obj["seed"] != null || obj["sigma_shift"] != null || obj["sampler"] != null
                    || obj["actions"] != null || obj["search_query"] != null || obj["civitai_query"] != null
                    || obj["use_init_image"] != null || obj["clear_init_image"] != null
                    || obj["init_creativity"] != null || obj["denoise"] != null
                    || obj["use_mask_image"] != null || obj["clear_mask_image"] != null
                    || obj["mask_blur"] != null || obj["mask_grow"] != null
                    || obj["look_at"] != null || obj["vision_from"] != null || obj["vision_slots"] != null
                    || obj["slot_to_init"] != null || obj["slot_to_mask"] != null
                    || obj["snapshot_generate"] != null || obj["select_slot"] != null
                    || obj["aspect"] != null || obj["images"] != null || obj["batch"] != null
                    || obj["vary"] != null || obj["lock_seed"] != null
                    || obj["creativity"] != null || obj["intensity"] != null
                    || obj["complexity"] != null || obj["movement"] != null
                    || obj["clear_prompt_images"] != null || obj["slot_to_prompt_image"] != null
                    || obj["pack"] != null))
                {
                    return obj;
                }
            }
            catch
            {
                // not json
            }
        }
        return null;
    }

    static string ExtractSearchQuery(JObject patch)
    {
        if (patch is null)
        {
            return null;
        }
        string q = (patch["search_query"] ?? patch["civitai_query"])?.ToString()?.Trim();
        if (!string.IsNullOrWhiteSpace(q))
        {
            return q;
        }
        if (patch["actions"] is JArray acts)
        {
            foreach (JToken a in acts)
            {
                if (string.Equals(a?.ToString(), "search_civitai", StringComparison.OrdinalIgnoreCase))
                {
                    return q; // may still be null — caller checks
                }
            }
        }
        return null;
    }

    static bool WantsCivitaiSearch(JObject patch)
    {
        if (patch is null)
        {
            return false;
        }
        if (!string.IsNullOrWhiteSpace(ExtractSearchQuery(patch)))
        {
            return true;
        }
        if (patch["actions"] is JArray acts)
        {
            foreach (JToken a in acts)
            {
                if (string.Equals(a?.ToString(), "search_civitai", StringComparison.OrdinalIgnoreCase))
                {
                    return true;
                }
            }
        }
        return false;
    }

    async Task<(string reply, JObject raw, JArray civitaiResults)> RunChatWithHops(
        Session session,
        string root,
        string modelName,
        string packName,
        bool includeBase,
        string contextJson,
        JArray userMessages,
        Func<string, Task> onDelta = null,
        Func<int, Task> onHopStart = null)
    {
        List<JObject> messages = BuildOllamaMessages(packName, includeBase, contextJson, userMessages);
        JArray civitaiResults = [];
        string reply = "";
        JObject lastRaw = null;
        for (int hop = 0; hop < MaxCivitaiHops; hop++)
        {
            if (onHopStart is not null)
            {
                await onHopStart(hop);
            }
            (reply, lastRaw) = await CallOllamaChat(root, modelName, messages, stream: onDelta is not null, onDelta);
            JObject patch = TryParsePatch(reply);
            if (hop + 1 >= MaxCivitaiHops || !WantsCivitaiSearch(patch))
            {
                break;
            }
            string query = ExtractSearchQuery(patch);
            if (string.IsNullOrWhiteSpace(query))
            {
                query = userMessages.LastOrDefault(m => m["role"]?.ToString() == "user")?["content"]?.ToString() ?? "";
            }
            if (string.IsNullOrWhiteSpace(query))
            {
                break;
            }
            JObject search = await AssistentSearchCivitai(session, query, 8);
            if (search["error"] is not null)
            {
                messages.Add(new JObject { ["role"] = "assistant", ["content"] = reply });
                messages.Add(new JObject
                {
                    ["role"] = "user",
                    ["content"] = $"Civitai search failed: {search["error"]}. Continue without download — use only available_loras from context.",
                });
                continue;
            }
            civitaiResults = search["results"] as JArray ?? [];
            messages.Add(new JObject { ["role"] = "assistant", ["content"] = reply });
            messages.Add(new JObject
            {
                ["role"] = "user",
                ["content"] =
                    "Civitai search results (JSON). Prefer `krea_likely: true`. Do NOT download yourself — the UI shows Confirm cards. " +
                    "Pick useful LoRAs from results or available_loras, emit a normal patch (prompt/loras). " +
                    "Omit search_civitai from actions unless you need a different query.\n```json\n" +
                    civitaiResults.ToString(Newtonsoft.Json.Formatting.None) + "\n```",
            });
        }
        return (reply, lastRaw, civitaiResults);
    }

    async Task<(string reply, JObject raw)> CallOllamaChat(
        string root,
        string modelName,
        List<JObject> ollamaMessages,
        bool stream,
        Func<string, Task> onDelta)
    {
        JObject payload = new()
        {
            ["model"] = modelName,
            ["stream"] = stream,
            ["messages"] = new JArray(ollamaMessages),
            ["options"] = new JObject
            {
                ["num_ctx"] = DefaultNumCtx,
            },
        };
        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 passes the whole request as the JObject param (not only a nested key).
    /// Support both flat fields and legacy nested <c>raw</c>.
    /// </summary>
    static void ExtractChatPayload(JObject raw, ref string baseUrl, ref string model, ref string pack, ref bool includeBase, out JArray userMessages, out string contextJson)
    {
        JObject whole = raw ?? [];
        JObject nested = whole["raw"] as JObject;
        if (string.IsNullOrWhiteSpace(baseUrl))
        {
            baseUrl = whole["base_url"]?.ToString()
                ?? whole["baseUrl"]?.ToString()
                ?? nested?["base_url"]?.ToString()
                ?? nested?["baseUrl"]?.ToString();
        }
        if (string.IsNullOrWhiteSpace(model))
        {
            model = whole["model"]?.ToString() ?? nested?["model"]?.ToString();
        }
        if (string.IsNullOrWhiteSpace(pack))
        {
            pack = whole["pack"]?.ToString() ?? nested?["pack"]?.ToString();
        }
        if (whole["includeBase"] is not null)
        {
            includeBase = whole.Value<bool?>("includeBase") ?? includeBase;
        }
        userMessages = (whole["messages"] as JArray) ?? (nested?["messages"] as JArray);
        contextJson = whole["context_json"]?.ToString() ?? nested?["context_json"]?.ToString();
    }

    /// <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);
        string root = NormalizeBaseUrl(baseUrl);
        string modelName = (model ?? "").Trim();
        if (string.IsNullOrWhiteSpace(modelName))
        {
            return new JObject { ["error"] = "model is required" };
        }
        if (userMessages is null || userMessages.Count == 0)
        {
            return new JObject { ["error"] = "messages required" };
        }
        string packName = (pack ?? "write_prompt").Trim();
        try
        {
            (string reply, JObject parsed, JArray civitai) = await RunChatWithHops(
                session, root, modelName, packName, includeBase, contextJson, userMessages);
            return new JObject
            {
                ["success"] = true,
                ["reply"] = reply,
                ["model"] = modelName,
                ["pack"] = packName,
                ["raw"] = parsed,
                ["civitai_results"] = civitai,
            };
        }
        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);
        string root = NormalizeBaseUrl(baseUrl);
        string modelName = (model ?? "").Trim();
        if (string.IsNullOrWhiteSpace(modelName))
        {
            await ws.SendJson(new JObject { ["error"] = "model is required" }, API.WebsocketTimeout);
            return null;
        }
        if (userMessages is null || userMessages.Count == 0)
        {
            await ws.SendJson(new JObject { ["error"] = "messages required" }, API.WebsocketTimeout);
            return null;
        }
        string packName = (pack ?? "write_prompt").Trim();
        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) = await RunChatWithHops(
                session, root, modelName, packName, includeBase, contextJson, userMessages, OnDelta, OnHopStart);
            await ws.SendJson(new JObject
            {
                ["success"] = true,
                ["done"] = true,
                ["reply"] = reply,
                ["model"] = modelName,
                ["pack"] = packName,
                ["raw"] = parsed,
                ["civitai_results"] = civitai,
            }, API.WebsocketTimeout);
        }
        catch (Exception ex)
        {
            await ws.SendJson(new JObject { ["error"] = $"Ollama chat failed: {ex.Message}" }, API.WebsocketTimeout);
        }
        return null;
    }
}