// The depot and chunk download flow follows DepotDownloader (https://github.com/SteamRE/DepotDownloader,
// commit 989f37b, GPL-2.0, (c) SteamRE). This file is therefore under GPL-2.0.

using System.Buffers;
using System.Collections.Concurrent;
using System.Net;
using System.Text.Json;
using CVRManager.Models;
using SteamKit2;
using SteamKit2.CDN;

namespace CVRManager.Services;

public static class SteamDownloader
{
    const uint AppId = (uint)Paths.AppId;
    const int MaxDownloads = 8;
    const string StateFolder = ".CVRManager";

    public sealed record PrivateBranch(string Name, string Code, byte[]? Key);

    const int MaxCodeChecksPerRefresh = 3;

    public static async Task<List<BranchInfo>> GetBranchesAsync(
        string account, string refreshToken, IReadOnlyList<PrivateBranch> privateBranches,
        Action<string, byte[]> keyResolved, Action<string> log, CancellationToken ct)
    {
        await using var link = await SteamLink.OpenAsync(account, refreshToken, ct);
        var (root, appToken) = await RequestAppInfoAsync(link);
        var depots = root["depots"];
        var branches = depots["branches"];
        if (branches == KeyValue.Invalid) throw new InvalidOperationException("Steam returned no branch data for ChilloutVR.");
        var list = branches.Children
            .Select(b => new BranchInfo(b.Name ?? "", b["buildid"].Value ?? "", b["pwdrequired"].Value == "1", b["description"].Value ?? ""))
            .ToList();

        var known = list.Select(b => b.Name).ToHashSet();
        var keys = new Dictionary<string, byte[]>();
        var checks = 0;
        var limited = false;
        foreach (var pb in privateBranches.Where(p => p.Code.Length > 0 && !known.Contains(p.Name)))
        {
            ct.ThrowIfCancellationRequested();
            var key = pb.Key ?? keys.GetValueOrDefault(pb.Name);
            if (key == null)
            {
                if (limited || checks >= MaxCodeChecksPerRefresh) continue;
                checks++;
                try
                {
                    foreach (var (name, k) in await BranchCodeService.ResolveAsync(link, pb.Code))
                    {
                        keys[name] = k;
                        keyResolved(name, k);
                    }
                }
                catch (InvalidOperationException e)
                {
                    log($"Could not check the code for '{pb.Name}': {e.Message}");
                    limited = e.Message.Contains("RateLimit");
                    continue;
                }
                key = keys.GetValueOrDefault(pb.Name);
                if (key == null) { log($"The saved code no longer unlocks '{pb.Name}'."); continue; }
            }

            var priv = await link.Apps.PICSGetPrivateBeta(AppId, appToken, pb.Name, key);
            if (priv.Result != EResult.OK) { log($"Steam would not list the files of '{pb.Name}' ({priv.Result})."); continue; }
            var selection = SelectDepots(depots, priv.DepotSection, pb.Name);
            if (selection.Count == 0) continue;
            list.Add(new BranchInfo(pb.Name, PrivateBuildId(priv.DepotSection, pb.Name) ?? PrivateMarker(selection), true, ""));
        }

        return list.OrderBy(b => b.Name == "public" ? 0 : 1).ThenBy(b => b.Name).ToList();
    }

    static string? PrivateBuildId(KeyValue? section, string branch)
    {
        var id = section?["branches"][branch]["buildid"].Value;
        return string.IsNullOrEmpty(id) ? null : id;
    }

    static string PrivateMarker(IEnumerable<(uint DepotId, ulong ManifestId)> depots)
    {
        var text = string.Join(";", depots.OrderBy(d => d.DepotId).Select(d => $"{d.DepotId}:{d.ManifestId}"));
        return Convert.ToHexString(System.Security.Cryptography.SHA1.HashData(System.Text.Encoding.ASCII.GetBytes(text)))[..12].ToLowerInvariant();
    }

