MSBT

using System.Runtime.InteropServices;
using System.Text;

namespace NLItemParamDump
{
    public sealed class MSBT
    {
        public enum MSBTEncoding : byte
        {
            UTF8 = 0,
            Unicode = 1,
            UTF32 = 3
        }

        [StructLayout(LayoutKind.Explicit, CharSet = CharSet.Ansi, Pack = 1, Size = 32)]
        public unsafe struct MSBTHeader
        {
            public const string MSBTMagic = "MsgStdBn";

            [FieldOffset(0x00)]
            public fixed byte Magic[8];

            [FieldOffset(0x08)]
            public ushort ByteOrderMark;

            [FieldOffset(0x0A)]
            public ushort Padding0A; // Verify this.

            [FieldOffset(0x0C)]
            public MSBTEncoding Encoding;

            [FieldOffset(0x0D)]
            public byte Version; // Unknown. Verify this.

            [FieldOffset(0x0E)]
            public ushort SectionCount; // Verify this.

            [FieldOffset(0x10)]
            public ushort Padding10; // Verify this.

            [FieldOffset(0x12)]
            public uint FileSize;

            // Check 0x14 - 0x1F.

            public string GetMagic()
            {
                char[] buf = new char[8];
                for (int i = 0; i < 8; i++)
                {
                    buf[i] = (char)Magic[i];
                }
                return new string(buf);
            }
        }

        [StructLayout(LayoutKind.Explicit, CharSet = CharSet.Ansi, Pack = 1, Size = 16)]
        public unsafe struct MSBTSectionHeader
        {
            [FieldOffset(0x0)]
            public fixed byte Magic[4];

            [FieldOffset(0x4)]
            public uint Size;

            public string GetMagic()
            {
                char[] buf = new char[4];
                for (int i = 0; i < 4; i++)
                {
                    buf[i] = (char)Magic[i];
                }
                return new string(buf);
            }
        }

        public static bool IsMSBT(in byte[] buffer) => Encoding.ASCII.GetString(buffer, 0, 8) == MSBTHeader.MSBTMagic;

        private Encoding? _Encoding;

        public MSBTHeader Header { get; private set; }
        public Endianness Endianness { get; private set; }
        public string? Name { get; private set; }
        public List<string> Messages => GetSection<TXTSection>()?.Entries.Select(bArr => Encoding.GetString(bArr)).ToList() ?? new List<string>();
        public List<byte[]> RawMessageData => GetSection<TXTSection>()?.Entries ?? new List<byte[]>();

        private readonly List<MSBTSection> Sections = new();

        public Encoding Encoding { get => _Encoding ?? Encoding.Default; private set => _Encoding = value; }

        public int GetMaximumMessageID()
        {
            if (HasSection<NLISection>())
                return (int)GetSection<NLISection>()!.Ids.Max(o => o.Value);
            else if (HasSection<TXTSection>())
                return GetSection<TXTSection>()!.Entries.Count - 1;
            else
                return -1;
        }

        public string? GetMessage(int messageId)
        {
            if (messageId < 0) return null;
            if (HasSection<TXTSection>())
            {
                var txt = GetSection<TXTSection>();
                if (txt!.Entries.Count > messageId) return txt!.GetMessage(messageId, Encoding);
            }

            return null;
        }

        public string? GetStrippedMessage(int messageId)
        {
            if (messageId < 0) return null;
            if (HasSection<TXTSection>())
            {
                var txt = GetSection<TXTSection>()!;
                if (txt.Entries.Count > messageId) return txt.GetStrippedMessage(messageId, Encoding);
            }

            return null;
        }

        public string? GetMessage(string label)
        {
            if (HasSection<LBLSection>() && HasSection<TXTSection>())
            {
                var lbl = GetSection<LBLSection>()!;
                var txt = GetSection<TXTSection>()!;
                int idx = (int?)lbl.Labels.FirstOrDefault(o => o.Name == label)?.Index ?? -1;
                if (idx == -1 || idx >= txt.EntryCount) return null;
                return txt.GetMessage(idx, Encoding);
            }
            return null;
        }

