using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Threading.Tasks;
using Newtonsoft.Json.Linq;
using SwarmUI.Accounts;
using SwarmUI.Utils;

namespace Mrleo1nid.SwarmAssistent;

/// <summary>Prompt assembly, memory retrieval/writeback and the Civitai search hop loop.</summary>
public partial class SwarmAssistentExtension
{
    const int MaxCivitaiHopsFallback = 2;
    const int MaxToolHopsFallback = 4;

    List<JObject> BuildOllamaMessages(string packName, bool includeBase, string contextJson, JArray userMessages, string extraSystem = null, string personaId = null, IEnumerable<string> skillIds = null)
    {
        List<JObject> ollamaMessages = [];
        StringBuilder system = new();
        string pid = AssistentConfig.SafeId(personaId) ?? Config.DefaultPersonaId();

        if (includeBase)
        {
            string core = Config.LoadCorePrompt(pid);
            if (!string.IsNullOrWhiteSpace(core))
            {
                system.AppendLine(core);
            }
        }

        JObject exact = Config.LoadExactForPrompt(pid);
        if (exact is not null && exact.Count > 0)
        {
            system.AppendLine();
            system.AppendLine("## Exact memory (canonical KV defaults — prefer over RAG for numbers)");
            system.AppendLine("```json");
            system.AppendLine(exact.ToString(Newtonsoft.Json.Formatting.None));
            system.AppendLine("```");
        }

        foreach (string skillId in skillIds ?? Config.ResolveEnabledSkills(pid, null))
        {
            string skillText = Config.LoadSkillPrompt(pid, skillId);
            if (!string.IsNullOrWhiteSpace(skillText))
            {
                system.AppendLine();
                system.AppendLine($"## Skill: {skillId}");
                system.AppendLine(skillText);
            }
        }

        string identity = Config.RenderIdentityBlock(pid);
        if (!string.IsNullOrWhiteSpace(identity))
        {
            system.AppendLine();
            system.AppendLine(identity);
        }

        if (!string.IsNullOrWhiteSpace(packName))
        {
            string situational = Config.LoadPackPrompt(pid, 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;
    }

    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,
        string personaId = null,
        JArray skillIds = null,
        string embedModel = null)
    {
        string pid = AssistentConfig.SafeId(personaId) ?? Config.DefaultPersonaId();
        List<string> skills = Config.ResolveEnabledSkills(pid, skillIds);
        string embed = string.IsNullOrWhiteSpace(embedModel)
            ? (Config.LoadSettings()["embed_model"]?.ToString()
               ?? Config.LoadAssistant(pid)["embed_model"]?.ToString()
               ?? "nomic-embed-text")
            : embedModel;

        try
        {
            await Memory.EnsureSeedAsync(root, Config, embed);
        }
        catch (Exception ex)
        {
            Logs.Debug($"Assistent memory seed: {ex.Message}");
        }

        string retrieveQuery = BuildRetrieveQuery(userMessages, contextJson, packName);
        JArray hits = [];
        try
        {
            AssistentMemory.RetrieveOptions opt = MemoryRetrieveOptions(pid);
            hits = await Memory.RetrieveAsync(root, retrieveQuery, opt.TopK, embed, Config.PersonaExtendsChain(pid), opt);
        }
        catch (Exception ex)
        {
            Logs.Debug($"Assistent memory retrieve: {ex.Message}");
        }

        string enrichedContext = InjectMemoryHits(contextJson, hits);
        enrichedContext = EnrichPersonaContext(enrichedContext, pid, packName);
        List<JObject> messages = BuildOllamaMessages(packName, includeBase, enrichedContext, userMessages, personaId: pid, skillIds: skills);
        JArray civitaiResults = [];
        string reply = "";
        JObject lastRaw = null;
        int maxHops = Math.Max(CfgInt("max_civitai_hops", MaxCivitaiHopsFallback), CfgInt("max_tool_hops", MaxToolHopsFallback));
        HashSet<string> hopDone = new(StringComparer.OrdinalIgnoreCase);
        var chain = Config.PersonaExtendsChain(pid);
        for (int hop = 0; hop < maxHops; hop++)
        {
            if (onHopStart is not null)
            {
                await onHopStart(hop);
            }
            (reply, lastRaw) = await CallOllamaChat(root, modelName, messages, stream: onDelta is not null, onDelta, pid);
            JObject patch = TryParsePatch(reply);
            await ApplyMemoryActions(root, patch, embed, pid);
            ApplyPersonaActions(patch, ref pid);
            if (hop + 1 >= maxHops)
            {
                break;
            }
            string tool = NextToolHop(patch);
            if (string.IsNullOrWhiteSpace(tool))
            {
                break;
            }
            (string follow, JArray civitaiHop) = await RunToolHop(session, root, embed, pid, chain, patch, tool, hopDone);
            if (follow is null)
            {
                break;
            }
            if (civitaiHop is { Count: > 0 })
            {
                civitaiResults = civitaiHop;
            }
            messages.Add(new JObject { ["role"] = "assistant", ["content"] = reply });
            messages.Add(new JObject { ["role"] = "user", ["content"] = follow });
        }
        return (reply, lastRaw, civitaiResults);
    }

    static string BuildRetrieveQuery(JArray userMessages, string contextJson, string packName = null)
    {
        StringBuilder sb = new();
        if (!string.IsNullOrWhiteSpace(packName))
        {
            sb.Append(packName).Append(' ');
        }
        if (!string.IsNullOrWhiteSpace(contextJson))
        {
            try
            {
                JObject ctx = JObject.Parse(contextJson);
                string ckpt = ctx["checkpoint"]?.ToString() ?? ctx["current_model"]?.ToString();
                if (!string.IsNullOrWhiteSpace(ckpt))
                {
                    sb.Append(ckpt).Append(' ');
                }
                if (ctx["enabled_loras"] is JArray en)
                {
                    foreach (JToken t in en.Take(12))
                    {
                        string n = t?["name"]?.ToString() ?? t?.ToString();
                        if (!string.IsNullOrWhiteSpace(n))
                        {
                            sb.Append(n).Append(' ');
                        }
                    }
                }
                if (ctx["krea_profile"] != null)
                {
                    sb.Append("krea ").Append(ctx["krea_profile"]).Append(' ');
                }
                string aspect = ctx["aspect"]?.ToString();
                if (!string.IsNullOrWhiteSpace(aspect))
                {
                    sb.Append(aspect).Append(' ');
                }
                string prompt = ctx["prompt"]?.ToString();
                if (!string.IsNullOrWhiteSpace(prompt))
                {
                    sb.Append(prompt.Length > 400 ? prompt[..400] : prompt).Append(' ');
                }
            }
            catch
            {
                // ignore
            }
        }
        foreach (JToken msg in (userMessages ?? []).Reverse().Take(3))
        {
            if (msg is JObject mo && string.Equals(mo["role"]?.ToString(), "user", StringComparison.OrdinalIgnoreCase))
            {
                string c = mo["content"]?.ToString() ?? "";
                sb.Append(c.Length > 500 ? c[..500] : c).Append(' ');
            }
        }
        string q = CollapseWs(sb.ToString());
        return string.IsNullOrWhiteSpace(q) ? "krea2 prompting" : q;
    }

    AssistentMemory.RetrieveOptions MemoryRetrieveOptions(string pid)
    {
        JObject a = Config.LoadAssistant(pid) ?? new JObject();
        AssistentMemory.RetrieveOptions opt = new()
        {
            TopK = a["memory_top_k"]?.Value<int?>() ?? 10,
            MinScore = a["memory_min_score"]?.Value<float?>() ?? 0.32f,
            ApplyQuotas = true,
        };
        if (a["memory_quotas"] is JObject quotas)
        {
            Dictionary<string, int> d = new(StringComparer.OrdinalIgnoreCase);
            foreach (JProperty p in quotas.Properties())
            {
                d[p.Name] = p.Value?.Value<int?>() ?? 2;
            }
            opt.Quotas = d;
        }
        return opt;
    }

    async Task<(string follow, JArray civitai)> RunToolHop(
        Session session,
        string root,
        string embed,
        string pid,
        IEnumerable<string> chain,
        JObject patch,
        string tool,
        HashSet<string> hopDone)
    {
        if (tool == "memory_get")
        {
            JArray got = [];
            foreach (JToken t in patch["memories"] as JArray ?? [])
            {
                if (t is not JObject mo)
                {
                    continue;
                }
                string kind = mo["kind"]?.ToString() ?? "note";
                string key = mo["key"]?.ToString() ?? "";
                if (string.IsNullOrWhiteSpace(key))
                {
                    continue;
                }
                string sig = $"get:{kind}:{key}";
                if (!hopDone.Add(sig))
                {
                    continue;
                }
                JObject row = Memory.Get(kind, key, chain);
                got.Add(row ?? new JObject { ["kind"] = kind, ["key"] = key, ["missing"] = true });
            }
            if (got.Count == 0)
            {
                return (null, null);
            }
            return (
                "memory_get results (JSON). Use these facts; omit memory_get unless you need a different key.\n```json\n"
                + got.ToString(Newtonsoft.Json.Formatting.None) + "\n```",
                null);
        }
        if (tool == "memory_search")
        {
            string q = ExtractMemoryQuery(patch);
            if (string.IsNullOrWhiteSpace(q) || !hopDone.Add("search:" + q))
            {
                return (null, null);
            }
            string kind = patch["memory_kind"]?.ToString();
            int topK = Config.LoadAssistant(pid)["memory_top_k"]?.Value<int?>() ?? 10;
            JArray rows = await Memory.SearchAsync(root, q, kind, topK, embed, chain);
            return (
                "memory_search results (JSON, hybrid FTS+vector). Omit memory_search unless you need a different query.\n```json\n"
                + rows.ToString(Newtonsoft.Json.Formatting.None) + "\n```",
                null);
        }
        if (tool == "lookup_tags")
        {
            string q = ExtractTagQuery(patch);
            if (string.IsNullOrWhiteSpace(q) || !hopDone.Add("tags:" + q))
            {
                return (null, null);
            }
            int lim = Config.LoadAssistant(pid)["tag_lookup_limit"]?.Value<int?>() ?? 20;
            JArray tags = Memory.LookupTags(q, lim);
            return (
                "lookup_tags results from Danbooru csv (canonical name, aliases, post_count). Krea prompts stay natural prose — use this to check spelling/aliases, do not dump tag soup.\n```json\n"
                + tags.ToString(Newtonsoft.Json.Formatting.None) + "\n```",
                null);
        }
        if (tool == "civitai")
        {
            string query = ExtractSearchQuery(patch);
            if (string.IsNullOrWhiteSpace(query) || !hopDone.Add("civitai:" + query))
            {
                return (null, null);
            }
            JObject search = await AssistentSearchCivitai(session, query, 8);
            if (search["error"] is not null)
            {
                return ($"Civitai search failed: {search["error"]}. Continue without download — use only available_loras from context.", null);
            }
            JArray civitaiResults = search["results"] as JArray ?? [];
            return (
                "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```",
                civitaiResults);
        }
        return (null, null);
    }

    static string InjectMemoryHits(string contextJson, JArray hits, JObject exact = null)
    {
        JObject ctx;
        try
        {
            ctx = string.IsNullOrWhiteSpace(contextJson) ? new JObject() : JObject.Parse(contextJson);
        }
        catch
        {
            ctx = new JObject { ["_raw_context"] = contextJson };
        }
        ctx["memory_hits"] = hits ?? new JArray();
        // Never re-inject full Exact into live context (already in system prompt).
        ctx.Remove("exact");
        if (ctx["session_exact"] is null)
        {
            ctx["session_exact"] = new JObject();
        }
        // Slim inventory for LLM: keep enabled + current, drop full dump if present
        if (ctx["available_loras"] is JArray allLoras && allLoras.Count > 24)
        {
            HashSet<string> keep = new(StringComparer.OrdinalIgnoreCase);
            if (ctx["enabled_loras"] is JArray en)
            {
                foreach (JToken t in en)
                {
                    string n = t?["name"]?.ToString() ?? t?.ToString();
                    if (!string.IsNullOrWhiteSpace(n))
                    {
                        keep.Add(n);
                    }
                }
            }
            foreach (JToken hit in hits ?? [])
            {
                if (string.Equals(hit?["kind"]?.ToString(), "lora", StringComparison.OrdinalIgnoreCase)
                    || string.Equals(hit?["kind"]?.ToString(), "card", StringComparison.OrdinalIgnoreCase))
                {
                    string k = hit?["key"]?.ToString();
                    if (!string.IsNullOrWhiteSpace(k))
                    {
                        keep.Add(k);
                    }
                }
            }
            JArray slim = [];
            foreach (JToken t in allLoras)
            {
                string n = t?["name"]?.ToString();
                if (!string.IsNullOrWhiteSpace(n) && (keep.Contains(n) || slim.Count < 12))
                {
                    if (keep.Contains(n) || t?["krea_likely"]?.Value<bool>() == true)
                    {
                        slim.Add(t);
                    }
                }
            }
            if (slim.Count == 0)
            {
                foreach (JToken t in allLoras.Take(12))
                {
                    slim.Add(t);
                }
            }
            ctx["available_loras"] = slim;
            ctx["available_loras_truncated"] = true;
            ctx["available_loras_total"] = allLoras.Count;
        }
        return ctx.ToString(Newtonsoft.Json.Formatting.None);
    }

    static string MemoryWritePersona(JObject mo, string currentPersonaId)
    {
        string scope = (mo?["scope"]?.ToString() ?? "").Trim().ToLowerInvariant();
        if (scope is "shared" or "common" or "global")
        {
            return AssistentMemory.SharedPersona;
        }
        // Personal only — never let the model write into another personality's store.
        return AssistentConfig.SafeId(currentPersonaId) ?? AssistentMemory.SharedPersona;
    }

    string EnrichPersonaContext(string contextJson, string personaId, string packName)
    {
        JObject ctx;
        try
        {
            ctx = string.IsNullOrWhiteSpace(contextJson) ? new JObject() : JObject.Parse(contextJson);
        }
        catch
        {
            ctx = new JObject { ["_raw_context"] = contextJson };
        }
        string pid = AssistentConfig.SafeId(personaId) ?? Config.DefaultPersonaId();
        ctx["persona_source"] = Config.PersonaSource(pid);
        JObject schema = Config.LoadControlsSchema(pid);
        JObject values = Config.LoadControlValues(pid);
        if (schema.Properties().Any())
        {
            ctx["persona_controls"] = new JObject
            {
                ["schema"] = schema,
                ["values"] = values,
            };
        }
        JArray catalog = [];
        foreach (var p in Config.ListPersonaCatalog())
        {
            catalog.Add(new JObject
            {
                ["id"] = p.id,
                ["title"] = p.title,
                ["source"] = p.source,
            });
        }
        ctx["personas"] = catalog;
        if (string.Equals(packName, "author_persona", StringComparison.OrdinalIgnoreCase)
            || string.Equals(packName, "persona", StringComparison.OrdinalIgnoreCase))
        {
            JObject shelves = Config.LoadIdentityParts(pid);
            shelves.Remove("extra");
            ctx["persona_shelves"] = shelves;
            ctx["persona_controls_schema"] = schema;
        }
        return ctx.ToString(Newtonsoft.Json.Formatting.None);
    }

    /// <summary>Apply overlay persona clone/write from patch. Ignores persona_delete. Updates pid ref after switch.</summary>
    void ApplyPersonaActions(JObject patch, ref string personaId)
    {
        if (patch is null || Config is null)
        {
            return;
        }
        // Never honor delete from the model.
        bool wantClone = false, wantWrite = false;
        if (patch["actions"] is JArray acts)
        {
            foreach (JToken a in acts)
            {
                string s = a?.ToString() ?? "";
                if (string.Equals(s, "persona_clone", StringComparison.OrdinalIgnoreCase))
                {
                    wantClone = true;
                }
                if (string.Equals(s, "persona_write", StringComparison.OrdinalIgnoreCase))
                {
                    wantWrite = true;
                }
            }
        }
        if (patch["persona_clone"] is JObject)
        {
            wantClone = true;
        }
        if (patch["persona_shelves"] is JObject)
        {
            wantWrite = true;
        }
        try
        {
            if (wantClone && patch["persona_clone"] is JObject clone)
            {
                string from = AssistentConfig.SafeId(clone["from"]?.ToString()) ?? personaId;
                string to = AssistentConfig.SafeId(clone["to"]?.ToString());
                string title = clone["title"]?.ToString();
                bool overwrite = clone["overwrite"]?.Value<bool?>() == true;
                if (to is not null)
                {
                    Config.ClonePersonaToOverlay(from, to, title, overwrite);
                    personaId = to;
                    patch["_persona_cloned"] = to;
                }
            }
            if (wantWrite && patch["persona_shelves"] is JObject shelves)
            {
                string target = AssistentConfig.SafeId(patch["persona"]?.ToString())
                    ?? AssistentConfig.SafeId(patch["persona_clone"]?["to"]?.ToString())
                    ?? personaId;
                if (target is not null)
                {
                    Config.SavePersonaShelves(target, shelves);
                    patch["_persona_written"] = target;
                }
            }
            // Control values from model patch (Exact).
            if (patch["controls"] is JObject ctrlVals)
            {
                string ctrlPid = AssistentConfig.SafeId(patch["persona"]?.ToString()) ?? personaId;
                Config.SaveControlValues(ctrlPid, ctrlVals);
                patch["_controls_saved"] = true;
            }
        }
        catch (Exception ex)
        {
            Logs.Warning($"Assistent persona actions: {ex.Message}");
            patch["_persona_error"] = ex.Message;
        }
    }

    async Task ApplyMemoryActions(string root, JObject patch, string embedModel, string personaId)
    {
        if (patch is null || Memory is null)
        {
            return;
        }
        bool upsert = false, forget = false;
        if (patch["actions"] is JArray acts)
        {
            foreach (JToken a in acts)
            {
                string s = a?.ToString() ?? "";
                if (string.Equals(s, "memory_upsert", StringComparison.OrdinalIgnoreCase))
                {
                    upsert = true;
                }
                if (string.Equals(s, "memory_forget", StringComparison.OrdinalIgnoreCase))
                {
                    forget = true;
                }
            }
        }
        JArray memories = patch["memories"] as JArray;
        if (memories is null || memories.Count == 0)
        {
            return;
        }
        foreach (JToken t in memories)
        {
            if (t is not JObject mo)
            {
                continue;
            }
            string kind = mo["kind"]?.ToString() ?? "note";
            string key = mo["key"]?.ToString() ?? "";
            string text = mo["text"]?.ToString() ?? "";
            string target = MemoryWritePersona(mo, personaId);
            try
            {
                if (forget && string.IsNullOrWhiteSpace(text))
                {
                    Memory.Forget(kind, key, persona: target);
                }
                else if (upsert || !string.IsNullOrWhiteSpace(text))
                {
                    await Memory.UpsertTextAsync(root, kind, key, text, "user", mo, embedModel, target);
                }
            }
            catch (Exception ex)
            {
                Logs.Debug($"ApplyMemoryActions: {ex.Message}");
            }
        }
    }
}