    static List<(uint DepotId, ulong ManifestId)> SelectDepots(KeyValue depotsSection, KeyValue? privateDepots, string branch)
    {
        var result = new List<(uint, ulong)>();
        foreach (var depot in depotsSection.Children)
        {
            if (!uint.TryParse(depot.Name, out var depotId) || depot.Children.Count == 0) continue;
            var cfg = depot["config"];
            if (cfg != KeyValue.Invalid)
            {
                var os = cfg["oslist"].Value;
                if (!string.IsNullOrWhiteSpace(os) && !os.Split(',').Contains("windows")) continue;
                var arch = cfg["osarch"].Value;
                if (!string.IsNullOrWhiteSpace(arch) && arch != "64") continue;
                var lang = cfg["language"].Value;
                if (!string.IsNullOrWhiteSpace(lang) && lang != "english") continue;
                if (cfg["lowviolence"].AsBoolean()) continue;
            }

            var gid = (privateDepots?[depot.Name!]["manifests"][branch]["gid"].Value) ?? depot["manifests"][branch]["gid"].Value;
            if (gid == null && branch != "public") gid = depot["manifests"]["public"]["gid"].Value;
            if (gid == null) continue;
            result.Add((depotId, ulong.Parse(gid)));
        }
        return result.OrderBy(d => d.Item1).ToList();
    }

    static async Task<(KeyValue Root, ulong Token)> RequestAppInfoAsync(SteamLink link)
    {
        var tokens = await link.Apps.PICSGetAccessTokens(new List<uint> { AppId }, new List<uint>());
        tokens.AppTokens.TryGetValue(AppId, out var token);
        var request = new SteamApps.PICSRequest(AppId) { AccessToken = token };
        var result = await link.Apps.PICSGetProductInfo(new List<SteamApps.PICSRequest> { request }, new List<SteamApps.PICSRequest>());
        var info = result.Results?.SelectMany(r => r.Apps).Where(a => a.Key == AppId).Select(a => a.Value).FirstOrDefault()
                   ?? throw new InvalidOperationException("Steam returned no app info for ChilloutVR.");
        return (info.KeyValues, token);
    }

    sealed record DepotPlan(uint DepotId, ulong ManifestId, byte[] Key);

    sealed class InstallState
    {
        public bool Complete { get; set; }
        public Dictionary<string, ulong> Manifests { get; set; } = new();
    }