        public string? GetStrippedMessage(string label)
        {
            if (HasSection<LBLSection>() && HasSection<TXTSection>())
            {
                var lbl = GetSection<LBLSection>()!;
                var txt = GetSection<TXTSection>()!;
                int idx = (int?)lbl.Labels.FirstOrDefault(o => o.Name == label)?.Index ?? -1;
                if (idx == -1 || idx >= txt.EntryCount) return null;
                return txt.GetStrippedMessage(idx, Encoding);
            }
            return null;
        }

        public string[]? GetMessages()
        {
            if (HasSection<TXTSection>())
                return GetSection<TXTSection>()!.Entries.Select(bArr => Encoding.GetString(bArr)).ToArray();

            return null;
        }

        public string[]? GetLabels()
        {
            if (HasSection<LBLSection>())
                return GetSection<LBLSection>()!.Labels.Select(o => o.Name).ToArray();
            return null;
        }

        public bool SetMessages(string[] messages)
        {
            if (messages == null || HasSection<TXTSection>() == false) return false;

            var txt = GetSection<TXTSection>()!;
            txt.Entries = messages.Select(s => Encoding.GetBytes(s)).ToList();
            txt.EntryCount = (uint)messages.Length;

            // TODO: Labels?
            return true;
        }

        public bool UpdateMessage(int messageId, string message)
        {
            if (messageId < 0 || HasSection<TXTSection>() == false) return false;

            // TODO: Do we want this? How would we handle an overlap where the message id is the same as an index in the TXT section?
            if (HasSection<NLISection>())
            {
                var nli = GetSection<NLISection>();
                if (nli != null && nli.Ids.ContainsValue((uint)messageId))
                    messageId = (int)nli.Ids.First(o => o.Value == (uint)messageId).Key;
            }

            var txt = GetSection<TXTSection>();
            if (txt == null || messageId >= txt.Entries.Count) return false;
            txt.Entries[messageId] = Encoding.GetBytes(message);

            return true;
        }

        private bool TryDetermineEndianness(Stream stream)
        {
            // Default to little-endian.
            using (var reader = new EndianReader(stream, Endianness.Little, true))
            {
                var header = reader.ReadStruct<MSBTHeader>();
                if (header.ByteOrderMark == 0xFEFF)
                {
                    Header = header;
                    Endianness = Endianness.Little;
                    return true;
                }
            }

            // Now try big-endian.
            stream.Seek(0, SeekOrigin.Begin);
            using (var reader = new EndianReader(stream, Endianness.Big, true))
            {
                var header = reader.ReadStruct<MSBTHeader>();
                if (header.ByteOrderMark == 0xFEFF)
                {
                    Header = header;
                    Endianness = Endianness.Big;
                    return true;
                }
            }

            return false; // Byte order mark is invalid.
        }

        private bool ValidateHeader(Stream stream)
        {
            if (Header.GetMagic() != MSBTHeader.MSBTMagic) return false;
            if (Header.ByteOrderMark != 0xFEFF) return false;
            if (Header.FileSize != stream.Length) return false;

            return true;
        }

        private bool HasSection<T>() where T : MSBTSection => Sections.Any(s => s is T);
        private T? GetSection<T>() where T : MSBTSection => Sections.First(s => s is T) as T;

        private uint GetTextHash(string text, uint numLabelGroups) => text.Aggregate(0u, (current, c) => current * 0x492 + c) % numLabelGroups;

        private void LoadSections(EndianReader reader)
        {
            reader.BaseStream.Seek(0x20, SeekOrigin.Begin);
            while (reader.BaseStream.Position < reader.BaseStream.Length)
            {
                switch (reader.PeekString(4))
                {
                    case "LBL1":
                        Sections.Add(new LBLSection(reader));
                        break;
                    case "ATR1":
                        Sections.Add(new ATRSection(reader));
                        break;
                    case "NLI1":
                        Sections.Add(new NLISection(reader));
                        break;
                    case "ATO1":
                        Sections.Add(new ATOSection(reader));
                        break;
                    case "TSY1":
                        Sections.Add(new TSYSection(reader));
                        break;
                    case "TXT2":
                        Sections.Add(new TXTSection(reader));
                        break;
                }
            }
        }

