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

namespace Mrleo1nid.SwarmAssistent;

/// <summary>Hugging Face datasets: search, compatibility gate, preview, import.</summary>
public partial class SwarmAssistentExtension
{
    static readonly Regex HfRepoIdRe = new(@"^[A-Za-z0-9][A-Za-z0-9._-]{0,95}(/[A-Za-z0-9][A-Za-z0-9._-]{0,95})?$", RegexOptions.Compiled);

    const string HfDatasetsServer = "https://datasets-server.huggingface.co";
    const string HfHubApi = "https://huggingface.co/api/datasets";

    static string GetHfToken(Session session)
        => session?.User?.GetGenericData("huggingface_api", "key")?.Trim();

    static HttpRequestMessage HfRequest(string url, Session session)
    {
        HttpRequestMessage req = new(HttpMethod.Get, url);
        string token = GetHfToken(session);
        if (!string.IsNullOrWhiteSpace(token))
        {
            req.Headers.Authorization = new System.Net.Http.Headers.AuthenticationHeaderValue("Bearer", token);
        }
        return req;
    }

    /// <summary>Normalize owner/name or HF datasets URL. Returns null when invalid.</summary>
    public static string NormalizeHfDatasetId(string raw)
    {
        string s = (raw ?? "").Trim();
        if (string.IsNullOrWhiteSpace(s))
        {
            return null;
        }
        if (s.StartsWith("http://", StringComparison.OrdinalIgnoreCase) || s.StartsWith("https://", StringComparison.OrdinalIgnoreCase))
        {
            if (!Uri.TryCreate(s, UriKind.Absolute, out Uri uri))
            {
                return null;
            }
            if (!string.Equals(uri.Host, "huggingface.co", StringComparison.OrdinalIgnoreCase))
            {
                return null;
            }
            string[] parts = uri.AbsolutePath.Trim('/').Split('/');
            if (parts.Length < 2 || !string.Equals(parts[0], "datasets", StringComparison.OrdinalIgnoreCase))
            {
                return null;
            }
            s = $"{parts[1]}/{parts[2]}";
        }
        s = s.Trim().TrimEnd('/');
        return HfRepoIdRe.IsMatch(s) ? s : null;
    }

    public async Task<JObject> AssistentSearchHfDatasets(Session session, string q = null, int limit = 20, bool show_all = false)
    {
        int take = Math.Clamp(limit, 1, 50);
        string search = (q ?? "").Trim();
        StringBuilder url = new($"{HfHubApi}?limit={take}&full=true");
        url.Append("&filter=task_categories:text-generation");
        url.Append("&filter=modality:text");
        if (!string.IsNullOrWhiteSpace(search))
        {
            url.Append("&search=").Append(Uri.EscapeDataString(search));
        }
        try
        {
            using HttpRequestMessage req = HfRequest(url.ToString(), session);
            using HttpResponseMessage resp = await HttpClient.SendAsync(req);
            string body = await resp.Content.ReadAsStringAsync();
            if (!resp.IsSuccessStatusCode)
            {
                return new JObject { ["error"] = $"HF search HTTP {(int)resp.StatusCode}: {Clip(body, 300)}" };
            }
            JArray rawList = JArray.Parse(body);
            JArray results = [];
            foreach (JToken item in rawList)
            {
                if (item is not JObject o)
                {
                    continue;
                }
                string id = o["id"]?.ToString();
                if (string.IsNullOrWhiteSpace(id))
                {
                    continue;
                }
                JObject check = await CheckHfDatasetInternal(session, id, useCache: true);
                string gate = check["gate"]?.ToString() ?? "rejected";
                if (!show_all && gate == "rejected")
                {
                    continue;
                }
                results.Add(new JObject
                {
                    ["id"] = id,
                    ["title"] = o["id"],
                    ["downloads"] = o["downloads"],
                    ["gate"] = gate,
                    ["reason"] = check["reason"],
                    ["schema"] = check["schema"],
                });
            }
            return new JObject { ["success"] = true, ["results"] = results, ["has_hf_token"] = !string.IsNullOrWhiteSpace(GetHfToken(session)) };
        }
        catch (Exception ex)
        {
            return new JObject { ["error"] = $"HF search: {ex.Message}" };
        }
    }

    public async Task<JObject> AssistentCheckHfDataset(Session session, string dataset)
    {
        string id = NormalizeHfDatasetId(dataset);
        if (id is null)
        {
            return new JObject { ["success"] = false, ["error"] = "Нужен owner/name или ссылка huggingface.co/datasets/…" };
        }
        JObject check = await CheckHfDatasetInternal(session, id, useCache: false);
        check["success"] = check["gate"]?.ToString() != "rejected";
        check["id"] = id;
        return check;
    }