    public static async Task InstallAsync(
        string gameDir, string branch, string branchCode, byte[]? branchKey, Action<string, byte[]> keyResolved,
        string account, string refreshToken, bool validate, Action<string> log, Action<double> progress, CancellationToken ct)
    {
        if (string.IsNullOrWhiteSpace(branch)) branch = "public";
        Directory.CreateDirectory(gameDir);
        var stateDir = Path.Combine(gameDir, StateFolder);
        Directory.CreateDirectory(stateDir);

        await using var link = await SteamLink.OpenAsync(account, refreshToken, ct);
        log(link.IsAnonymous ? "Connected to Steam anonymously." : $"Signed in to Steam as {account}.");

        var (root, appToken) = await RequestAppInfoAsync(link);
        var depotsSection = root["depots"];
        var branchNode = depotsSection["branches"][branch];
        var listed = branchNode != KeyValue.Invalid;
        var buildId = listed ? branchNode["buildid"].Value ?? "" : "";

        KeyValue? privateDepots = null;
        if (!listed || branchNode["pwdrequired"].Value == "1")
        {
            if (string.IsNullOrEmpty(branchCode))
                throw new InvalidOperationException($"The branch '{branch}' is not public and needs its beta code. Enter the code from the branch list.");
            if (branchKey == null)
            {
                var keys = await BranchCodeService.ResolveAsync(link, branchCode);
                foreach (var (name, k) in keys) keyResolved(name, k);
                if (!keys.TryGetValue(branch, out branchKey))
                    throw new InvalidOperationException($"The saved code does not unlock '{branch}'. Enter the code again from the branch list.");
            }
            var priv = await link.Apps.PICSGetPrivateBeta(AppId, appToken, branch, branchKey);
            if (priv.Result != EResult.OK) throw new InvalidOperationException($"Steam would not give out the files for '{branch}' ({priv.Result}).");
            privateDepots = priv.DepotSection;
        }

        var plans = new List<DepotPlan>();
        foreach (var (depotId, manifestId) in SelectDepots(depotsSection, privateDepots, branch))
        {
            var keyReply = await link.Apps.GetDepotDecryptionKey(depotId, AppId);
            if (keyReply.Result != EResult.OK) throw new InvalidOperationException($"Steam refused the key for depot {depotId} ({keyReply.Result}).");
            plans.Add(new DepotPlan(depotId, manifestId, keyReply.DepotKey));
        }
        if (plans.Count == 0) throw new InvalidOperationException("Steam listed no downloadable files for this branch.");
        if (buildId.Length == 0) buildId = PrivateBuildId(privateDepots, branch) ?? PrivateMarker(plans.Select(p => (p.DepotId, p.ManifestId)));

        var statePath = Path.Combine(stateDir, "state.json");
        var state = new InstallState();
        try { if (File.Exists(statePath)) state = JsonSerializer.Deserialize<InstallState>(File.ReadAllText(statePath)) ?? new(); } catch { }
        var forceValidate = validate || !state.Complete;
        state.Complete = false;
        File.WriteAllText(statePath, JsonSerializer.Serialize(state));

        var servers = await LoadServersAsync(link);
        var cdn = new Client(link.Client);
        var authTokens = new ConcurrentDictionary<(uint, string), string?>();
        var serverCursor = 0;

        Server NextServer() => servers[Math.Abs(Interlocked.Increment(ref serverCursor)) % servers.Count];

        async Task<string?> TokenFor(uint depotId, Server server)
        {
            if (authTokens.TryGetValue((depotId, server.Host!), out var cached)) return cached;
            var reply = await link.Content.GetCDNAuthToken(AppId, depotId, server.Host!);
            var token = reply.Result == EResult.OK ? reply.Token : null;
            authTokens[(depotId, server.Host!)] = token;
            return token;
        }

        var fullGameDir = Path.GetFullPath(gameDir) + Path.DirectorySeparatorChar;
        string Target(string fileName)
        {
            var full = Path.GetFullPath(Path.Combine(gameDir, fileName));
            if (!full.StartsWith(fullGameDir, StringComparison.OrdinalIgnoreCase)) throw new InvalidOperationException($"Refusing to write outside the install folder: {fileName}");
            return full;
        }

        var newManifests = new Dictionary<uint, DepotManifest>();
        var oldManifests = new Dictionary<uint, DepotManifest?>();
        foreach (var plan in plans)
        {
            ct.ThrowIfCancellationRequested();
            DepotManifest? old = null;
            var oldFile = Path.Combine(stateDir, $"{plan.DepotId}.manifest");
            if (state.Manifests.ContainsKey(plan.DepotId.ToString()) && File.Exists(oldFile))
                try { old = DepotManifest.LoadFromFile(oldFile); } catch { }
            oldManifests[plan.DepotId] = old;

            var manifest = old is { } o && o.ManifestGID == plan.ManifestId ? o : await DownloadManifestAsync(plan);
            newManifests[plan.DepotId] = manifest;
            log($"Depot {plan.DepotId}: {manifest.Files!.Count} files, {FileOps.FormatSize((long)manifest.TotalUncompressedSize)}.");
        }

        async Task<DepotManifest> DownloadManifestAsync(DepotPlan plan)
        {
            for (var attempt = 0; attempt < 8; attempt++)
            {
                ct.ThrowIfCancellationRequested();
                var server = NextServer();
                try
                {
                    var code = await link.Content.GetManifestRequestCode(plan.DepotId, AppId, plan.ManifestId, branch);
                    if (code == 0) throw new InvalidOperationException($"Steam gave no manifest code for depot {plan.DepotId}.");
                    return await cdn.DownloadManifestAsync(plan.DepotId, plan.ManifestId, code, server, plan.Key, null, authTokens.GetValueOrDefault((plan.DepotId, server.Host!)));
                }
                catch (SteamKitWebRequestException e) when (e.StatusCode == HttpStatusCode.Forbidden && !authTokens.ContainsKey((plan.DepotId, server.Host!)))
                {
                    await TokenFor(plan.DepotId, server);
                }
                catch (Exception e) when (e is not OperationCanceledException && e is not InvalidOperationException)
                {
                    log($"Manifest download from {server.Host} failed ({e.Message}); trying another server.");
                }
            }
            throw new InvalidOperationException($"Could not download the file list for depot {plan.DepotId}.");
        }

        var work = new List<(uint DepotId, byte[] Key, DepotManifest.FileData File, string Path, List<DepotManifest.ChunkData> Chunks)>();
        long totalFiles = newManifests.Values.Sum(m => m.Files!.Count(f => !f.Flags.HasFlag(EDepotFileFlag.Directory)));
        long scanned = 0;
        log(forceValidate ? "Checking every file against Steam..." : "Looking for changed files...");
        progress(0);

        var workLock = new object();
        foreach (var plan in plans)
        {
            var manifest = newManifests[plan.DepotId];
            var old = oldManifests[plan.DepotId];
            var oldByName = old?.Files!.ToDictionary(f => f.FileName, StringComparer.OrdinalIgnoreCase);

            foreach (var dir in manifest.Files!.Where(f => f.Flags.HasFlag(EDepotFileFlag.Directory)))
                Directory.CreateDirectory(Target(dir.FileName));

            await Parallel.ForEachAsync(manifest.Files!.Where(f => !f.Flags.HasFlag(EDepotFileFlag.Directory)),
                new ParallelOptions { MaxDegreeOfParallelism = MaxDownloads, CancellationToken = ct },
                async (file, token) =>
                {
                    await Task.Yield();
                    var path = Target(file.FileName);
                    Directory.CreateDirectory(Path.GetDirectoryName(path)!);
                    var needed = NeededChunks(path, file, oldByName, forceValidate);
                    if (needed.Count > 0) lock (workLock) work.Add((plan.DepotId, plan.Key, file, path, needed));
                    var done = Interlocked.Increment(ref scanned);
                    if (done % 25 == 0 || done == totalFiles) progress(done * 100.0 / totalFiles);
                });
        }

        var totalBytes = work.Sum(w => w.Chunks.Sum(c => (long)c.UncompressedLength));
        if (work.Count == 0) log("Everything is already up to date.");
        else
        {
            log($"Downloading {FileOps.FormatSize(totalBytes)} in {work.Count} file(s)...");
            progress(0);
            long downloaded = 0;
            var queue = work.SelectMany(w => w.Chunks.Select(c => (w.DepotId, w.Key, w.Path, Chunk: c))).ToList();
            await Parallel.ForEachAsync(queue, new ParallelOptions { MaxDegreeOfParallelism = MaxDownloads, CancellationToken = ct }, async (item, token) =>
            {
                var buffer = ArrayPool<byte>.Shared.Rent((int)item.Chunk.UncompressedLength);
                try
                {
                    var written = 0;
                    for (var attempt = 0; attempt < 10 && written == 0; attempt++)
                    {
                        token.ThrowIfCancellationRequested();
                        var server = NextServer();
                        try
                        {
                            written = await cdn.DownloadDepotChunkAsync(item.DepotId, item.Chunk, server, buffer, item.Key, null, authTokens.GetValueOrDefault((item.DepotId, server.Host!)));
                        }
                        catch (SteamKitWebRequestException e) when (e.StatusCode == HttpStatusCode.Forbidden && !authTokens.ContainsKey((item.DepotId, server.Host!)))
                        {
                            await TokenFor(item.DepotId, server);
                        }
                        catch (Exception e) when (e is not OperationCanceledException)
                        {
                            if (attempt >= 4) log($"Chunk download from {server.Host} failed ({e.Message}).");
                        }
                    }
                    if (written == 0) throw new InvalidOperationException("No Steam download server could deliver a file piece. Try again in a minute.");

                    await using var fs = new FileStream(item.Path, FileMode.Open, FileAccess.Write, FileShare.ReadWrite);
                    fs.Seek((long)item.Chunk.Offset, SeekOrigin.Begin);
                    await fs.WriteAsync(buffer.AsMemory(0, written), token);
                    progress(Interlocked.Add(ref downloaded, written) * 100.0 / totalBytes);
                }
                finally { ArrayPool<byte>.Shared.Return(buffer); }
            });
        }

        foreach (var plan in plans)
        {
            var old = oldManifests[plan.DepotId];
            if (old != null)
            {
                var keep = newManifests[plan.DepotId].Files!.Select(f => f.FileName).ToHashSet(StringComparer.OrdinalIgnoreCase);
                foreach (var gone in old.Files!.Where(f => !f.Flags.HasFlag(EDepotFileFlag.Directory) && !keep.Contains(f.FileName)))
                {
                    var path = Target(gone.FileName);
                    if (File.Exists(path)) { File.Delete(path); log($"Removed {gone.FileName}"); }
                }
            }
            newManifests[plan.DepotId].SaveToFile(Path.Combine(stateDir, $"{plan.DepotId}.manifest"));
            state.Manifests[plan.DepotId.ToString()] = plan.ManifestId;
        }
        state.Complete = true;
        File.WriteAllText(statePath, JsonSerializer.Serialize(state));
        AppManifest.Write(gameDir, buildId, branch);
        progress(100);
        log($"Installed build {buildId} of '{branch}'.");
    }