        public bool Load(in byte[] data)
        {
            using MemoryStream ms = new(data);
            return Load(ms);
        }

        public bool Load(Stream stream)
        {
            // Sanity checks first.
            if (stream.Length < 32) throw new ArgumentException("The input stream cannot be smaller than the size of an MSBT header (32 bytes)!");
            if (TryDetermineEndianness(stream) == false) throw new ArgumentException("The header's byte order mark was invalid! It must be FFFE or FEFF!");
            if (ValidateHeader(stream) == false) throw new ArgumentException("The header failed to pass validation checks!");

            Name = (stream is FileStream fs) ? Path.GetFileNameWithoutExtension(fs.Name) : "msbt";

            // Set encoding based on encoding value in header & endianness.
            Encoding = Header.Encoding switch
            {
                MSBTEncoding.UTF8 => Encoding.UTF8,
                MSBTEncoding.Unicode => Endianness == Endianness.Big ? Encoding.BigEndianUnicode : Encoding.Unicode,
                MSBTEncoding.UTF32 => Encoding.UTF32,
                _ => throw new InvalidOperationException("This MSBT file's encoding format is unknown. Please open an issue on the GitHub page!"),
            };

            // Load the MSBT sections.
            try
            {
                using var reader = new EndianReader(stream, Endianness);
                LoadSections(reader);
            }
            catch
            {
                return false;
            }

            if (Sections.Count != Header.SectionCount) throw new ArgumentException("The header's section count does not match the real section count!");
            return true;
        }

        public bool Save(string path)
        {
            throw new NotImplementedException();
        }

        public void Dump(string path)
        {
            if (!HasSection<TXTSection>() || Name == null) return;

            string dir = Path.Combine(path, Name);
            Directory.CreateDirectory(dir);
            TXTSection? textData = GetSection<TXTSection>();
            LBLSection? labels = HasSection<LBLSection>() ? GetSection<LBLSection>() : null;
            if (textData == null || labels == null) return;
            for (int i = 0; i < textData.EntryCount; i++)
            {
                string name = labels?.Labels.First(l => l.Index == i).Name ?? i.ToString();
                File.WriteAllText(Path.Combine(dir, name), textData.GetMessage(i, Encoding));
            }
        }

        private abstract class MSBTSection
        {
            public MSBTSectionHeader Header;
            public MSBTSection(EndianReader reader) => Header = reader.ReadStruct<MSBTSectionHeader>();
        }

        private sealed class LBLSection : MSBTSection
        {
            [StructLayout(LayoutKind.Explicit, Pack = 1, Size = 8)]
            public struct Group
            {
                [FieldOffset(0)]
                public uint LabelCount;
                [FieldOffset(4)]
                public uint Offset;
            }

            public sealed class Label
            {
                public string Name;
                public uint Index;
                public uint Checksum;

                public Label(string name, uint index, uint checksum)
                {
                    Name = name;
                    Index = index;
                    Checksum = checksum;
                }
            }

            public uint GroupCount;
            public List<Group> Groups;
            public List<Label> Labels;

            public LBLSection(EndianReader reader) : base(reader)
            {
                Groups = new List<Group>();
                Labels = new List<Label>();

                GroupCount = reader.ReadUInt32();

                // Load groups
                for (var i = 0; i < GroupCount; i++)
                    Groups.Add(reader.ReadStruct<Group>());

                // Load labels
                for (var x = 0; x < Groups.Count; x++)
                    for (var i = 0; i < Groups[x].LabelCount; i++)
                        Labels.Add(new Label(Encoding.ASCII.GetString(reader.ReadBytes(reader.ReadByte())), reader.ReadUInt32(), (uint)x));

                reader.Align(16);
            }
        }