    async Task<JObject> CheckHfDatasetInternal(Session session, string datasetId, bool useCache)
    {
        string cacheKey = $"hf:{datasetId}";
        if (useCache)
        {
            JObject cached = Memory.GetKvObject(KvHfDatasetCache)?[cacheKey] as JObject;
            if (cached is not null && cached["checked_at"]?.Value<long?>() > DateTimeOffset.UtcNow.AddHours(-6).ToUnixTimeMilliseconds())
            {
                return cached;
            }
        }
        JObject result = new() { ["id"] = datasetId, ["gate"] = "rejected", ["reason"] = "unknown" };
        try
        {
            using HttpRequestMessage validReq = HfRequest($"{HfDatasetsServer}/is-valid?dataset={Uri.EscapeDataString(datasetId)}", session);
            using HttpResponseMessage validResp = await HttpClient.SendAsync(validReq);
            string validBody = await validResp.Content.ReadAsStringAsync();
            if (!validResp.IsSuccessStatusCode)
            {
                result["reason"] = $"is-valid HTTP {(int)validResp.StatusCode}";
                return CacheHfCheck(cacheKey, result);
            }
            JObject valid = JObject.Parse(validBody);
            bool viewer = valid["viewer"]?.Value<bool?>() == true;
            bool preview = valid["preview"]?.Value<bool?>() == true;
            if (!viewer && !preview)
            {
                result["reason"] = "Набор не читается через datasets (viewer/preview = false)";
                return CacheHfCheck(cacheKey, result);
            }
            using HttpRequestMessage splitReq = HfRequest($"{HfDatasetsServer}/splits?dataset={Uri.EscapeDataString(datasetId)}", session);
            using HttpResponseMessage splitResp = await HttpClient.SendAsync(splitReq);
            string splitBody = await splitResp.Content.ReadAsStringAsync();
            if (!splitResp.IsSuccessStatusCode)
            {
                result["reason"] = $"splits HTTP {(int)splitResp.StatusCode}";
                return CacheHfCheck(cacheKey, result);
            }
            JArray splits = JObject.Parse(splitBody)["splits"] as JArray ?? [];
            if (splits.Count == 0)
            {
                result["reason"] = "Нет splits";
                return CacheHfCheck(cacheKey, result);
            }
            JObject first = splits[0] as JObject;
            string config = first?["config"]?.ToString() ?? "default";
            string split = first?["split"]?.ToString() ?? "train";
            using HttpRequestMessage rowsReq = HfRequest($"{HfDatasetsServer}/first-rows?dataset={Uri.EscapeDataString(datasetId)}&config={Uri.EscapeDataString(config)}&split={Uri.EscapeDataString(split)}", session);
            using HttpResponseMessage rowsResp = await HttpClient.SendAsync(rowsReq);
            string rowsBody = await rowsResp.Content.ReadAsStringAsync();
            if (!rowsResp.IsSuccessStatusCode)
            {
                result["reason"] = $"first-rows HTTP {(int)rowsResp.StatusCode}: {Clip(rowsBody, 200)}";
                return CacheHfCheck(cacheKey, result);
            }
            JObject rowsData = JObject.Parse(rowsBody);
            JObject features = rowsData["features"] as JObject;
            (string gate, string reason, JObject schema) = ClassifyHfFeatures(features);
            result["gate"] = gate;
            result["reason"] = reason;
            result["schema"] = schema;
            result["config"] = config;
            result["split"] = split;
            result["features"] = features;
            result["sample_rows"] = rowsData["rows"];
            result["runner_only"] = HasHugeSizeTag(datasetId);
            return CacheHfCheck(cacheKey, result);
        }
        catch (Exception ex)
        {
            result["reason"] = ex.Message;
            return CacheHfCheck(cacheKey, result);
        }
    }

    static bool HasHugeSizeTag(string datasetId) => false;

    JObject CacheHfCheck(string cacheKey, JObject result)
    {
        result["checked_at"] = DateTimeOffset.UtcNow.ToUnixTimeMilliseconds();
        try
        {
            JObject bag = Memory.GetKvObject(KvHfDatasetCache) ?? new JObject();
            bag[cacheKey] = result;
            Memory.SetKvObject(KvHfDatasetCache, bag);
        }
        catch (Exception ex)
        {
            Logs.Debug($"CacheHfCheck: {ex.Message}");
        }
        return result;
    }

