using System;
using System.Diagnostics.CodeAnalysis;
using System.Runtime.InteropServices;
using System.Text;

using PeNet;
using PeNet.Header.Pe;

namespace InnoPatcher;

sealed class InnoLoader {
  //
  // Based on innounp repository, file innounp.dpr
  //

  private const int HeaderSize = 64;

  private static ReadOnlySpan<byte> SignaturePrefix => "Inno Setup Setup Data ("u8;
  private static ReadOnlySpan<byte> SignatureIsxPrefix => "My Inno Setup Extensions Setup Data ("u8;
  private static ReadOnlySpan<byte> UnicodeFlagLower => "(u)"u8;
  private static ReadOnlySpan<byte> UnicodeFlagUpper => "(U)"u8;
  private static ReadOnlySpan<byte> ISXFlag => "with ISX"u8;

  private static readonly byte[] SignatureOffsetTableId6 = new byte[] { 0x72, 0x44, 0x6C, 0x50, 0x74, 0x53, 0x30, 0x36, 0x87, 0x65, 0x56, 0x78 }; // rDlPtS06 and some hex mess 
  private static readonly byte[] SignatureOffsetTableId7 = new byte[] { 0x72, 0x44, 0x6C, 0x50, 0x74, 0x53, 0x30, 0x37, 0x87, 0x65, 0x56, 0x78 }; // rDlPtS07 and some hex mess 
  private static readonly byte[] SignatureOffsetTableId8 = new byte[] { 0x72, 0x44, 0x6C, 0x50, 0x74, 0x53, 0xCD, 0xE6, 0xD7, 0x7B, 0x0B, 0x2A }; // rDlPtS and some hex mess

  private const int RCDATA_Id = 10;
  private const int InnoResource_Id = 11111;

  private readonly Stream _file;
  private readonly InnoLoaderOffsetTable _offsets;
  private readonly Version _version;
  private readonly bool _isUnicode;
  private readonly bool _isISX;

  private InnoLoader(Stream file, InnoLoaderOffsetTable offsets, Version version, bool isUnicode, bool isISX) {
    this._file = file;
    this._offsets = offsets;
    this._version = version;
    this._isUnicode = isUnicode;
    this._isISX = isISX;
  }

  /// <summary>
  /// Creates the Inno loader from a file with PE format.
  /// </summary>
  internal static InnoLoader FromPE(Stream peFile) {
    if (!peFile.CanSeek)
      throw new ArgumentOutOfRangeException();

    InnoLoaderOffsetTable offsets;
    if (!TryGetOffsetTableFromResource(peFile, out offsets)
      && !TryGetOffsetTableFromFile(peFile, out offsets))
      throw new FormatException("Failed to find the offsets table.");

    byte[] header = new byte[HeaderSize];
    peFile.Seek(offsets.Offset0, SeekOrigin.Begin);
    int read, offset = 0;
    while ((read = peFile.Read(header, offset, HeaderSize - offset)) != 0)
      offset += read;
    if (offset < HeaderSize)
      throw new FormatException();

    if (!TryParseHeader(header, out Version? version, out bool isUnicode, out bool isISX))
      throw new FormatException("Failed to parse header.");
    
    return new InnoLoader(peFile, offsets, version, isUnicode, isISX);
  }

  public Version Version => _version;

  internal int Version2 => _version.Major * 1000 + _version.Minor * 100 + _version.Build;

  public bool IsUnicode => _isUnicode;

  public bool IsISX => _isISX;

  public InnoSetup GetSetup() {
    _file.Seek(_offsets.Offset0 + HeaderSize, SeekOrigin.Begin);
    if (!InnoSetup.TryParse(_file, Version2, out InnoSetup? result))
      throw new FormatException("Failed to parse Inno Setup");
    return result;
  }

  private static bool TryGetOffsetTableFromFile(Stream peFile, out InnoLoaderOffsetTable result) {
    throw new NotImplementedException();
  }

