BeatSaberManager/Bsm.Arcache.Runner/Program.cs

using System;
using System.Collections.Concurrent;
using System.Collections.Generic;
using System.IO;
using System.Text;
using System.Threading;
using System.Threading.Tasks;
using Bsm.Core;
using Bsm.Core.Arcache;

namespace Bsm.Arcache.Runner
{
    internal class Program
    {
        private static void Main(string[] args)
        {
#if DEBUG
            var dir = @"Y:\BeatSaverAll";
            var arcBase = @"Z:\bsa";
            var lenArc = @"Z:\bslp";
#else
            var dir = AppDomain.CurrentDomain.BaseDirectory;
            var arcBase = Path.Combine(dir, "bsa");
            var lenArc = Path.Combine(dir, "bslp");
#endif

            //TODO: speed up, only once lookup, extract egg and dat at once

            CreateJsonArchive(dir, arcBase);
            CreateLenArchiveDirectly(dir, lenArc);

            Console.WriteLine();
            Console.Write("Press ENTER to exit...");
            Console.ReadLine();
        }

        internal static void CreateLenArchiveDirectly(string dir, string lenArc)
        {
            var lst = new List<string>();
            Console.Write("Looking up ogg files...");
            IoUtil.LookupFilesRecursively(dir, (p, f) =>
            {
                var ext = Path.GetExtension(p.cFileName)?.ToLower();
                if (ext != ".egg") return;
                var archivePath = f.TrimStart('/') + "/" + p.cFileName;
                lst.Add(archivePath.ToLower());
            });
            Console.WriteLine($"{lst.Count}");
            var processed = 0;

            var ms = new MemoryStream();
            var writer = new BinaryWriter(ms, Encoding.UTF8, true);

            void AppendEntry(string item)
            {
                var fullPath = Path.Combine(dir, item);

                float len;
                Exception ex = null;
                try
                {
                    len = GetLengthFromOgg(fullPath);
                }
                catch (Exception e)
                {
                    len = -1;
                    ex = e;
                }

                lock (writer)
                {
                    writer.Write(item);
                    writer.Write(len);
                }

                Console.WriteLine($"[{Interlocked.Increment(ref processed)}/{lst.Count}]Added:{fullPath}:{len}:{ex?.Message}");
            }

            Parallel.ForEach(lst, AppendEntry);

            writer.Flush();
            writer.Close();

            var bufRaw = ms.ToArray();
            var bufCompress = Lz4ArchiveBase.CompressPickle(bufRaw);
            File.WriteAllBytes(lenArc, bufCompress);
        }

        internal static void CreateJsonArchive(string dir, string arcBase)
        {
            Console.WriteLine("Creating archive file...");
            var arc = new Lz4ArchiveCreator(arcBase);

            var lst = new List<string>();
            Console.Write("Looking up json files...");
            IoUtil.LookupFilesRecursively(dir, (p, f) =>
            {
                var ext = Path.GetExtension(p.cFileName)?.ToLower();
                if (ext != ".dat") return;
                var archivePath = f.TrimStart('/') + "/" + p.cFileName;
                lst.Add(archivePath.ToLower());
            });

            Console.WriteLine($"{lst.Count}");
            var processed = 0;

            void AddToArchive(string archivePath)
            {
                var fullPath = Path.Combine(dir, archivePath);
                arc.AppendFile(fullPath, archivePath.ToLower());
                Console.WriteLine($"[{Interlocked.Increment(ref processed)}/{lst.Count}]Added:{fullPath}");
            }

            //foreach (var item in lst) AddToArchive(item);
            Parallel.ForEach(lst, AddToArchive);

            arc.Close();
        }

        internal static void CreateLenArchive(string dir, string lenArc)
        {
            var lst = new List<string>();
            Console.Write("Looking up len files...");
            IoUtil.LookupFilesRecursively(dir, (p, f) =>
            {
                var ext = Path.GetExtension(p.cFileName)?.ToLower();
                if (ext != ".len") return;
                var archivePath = f.TrimStart('/') + "/" + p.cFileName;
                lst.Add(archivePath);
            });
            Console.WriteLine($"{lst.Count}");

            var ms = new MemoryStream();
            var writer = new BinaryWriter(ms, Encoding.UTF8, true);

            for (var i = 0; i < lst.Count; i++)
            {
                var fullPath = Path.Combine(dir, lst[i]);
                writer.Write(lst[i]);
                writer.Write(File.ReadAllBytes(fullPath));
                Console.WriteLine($"[{i + 1}/{lst.Count}]Added:{fullPath}");
            }

            writer.Flush();
            writer.Close();

            var bufRaw = ms.ToArray();
            var bufCompress = Lz4ArchiveBase.CompressPickle(bufRaw);
            File.WriteAllBytes(lenArc, bufCompress);
        }

