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

namespace Mrleo1nid.SwarmAssistent;

/// <summary>OpenRouter chat transport (OpenAI-compatible /chat/completions). Chat only —
/// memory embeddings stay on Ollama. The API key lives server-side in settings.json
/// (or OPENROUTER_API_KEY) and is never sent back to the browser.</summary>
public partial class SwarmAssistentExtension
{
    const string OpenRouterRoot = "https://openrouter.ai/api/v1";

    const string ProviderOllama = "ollama";

    const string ProviderOpenRouter = "openrouter";

    string ChatProvider()
    {
        string p = Config?.LoadSettings()["provider"]?.ToString()?.Trim().ToLowerInvariant();
        return p == ProviderOpenRouter ? ProviderOpenRouter : ProviderOllama;
    }

    bool UseOpenRouter() => ChatProvider() == ProviderOpenRouter;

    string ProviderLabel() => UseOpenRouter() ? "OpenRouter" : "Ollama";

    string OpenRouterKey()
    {
        string key = Config?.LoadSettings()["openrouter_api_key"]?.ToString()?.Trim();
        return string.IsNullOrWhiteSpace(key) ? Environment.GetEnvironmentVariable("OPENROUTER_API_KEY")?.Trim() ?? "" : key;
    }

    /// <summary>Provider state for the settings pane. Only a masked tail of the key is exposed.</summary>
    public async Task<JObject> AssistentGetProvider(Session session)
    {
        await Task.CompletedTask;
        string key = OpenRouterKey();
        return new JObject
        {
            ["success"] = true,
            ["provider"] = ChatProvider(),
            ["openrouter_key_set"] = key.Length > 0,
            ["openrouter_key_hint"] = key.Length > 8 ? $"…{key[^4..]}" : "",
            ["openrouter_model"] = Config.LoadSettings()["openrouter_model"]?.ToString() ?? "",
        };
    }

    /// <summary>Sets provider / key / model. Null fields are left untouched; an empty <paramref name="api_key"/> clears the key.</summary>
    public async Task<JObject> AssistentSetProvider(Session session, string provider = null, string api_key = null, string model = null)
    {
        await Task.CompletedTask;
        JObject patch = [];
        if (provider is not null)
        {
            string p = provider.Trim().ToLowerInvariant();
            if (p != ProviderOllama && p != ProviderOpenRouter)
            {
                return new JObject { ["error"] = "provider must be ollama or openrouter" };
            }
            patch["provider"] = p;
        }
        if (api_key is not null)
        {
            patch["openrouter_api_key"] = api_key.Trim();
        }
        if (model is not null)
        {
            patch["openrouter_model"] = model.Trim();
        }
        Config.SaveSettings(patch);
        return await AssistentGetProvider(session);
    }

    HttpRequestMessage OpenRouterRequest(HttpMethod method, string path, HttpContent content = null)
    {
        HttpRequestMessage req = new(method, $"{OpenRouterRoot}{path}") { Content = content };
        string key = OpenRouterKey();
        if (key.Length > 0)
        {
            req.Headers.TryAddWithoutValidation("Authorization", $"Bearer {key}");
        }
        req.Headers.TryAddWithoutValidation("HTTP-Referer", "https://github.com/mcmonkeyprojects/SwarmUI");
        req.Headers.TryAddWithoutValidation("X-Title", "SwarmUI Assistent");
        return req;
    }

    /// <summary>OpenRouter chat models + Ollama embed models (best effort) for memory.</summary>
    async Task<JObject> ListOpenRouterModels(string ollamaRoot)
    {
        if (OpenRouterKey().Length == 0)
        {
            return new JObject { ["error"] = "OpenRouter: API key not set" };
        }
        JArray models = [];
        JArray visionModels = [];
        try
        {
            using HttpRequestMessage req = OpenRouterRequest(HttpMethod.Get, "/models");
            using HttpResponseMessage resp = await HttpClient.SendAsync(req);
            string body = await resp.Content.ReadAsStringAsync();
            if (!resp.IsSuccessStatusCode)
            {
                return new JObject { ["error"] = $"OpenRouter /models HTTP {(int)resp.StatusCode}: {Clip(body, 400)}" };
            }
            IEnumerable<JObject> entries = (JObject.Parse(body)["data"] as JArray ?? []).OfType<JObject>()
                .Where(m => !string.IsNullOrWhiteSpace(m["id"]?.ToString()))
                .Where(m => (m["architecture"]?["output_modalities"] as JArray)?.Any(t => t.ToString() == "text") ?? true)
                .OrderBy(m => m["id"].ToString(), StringComparer.OrdinalIgnoreCase);
            foreach (JObject m in entries)
            {
                string id = m["id"].ToString();
                models.Add(id);
                if ((m["architecture"]?["input_modalities"] as JArray)?.Any(t => t.ToString() == "image") == true)
                {
                    visionModels.Add(id);
                }
            }
        }
        catch (Exception ex)
        {
            return new JObject { ["error"] = $"OpenRouter unreachable: {ex.Message}" };
        }
        JArray memoryModels = [];
        JObject ollama = await ListOllamaModels(ollamaRoot);
        if (ollama["memory_models"] is JArray mem)
        {
            memoryModels = mem;
        }
        string preferred = Config.LoadSettings()["openrouter_model"]?.ToString()?.Trim() ?? "";
        return new JObject
        {
            ["success"] = true,
            ["provider"] = ProviderOpenRouter,
            ["base_url"] = OpenRouterRoot,
            ["models"] = models,
            ["vision_models"] = visionModels,
            ["memory_models"] = memoryModels,
            ["memory_error"] = ollama["error"],
            ["preferred"] = models.Any(t => t.ToString() == preferred) ? preferred : "",
        };
    }