    static List<DepotManifest.ChunkData> NeededChunks(
        string path, DepotManifest.FileData file, Dictionary<string, DepotManifest.FileData>? oldByName, bool forceValidate)
    {
        var chunks = file.Chunks!.OrderBy(c => c.Offset).ToList();
        var info = new FileInfo(path);
        if (!info.Exists)
        {
            using var created = File.Create(path);
            created.SetLength((long)file.TotalSize);
            return chunks;
        }

        if (!forceValidate && info.Length == (long)file.TotalSize
            && oldByName != null && oldByName.TryGetValue(file.FileName, out var before) && before.FileHash.AsSpan().SequenceEqual(file.FileHash))
            return new();

        using var fs = new FileStream(path, FileMode.Open, FileAccess.ReadWrite, FileShare.ReadWrite);
        if (fs.Length != (long)file.TotalSize) fs.SetLength((long)file.TotalSize);
        var needed = new List<DepotManifest.ChunkData>();
        var buffer = ArrayPool<byte>.Shared.Rent(chunks.Count == 0 ? 1 : (int)chunks.Max(c => c.UncompressedLength));
        try
        {
            foreach (var chunk in chunks)
            {
                fs.Seek((long)chunk.Offset, SeekOrigin.Begin);
                var length = (int)chunk.UncompressedLength;
                var read = 0;
                while (read < length)
                {
                    var n = fs.Read(buffer, read, length - read);
                    if (n == 0) break;
                    read += n;
                }
                if (read < length || Adler(buffer, length) != chunk.Checksum) needed.Add(chunk);
            }
        }
        finally { ArrayPool<byte>.Shared.Return(buffer); }
        return needed;
    }

    static uint Adler(byte[] data, int length)
    {
        uint a = 0, b = 0;
        for (var i = 0; i < length; i++)
        {
            a = (a + data[i]) % 65521;
            b = (b + a) % 65521;
        }
        return a | (b << 16);
    }

    static async Task<List<Server>> LoadServersAsync(SteamLink link)
    {
        var all = await link.Content.GetServersForSteamPipe();
        var list = new List<Server>();
        foreach (var s in all.Where(s => (s.AllowedAppIds.Length == 0 || s.AllowedAppIds.Contains(AppId)) && (s.Type == "SteamCache" || s.Type == "CDN"))
                     .OrderBy(s => s.WeightedLoad))
            for (var i = 0; i < s.NumEntries; i++) list.Add(s);
        if (list.Count == 0) throw new InvalidOperationException("Steam gave no download servers.");
        return list;
    }
}