    static (string gate, string reason, JObject schema) ClassifyHfFeatures(JObject features)
    {
        if (features is null || !features.Properties().Any())
        {
            return ("rejected", "Нет колонок (features пуст)", null);
        }
        HashSet<string> names = new(StringComparer.OrdinalIgnoreCase);
        foreach (JProperty p in features.Properties())
        {
            names.Add(p.Name);
            JToken dtype = p.Value?["dtype"] ?? p.Value?["type"];
            string dt = dtype?.ToString() ?? "";
            if (dt.Contains("image", StringComparison.OrdinalIgnoreCase)
                || dt.Contains("audio", StringComparison.OrdinalIgnoreCase)
                || dt.Contains("video", StringComparison.OrdinalIgnoreCase))
            {
                return ("rejected", $"Мультимодальная колонка {p.Name} ({dt})", null);
            }
        }
        if (names.Contains("chosen") && names.Contains("rejected"))
        {
            return ("rejected", "DPO-набор (chosen/rejected) — не для SFT", null);
        }
        if (names.Count == 1 && names.Contains("text"))
        {
            return ("rejected", "Предобучение (одна колонка text), не диалоги", null);
        }
        if (names.Contains("messages"))
        {
            return ("ok", "OpenAI messages", new JObject { ["kind"] = "messages" });
        }
        if (names.Contains("conversations"))
        {
            return ("ok", "ShareGPT conversations", new JObject { ["kind"] = "conversations" });
        }
        if (names.Contains("instruction") && names.Contains("output"))
        {
            return ("ok", "Alpaca instruction/output", new JObject { ["kind"] = "alpaca" });
        }
        if (names.Contains("prompt") && (names.Contains("response") || names.Contains("completion") || names.Contains("answer")))
        {
            string respCol = names.Contains("response") ? "response" : names.Contains("completion") ? "completion" : "answer";
            return ("ok", "prompt/response", new JObject { ["kind"] = "prompt_response", ["response_col"] = respCol });
        }
        if (names.Contains("question") && names.Contains("answer"))
        {
            return ("ok", "question/answer", new JObject { ["kind"] = "qa" });
        }
        List<string> stringCols = [];
        foreach (JProperty p in features.Properties())
        {
            JToken dtype = p.Value?["dtype"] ?? p.Value?["type"];
            string dt = dtype?.ToString() ?? "";
            if (dt.Contains("string", StringComparison.OrdinalIgnoreCase) || dt == "value")
            {
                stringCols.Add(p.Name);
            }
        }
        if (stringCols.Count >= 2)
        {
            return ("mapping", "Нужен ручной маппинг колонок", new JObject
            {
                ["kind"] = "custom",
                ["columns"] = new JArray(stringCols),
            });
        }
        return ("rejected", "Схема не подходит для SFT", null);
    }

    public async Task<JObject> AssistentPreviewHfDataset(Session session, string dataset, string config = null, string split = null)
    {
        string id = NormalizeHfDatasetId(dataset);
        if (id is null)
        {
            return new JObject { ["error"] = "invalid dataset id" };
        }
        JObject check = await CheckHfDatasetInternal(session, id, useCache: true);
        if (check["gate"]?.ToString() == "rejected")
        {
            return new JObject { ["error"] = check["reason"]?.ToString() ?? "rejected", ["check"] = check };
        }
        return new JObject { ["success"] = true, ["check"] = check };
    }

    public async Task<JObject> AssistentImportHfDataset(Session session, JObject raw)
    {
        string dataset = raw?["dataset"]?.ToString();
        int limit = raw?["limit"]?.Value<int?>() ?? 200;
        JObject mapping = raw?["mapping"] as JObject;
        string id = NormalizeHfDatasetId(dataset);
        if (id is null)
        {
            return new JObject { ["error"] = "invalid dataset id" };
        }
        JObject check = await CheckHfDatasetInternal(session, id, useCache: true);
        string gate = check["gate"]?.ToString();
        if (gate == "rejected")
        {
            return new JObject { ["error"] = check["reason"]?.ToString() ?? "rejected" };
        }
        if (gate == "mapping" && (mapping is null || mapping.Count == 0))
        {
            return new JObject { ["error"] = "Нужен маппинг колонок", ["check"] = check };
        }
        int take = Math.Clamp(limit, 1, 5000);
        JArray rows = [];
        string config = check["config"]?.ToString() ?? "default";
        string split = check["split"]?.ToString() ?? "train";
        int offset = 0;
        while (rows.Count < take)
        {
            int chunk = Math.Min(100, take - rows.Count);
            using HttpRequestMessage rowsReq = HfRequest($"{HfDatasetsServer}/rows?dataset={Uri.EscapeDataString(id)}&config={Uri.EscapeDataString(config)}&split={Uri.EscapeDataString(split)}&offset={offset}&length={chunk}", session);
            using HttpResponseMessage rowsResp = await HttpClient.SendAsync(rowsReq);
            string rowsBody = await rowsResp.Content.ReadAsStringAsync();
            if (!rowsResp.IsSuccessStatusCode)
            {
                break;
            }
            JObject parsed = JObject.Parse(rowsBody);
            JArray batch = parsed["rows"] as JArray ?? [];
            if (batch.Count == 0)
            {
                break;
            }
            foreach (JToken t in batch)
            {
                rows.Add(t);
            }
            offset += batch.Count;
            if (batch.Count < chunk)
            {
                break;
            }
        }
        if (rows.Count == 0)
        {
            rows = check["sample_rows"] as JArray ?? [];
        }
        int imported = 0;
        foreach (JToken rowTok in rows.Take(take))
        {
            if (rowTok is not JObject row)
            {
                continue;
            }
            JObject rowData = row["row"] as JObject ?? row;
            JArray messages = ConvertHfRowToMessages(rowData, check["schema"] as JObject, mapping);
            if (messages is null || messages.Count == 0)
            {
                continue;
            }
            Memory.UpsertTrainSample(new JObject
            {
                ["source"] = "hf",
                ["hf_repo"] = id,
                ["messages"] = messages,
                ["status"] = "draft",
            });
            imported++;
        }
        if (imported == 0 && check["runner_only"]?.Value<bool?>() == true)
        {
            return new JObject { ["success"] = true, ["imported"] = 0, ["runner_only"] = true, ["id"] = id, ["note"] = "Большой набор — используй HF id в QLoRA-раннере" };
        }
        return new JObject { ["success"] = true, ["imported"] = imported, ["id"] = id };
    }