        internal static void CreateOggLengthFile(string dir)
        {
            var dic = new ConcurrentDictionary<string, float>();
            Console.Write("Looking up ogg files...");
            IoUtil.LookupFilesRecursively(dir, (p, f) =>
            {
                if (Path.GetExtension(p.cFileName)?.ToLower() != ".egg") return;
                var archivePath = f.TrimStart('/') + "/" + p.cFileName;
                dic.TryAdd(archivePath, float.NaN);
            });

            Console.WriteLine($"{dic.Count}");
            var processed = 0;

            void Process(KeyValuePair<string, float> keyValuePair)
            {
                var fullPath = Path.Combine(dir, keyValuePair.Key);
                var lenPath = fullPath + ".len";
                if (File.Exists(lenPath))
                {
                    Console.WriteLine($"[{Interlocked.Increment(ref processed)}/{dic.Count}]{keyValuePair.Key}:Skip");
                    return;
                }

                try
                {
                    var len = GetLengthFromOgg(fullPath);
                    File.WriteAllBytes(lenPath, BitConverter.GetBytes(len));
                    Console.WriteLine($"[{Interlocked.Increment(ref processed)}/{dic.Count}]{keyValuePair.Key}:{len}");
                    dic[keyValuePair.Key] = len;
                }
                catch (Exception e)
                {
                    dic[keyValuePair.Key] = -1;
                    Console.WriteLine($"[{Interlocked.Increment(ref processed)}/{dic.Count}]{keyValuePair.Key}:{e.Message}");
                }
            }

            Parallel.ForEach(dic, Process);
            //foreach (var item in dic)
            //{
            //    Process(item);
            //}
        }

        private static byte[] oggBytes = new byte[] { 0x4F, 0x67, 0x67, 0x53, 0x00, 0x04 };

        public static float GetLengthFromOgg(string oggFile)
        {
            using (var fs = File.OpenRead(oggFile))
            using (var br = new BinaryReader(fs, Encoding.ASCII))
            {
                /*
                 * Tries to find the array of bytes from the stream
                 */
                bool FindBytes(byte[] bytes, int searchLength)
                {
                    for (var i = 0; i < searchLength; i++)
                    {
                        var b = br.ReadByte();
                        if (b != bytes[0]) continue;
                        var by = br.ReadBytes(bytes.Length - 1);
                        // hardcoded 6 bytes compare, is fine because all inputs used are 6 bytes
                        // bitwise AND the last byte to read only the flag bit for lastSample searching
                        // shouldn't cause issues finding rate, hopefully
                        if (by[0] == bytes[1] && by[1] == bytes[2] && by[2] == bytes[3] && by[3] == bytes[4] && (by[4] & bytes[5]) == bytes[5]) return true;
                        var index = Array.IndexOf(@by, bytes[0]);
                        if (index != -1)
                        {
                            fs.Position += index - (bytes.Length - 1);
                            i += index;
                        }
                        else
                            i += (bytes.Length - 1);
                    }
                    return false;
                }

                var rate = -1;
                long lastSample = -1;

                //Skip Capture Pattern
                fs.Position = 24;

                //{0x76, 0x6F, 0x72, 0x62, 0x69, 0x73} = "vorbis" in byte values
                var foundVorbis = FindBytes(new byte[] { 0x76, 0x6F, 0x72, 0x62, 0x69, 0x73 }, 256);
                if (foundVorbis)
                {
                    fs.Position += 5;
                    rate = br.ReadInt32();
                }
                else
                {
                    return -1;
                }

                /*
                 * this finds the last occurrence of OggS in the file by checking for a bit flag (0x04)
                 * reads in blocks determined by seekBlockSize
                 * 6144 does not add significant overhead and speeds up the search significantly
                 */
                const int seekBlockSize = 6144;
                const int seekTries = 10; // 60 KiB should be enough for any sane ogg file
                for (var i = 0; i < seekTries; i++)
                {
                    var seekPos = (i + 1) * seekBlockSize * -1;
                    var overshoot = Math.Max((int)(-seekPos - fs.Length), 0);
                    if (overshoot >= seekBlockSize)
                    {
                        break;
                    }
                    fs.Seek(seekPos + overshoot, SeekOrigin.End);
                    var foundOggS = FindBytes(oggBytes, seekBlockSize - overshoot);
                    if (foundOggS)
                    {
                        lastSample = br.ReadInt64();
                        break;
                    }
                }

                if (lastSample == -1)
                {
                    return -1;
                }

                var length = lastSample / (float)rate;
                return length;
            }
        }
    }
}