  private static bool TryGetOffsetTableFromResource(Stream peFile, out InnoLoaderOffsetTable result) {
    if (PeFile.TryParse(peFile, out PeFile? pe) && pe!.ImageResourceDirectory != null) {
      foreach (var i in pe!.ImageResourceDirectory!.DirectoryEntries!)
        if (i.IsIdEntry && i.ID == RCDATA_Id && i.DataIsDirectory) {
          foreach (var j in i.ResourceDirectory!.DirectoryEntries!)
            if (j.IsIdEntry && j.ID == InnoResource_Id
              && j.DataIsDirectory && j.ResourceDirectory!.NumberOfIdEntries == 1) {
                ImageResourceDataEntry a = j.ResourceDirectory!.DirectoryEntries![0]!.ResourceDataEntry!;
                if (TryGetOffsetFromRva(a.OffsetToData, pe, out uint tableOffset)) {
                  int read, offset = 0;
                  byte[] buffer = new byte[a.Size1];
                  peFile.Seek(tableOffset, SeekOrigin.Begin);
                  while ((read = peFile.Read(buffer, offset, buffer.Length - offset)) != 0)
                    offset += read;
                  if (offset == buffer.Length && InnoLoaderOffsetTable.TryParse(buffer, out result))
                    return true;
                }
                break;
              }
          break;
        }
    }
    result = default;
    return false;
  }

  private static bool TryParseHeader(ReadOnlySpan<byte> header, [NotNullWhen(true)] out Version? ver, out bool isUnicode, out bool isISX) {
    ver = null;
    isUnicode = false;
    isISX = false;
    if (header.StartsWith(SignaturePrefix))
      header = header[SignaturePrefix.Length..];
    else if (header.StartsWith(SignatureIsxPrefix))
      header = header[SignatureIsxPrefix.Length..];
    else
      return false;

    if (header.Length < 2 || header[0] < '1' || header[0] > '9' || header[1] != '.')
      return false;

    int verLen = 2;
    while (verLen < header.Length)
      if (header[verLen] == ')')
        break;
      else
        verLen++;
    if (verLen == header.Length)
      return false;

    var versionB = header.Slice(0, verLen);
    header = header.Slice(verLen + 1);
    if (!Version.TryParse(Encoding.ASCII.GetString(versionB), out ver))
      return false;

    isUnicode = header.IndexOf(UnicodeFlagLower) >= 0 || header.IndexOf(UnicodeFlagUpper) >= 0;
    isISX = header.IndexOf(ISXFlag) >= 0;
    return true;
  }

  /// <summary>
  /// Converts relative virtual address (RVA) to a file offset.
  /// </summary>
  private static bool TryGetOffsetFromRva(ulong rva, PeFile pe, out uint offset) {
    foreach (var k in pe.ImageSectionHeaders!) {
      if (k.VirtualAddress > rva) {
        break;
      } else if (rva < k.VirtualAddress + k.VirtualSize) {
        offset = (uint)(rva - k.VirtualAddress) + k.PointerToRawData;
        return true;
      }
    }
    offset = 0;
    return false;
  }

  internal static bool VerifyCRC(ReadOnlySpan<byte> data, uint expectedCRC) {
    return System.IO.Hashing.Crc32.HashToUInt32(data) == expectedCRC;
  }

  struct InnoLoaderOffsetTable {
    internal uint Version; // SetupLdrOffsetTableVersion
    internal uint TotalSize; // Minimum expected size of setup.exe
    internal uint OffsetEXE; // Offset of compressed setup.e32
    internal uint UncompressedSizeEXE; // Size of setup.e32 before compression
    internal int CRCEXE; // CRC of setup.e32 before compression
    internal uint Offset0; // Offset of embedded setup-0.bin data
    internal uint Offset1; // Offset of embedded setup-1.bin data or 0 when DiskSpanning=yes
    internal uint TableCRC; // CRC of all prior fields in this record