        private sealed class NLISection : MSBTSection
        {
            public uint IdCount;
            public Dictionary<uint, uint> Ids;

            public NLISection(EndianReader reader) : base(reader)
            {
                Ids = new Dictionary<uint, uint>();
                IdCount = reader.ReadUInt32();

                for (var i = 0; i < IdCount; i++)
                {
                    var id = reader.ReadUInt32();
                    var index = reader.ReadUInt32();
                    Ids.Add(index, id);
                }

                reader.Align(16);
            }
        }

        private sealed class ATOSection : MSBTSection
        {
            // TODO: ATO1 section needs to be researched.
            public byte[] Data;

            public ATOSection(EndianReader reader) : base(reader)
            {
                Data = reader.ReadBytes((int)Header.Size);
                reader.Align(16);
            }
        }

        private sealed class ATRSection : MSBTSection
        {
            public uint AttributeCount;
            public uint AttributeSize;

            public List<byte[]> Attributes;

            public ATRSection(EndianReader reader) : base(reader)
            {
                Attributes = new List<byte[]>();

                AttributeCount = reader.ReadUInt32();
                AttributeSize = reader.ReadUInt32();

                for (var i = 0; i < AttributeCount; i++)
                    Attributes.Add(reader.ReadBytes((int)AttributeSize));

                reader.Align(16);
            }
        }

        private sealed class TSYSection : MSBTSection
        {
            // TODO: TSY1 section needs to be researched.
            public uint[] Data;

            public TSYSection(EndianReader reader) : base(reader)
            {
                Data = new uint[Header.Size / sizeof(uint)];
                for (int i = 0; i < Data.Length; i++)
                {
                    Data[i] = reader.ReadUInt32();
                }

                reader.Align(16);
            }
        }

        private sealed class TXTSection : MSBTSection
        {
            public uint EntryCount;

            public List<byte[]> Entries;

            public TXTSection(EndianReader reader) : base(reader)
            {
                Entries = new List<byte[]>();
                EntryCount = reader.ReadUInt32();

                // Load string offsets.
                var offsets = new uint[EntryCount];
                for (var i = 0; i < EntryCount; i++)
                    offsets[i] = reader.ReadUInt32();

                for (var i = 0; i < EntryCount; i++)
                {
                    uint size;
                    if (i == EntryCount - 1)
                        size = Header.Size - offsets[i];
                    else
                        size = offsets[i + 1] - offsets[i];

                    Entries.Add(reader.ReadBytes((int)size));
                }

                reader.Align(16);
            }

            public string? GetMessage(int idx, Encoding encoding)
            {
                return idx > -1 && idx < Entries.Count ? encoding.GetString(Entries[idx]) : null;
            }

            public byte[]? GetRawMessage(int idx)
            {
                return idx > -1 && idx < Entries.Count ? Entries[idx] : null;
            }

            public byte[]? GetRawMessageStripped(int idx)
            {
                byte[]? data = GetRawMessage(idx);
                if (data == null) return null;
                List<byte> stripped = new(data.Length);
                for (int i = 0; i < data.Length; i += 2)
                {
                    if (data[i + 1] == 0 && (data[i] == 0x0E || data[i] == 0x0F))
                    {
                        int size = 4 * sizeof(ushort) + BitConverter.ToUInt16(data, i + 3 * sizeof(ushort)); /* u16 cmd, u16 group, u16 index, u16 arg_size, u8 args[] */
                        i += size - 2;
                    }
                    else
                    {
                        stripped.Add(data[i]);
                        stripped.Add(data[i + 1]);
                    }
                }
                return stripped.ToArray();
            }

            public string? GetStrippedMessage(int idx, Encoding encoding)
            {
                byte[]? data = GetRawMessageStripped(idx);
                return data == null ? null : encoding.GetString(data).Replace("\0", "");
            }
        }
    }
}