    static JArray ConvertHfRowToMessages(JObject row, JObject schema, JObject mapping)
    {
        string kind = schema?["kind"]?.ToString() ?? mapping?["kind"]?.ToString();
        if (kind == "messages" && row["messages"] is JArray msgs)
        {
            return NormalizeMessagesArray(msgs);
        }
        if (kind == "conversations" && row["conversations"] is JArray conv)
        {
            JArray outArr = [];
            foreach (JToken c in conv)
            {
                if (c is not JObject co)
                {
                    continue;
                }
                string from = co["from"]?.ToString() ?? "";
                string val = co["value"]?.ToString() ?? "";
                string role = from is "human" or "user" ? "user" : from is "gpt" or "assistant" or "chatgpt" ? "assistant" : "user";
                outArr.Add(new JObject { ["role"] = role, ["content"] = val });
            }
            return outArr.Count > 0 ? outArr : null;
        }
        if (kind == "alpaca")
        {
            string instr = row["instruction"]?.ToString() ?? "";
            string inp = row["input"]?.ToString() ?? "";
            string output = row["output"]?.ToString() ?? "";
            string user = string.IsNullOrWhiteSpace(inp) ? instr : $"{instr}\n{inp}";
            return new JArray
            {
                new JObject { ["role"] = "user", ["content"] = user },
                new JObject { ["role"] = "assistant", ["content"] = output },
            };
        }
        if (kind == "prompt_response")
        {
            string respCol = schema?["response_col"]?.ToString() ?? "response";
            return new JArray
            {
                new JObject { ["role"] = "user", ["content"] = row["prompt"]?.ToString() ?? "" },
                new JObject { ["role"] = "assistant", ["content"] = row[respCol]?.ToString() ?? "" },
            };
        }
        if (kind == "qa")
        {
            return new JArray
            {
                new JObject { ["role"] = "user", ["content"] = row["question"]?.ToString() ?? "" },
                new JObject { ["role"] = "assistant", ["content"] = row["answer"]?.ToString() ?? "" },
            };
        }
        if (kind == "custom" && mapping is not null)
        {
            string userCol = mapping["user_col"]?.ToString();
            string asstCol = mapping["assistant_col"]?.ToString();
            if (!string.IsNullOrWhiteSpace(userCol) && !string.IsNullOrWhiteSpace(asstCol))
            {
                return new JArray
                {
                    new JObject { ["role"] = "user", ["content"] = row[userCol]?.ToString() ?? "" },
                    new JObject { ["role"] = "assistant", ["content"] = row[asstCol]?.ToString() ?? "" },
                };
            }
        }
        return null;
    }

    static JArray NormalizeMessagesArray(JArray msgs)
    {
        JArray outArr = [];
        foreach (JToken m in msgs)
        {
            if (m is not JObject mo)
            {
                continue;
            }
            string role = mo["role"]?.ToString() ?? "user";
            string content = mo["content"]?.ToString() ?? mo["text"]?.ToString() ?? "";
            if (string.IsNullOrWhiteSpace(content))
            {
                continue;
            }
            outArr.Add(new JObject { ["role"] = role, ["content"] = content });
        }
        return outArr.Count > 0 ? outArr : null;
    }
}