    internal static bool TryParse(ReadOnlySpan<byte> data, [NotNullWhen(true)] out InnoLoaderOffsetTable result) {
      const int IdLength = 12;
      if (data.Length >= IdLength) {
        ReadOnlySpan<byte> id = data.Slice(0, IdLength);
        ReadOnlySpan<byte> payload = data.Slice(IdLength);

        if (id.SequenceEqual(SignatureOffsetTableId8) && payload.Length == Marshal.SizeOf<InnoLoaderOffsetTableV8>()) {
          result = MemoryMarshal.Cast<byte, InnoLoaderOffsetTableV8>(payload)[0].Unify();
          return VerifyCRC(data.Slice(0, data.Length - 4), result.TableCRC);
        } else if (id.SequenceEqual(SignatureOffsetTableId7) && payload.Length == Marshal.SizeOf<InnoLoaderOffsetTableV7>()) {
          result = MemoryMarshal.Cast<byte, InnoLoaderOffsetTableV7>(payload)[0].Unify();
          return VerifyCRC(data.Slice(0, data.Length - 4), result.TableCRC);
        } else if (id.SequenceEqual(SignatureOffsetTableId6) && payload.Length == Marshal.SizeOf<InnoLoaderOffsetTableV6>()) {
          result = MemoryMarshal.Cast<byte, InnoLoaderOffsetTableV6>(payload)[0].Unify();
          return VerifyCRC(data.Slice(0, data.Length - 4), result.TableCRC);
        }
      }

      result = default;
      return false;
    }
  }

  [StructLayout(LayoutKind.Sequential, Pack = 1)]
  struct InnoLoaderOffsetTableV6 {
    internal uint TotalSize; // Minimum expected size of setup.exe
    internal uint OffsetEXE; // Offset of compressed setup.e32
    internal uint CompressedSizeEXE;
    internal uint UncompressedSizeEXE; // Size of setup.e32 before compression
    internal int CRCEXE; // CRC of setup.e32 before compression
    internal uint Offset0; // Offset of embedded setup-0.bin data
    internal uint Offset1; // Offset of embedded setup-1.bin data or 0 when DiskSpanning=yes
    internal uint TableCRC; // CRC of all prior fields in this record

    internal InnoLoaderOffsetTable Unify() {
      InnoLoaderOffsetTable a;
      a.Version = 0;
      a.TotalSize = TotalSize;
      a.OffsetEXE = OffsetEXE;
      a.UncompressedSizeEXE = UncompressedSizeEXE;
      a.CRCEXE = CRCEXE;
      a.Offset0 = Offset0;
      a.Offset1 = Offset1;
      a.TableCRC = TableCRC;
      return a;
    }
  }

  [StructLayout(LayoutKind.Sequential, Pack = 1)]
  struct InnoLoaderOffsetTableV7 {
    internal uint TotalSize; // Minimum expected size of setup.exe
    internal uint OffsetEXE; // Offset of compressed setup.e32
    internal uint UncompressedSizeEXE; // Size of setup.e32 before compression
    internal int CRCEXE; // CRC of setup.e32 before compression
    internal uint Offset0; // Offset of embedded setup-0.bin data
    internal uint Offset1; // Offset of embedded setup-1.bin data or 0 when DiskSpanning=yes
    internal uint TableCRC; // CRC of all prior fields in this record

    internal InnoLoaderOffsetTable Unify() {
      InnoLoaderOffsetTable a;
      a.Version = 0;
      a.TotalSize = TotalSize;
      a.OffsetEXE = OffsetEXE;
      a.UncompressedSizeEXE = UncompressedSizeEXE;
      a.CRCEXE = CRCEXE;
      a.Offset0 = Offset0;
      a.Offset1 = Offset1;
      a.TableCRC = TableCRC;
      return a;
    }
  }

  [StructLayout(LayoutKind.Sequential, Pack = 1)]
  struct InnoLoaderOffsetTableV8 {
    internal uint Version; // SetupLdrOffsetTableVersion
    internal uint TotalSize; // Minimum expected size of setup.exe
    internal uint OffsetEXE; // Offset of compressed setup.e32
    internal uint UncompressedSizeEXE; // Size of setup.e32 before compression
    internal int CRCEXE; // CRC of setup.e32 before compression
    internal uint Offset0; // Offset of embedded setup-0.bin data
    internal uint Offset1; // Offset of embedded setup-1.bin data or 0 when DiskSpanning=yes
    internal uint TableCRC; // CRC of all prior fields in this record

    internal InnoLoaderOffsetTable Unify() {
      InnoLoaderOffsetTable a;
      a.Version = Version;
      a.TotalSize = TotalSize;
      a.OffsetEXE = OffsetEXE;
      a.UncompressedSizeEXE = UncompressedSizeEXE;
      a.CRCEXE = CRCEXE;
      a.Offset0 = Offset0;
      a.Offset1 = Offset1;
      a.TableCRC = TableCRC;
      return a;
    }
  }
}