    /// <summary>Ollama-style messages (<c>images</c> = raw base64) → OpenAI content parts.</summary>
    static JArray ToOpenAiMessages(List<JObject> ollamaMessages)
    {
        JArray result = [];
        foreach (JObject m in ollamaMessages)
        {
            string role = m["role"]?.ToString() ?? "user";
            string text = m["content"]?.ToString() ?? "";
            if (m["images"] is JArray images && images.Count > 0)
            {
                JArray parts = [];
                if (text.Length > 0)
                {
                    parts.Add(new JObject { ["type"] = "text", ["text"] = text });
                }
                foreach (JToken img in images)
                {
                    string s = img.ToString();
                    string url = s.StartsWith("data:", StringComparison.OrdinalIgnoreCase) || s.StartsWith("http", StringComparison.OrdinalIgnoreCase)
                        ? s
                        : $"data:image/jpeg;base64,{s}";
                    parts.Add(new JObject { ["type"] = "image_url", ["image_url"] = new JObject { ["url"] = url } });
                }
                result.Add(new JObject { ["role"] = role, ["content"] = parts });
            }
            else
            {
                result.Add(new JObject { ["role"] = role, ["content"] = text });
            }
        }
        return result;
    }

    /// <summary>Copies OpenAI usage into the Ollama field names the pipeline / UI already read.</summary>
    static void MapUsage(JObject raw, JToken usage)
    {
        if (usage is JObject u && u["prompt_tokens"]?.Value<int?>() is int prompt && prompt > 0)
        {
            raw["prompt_eval_count"] = prompt;
        }
    }

    async Task<(string reply, JObject raw)> CallOpenRouterChat(
        string modelName,
        List<JObject> ollamaMessages,
        bool stream,
        Func<string, Task> onDelta,
        string personaId = null)
    {
        if (OpenRouterKey().Length == 0)
        {
            throw new Exception("OpenRouter API key not set (⚙ → Модели)");
        }
        int numPredict = Config.LoadAssistant(AssistentConfig.SafeId(personaId) ?? Config.DefaultPersonaId())["num_predict"]?.Value<int?>()
            ?? 3072;
        if (numPredict < 512)
        {
            numPredict = 512;
        }
        JObject payload = new()
        {
            ["model"] = modelName,
            ["stream"] = stream,
            ["messages"] = ToOpenAiMessages(ollamaMessages),
            ["max_tokens"] = numPredict,
            ["usage"] = new JObject { ["include"] = true },
        };
        using StringContent content = new(payload.ToString(Newtonsoft.Json.Formatting.None), Encoding.UTF8, "application/json");
        using HttpRequestMessage req = OpenRouterRequest(HttpMethod.Post, "/chat/completions", 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($"OpenRouter HTTP {(int)resp.StatusCode}: {Clip(errBody, 800)}");
        }
        if (!stream)
        {
            string body = await resp.Content.ReadAsStringAsync();
            JObject parsed = JObject.Parse(body);
            if (parsed["error"] is JToken err)
            {
                throw new Exception($"OpenRouter: {Clip(err["message"]?.ToString() ?? err.ToString(), 800)}");
            }
            MapUsage(parsed, parsed["usage"]);
            string reply = parsed["choices"]?[0]?["message"]?["content"]?.ToString() ?? "";
            if (TryTruncateAtCompleteFence(reply, out string cut))
            {
                reply = cut;
            }
            return (reply, parsed);
        }
        StringBuilder full = new();
        JObject last = [];
        await using Stream streamBody = await resp.Content.ReadAsStreamAsync();
        using StreamReader reader = new(streamBody, Encoding.UTF8);
        while (true)
        {
            string line = await reader.ReadLineAsync();
            if (line is null)
            {
                break;
            }
            // SSE: skip blanks and ": OPENROUTER PROCESSING" keep-alive comments.
            if (!line.StartsWith("data:", StringComparison.Ordinal))
            {
                continue;
            }
            string data = line[5..].Trim();
            if (data == "[DONE]")
            {
                break;
            }
            JObject chunk;
            try
            {
                chunk = JObject.Parse(data);
            }
            catch (Exception)
            {
                continue;
            }
            if (chunk["error"] is JToken err)
            {
                throw new Exception($"OpenRouter: {Clip(err["message"]?.ToString() ?? err.ToString(), 800)}");
            }
            last["id"] = chunk["id"];
            last["model"] = chunk["model"];
            MapUsage(last, chunk["usage"]);
            string delta = chunk["choices"]?[0]?["delta"]?["content"]?.ToString() ?? "";
            if (string.IsNullOrEmpty(delta))
            {
                continue;
            }
            full.Append(delta);
            // Same early stop as the Ollama path: closed ```json``` patch → stop reading.
            if (TryTruncateAtCompleteFence(full.ToString(), out string cut))
            {
                string extra = full.Length > cut.Length ? full.ToString(cut.Length, full.Length - cut.Length) : "";
                full.Clear();
                full.Append(cut);
                int keep = delta.Length - extra.Length;
                if (onDelta is not null && keep > 0)
                {
                    await onDelta(delta[..keep]);
                }
                break;
            }
            if (onDelta is not null)
            {
                await onDelta(delta);
            }
        }
        return (full.ToString(), last);
    }
}