using System.Buffers.Binary;
using System.Security.Cryptography;
using System.Text;
using System.Text.Json;
using XFESystem.Core;

if (args.Length is < 2 or > 4 || (args.Length == 4 && args[3] != "--independent-kernel"))
    throw new ArgumentException("Usage: PackImage <Console or Desktop Kernel.elf> <new System.img> [widthxheight] [--independent-kernel]. Each image contains one system; dual-system images are no longer generated.");
var independentKernel = args.Length >= 4;
var resolution = args.Length >= 3 ? args[2] : "0x0";
if (resolution.Split('x') is not [var resolutionWidth, var resolutionHeight] ||
    !int.TryParse(resolutionWidth, out var width) || !int.TryParse(resolutionHeight, out var height) ||
    width < 0 || height < 0 || width > 16384 || height > 16384 || (width == 0) != (height == 0))
    throw new ArgumentException("Resolution must be 0x0 (native/default) or a bounded widthxheight.");
var kernelPath = Path.GetFullPath(args[0]);
var outputPath = Path.GetFullPath(args[1]);
var rootPath = Path.GetFullPath(Path.Combine(AppContext.BaseDirectory, "../../../../../")).TrimEnd(Path.DirectorySeparatorChar);
if (!File.Exists(Path.Combine(rootPath, "XFESystem.slnx")))
    throw new InvalidOperationException("Run the tool from this repository's build tree.");
var expectedPrefix = rootPath + Path.DirectorySeparatorChar;
if (!outputPath.StartsWith(expectedPrefix, StringComparison.OrdinalIgnoreCase) ||
    !string.Equals(Path.GetExtension(outputPath), ".img", StringComparison.OrdinalIgnoreCase))
    throw new ArgumentException("The output must be a new .img file inside the repository.");
var kernel = File.ReadAllBytes(kernelPath);
if (kernel.Length < 64 || !kernel.AsSpan(0, 4).SequenceEqual("\u007fELF"u8) || kernel[4] != 2 ||
    BinaryPrimitives.ReadUInt16LittleEndian(kernel.AsSpan(18)) != 62)
    throw new InvalidDataException("Expected an ELF64 x86_64 kernel.");
JsonElement? independentBuild = null;
if (independentKernel)
{
    using var buildDocument = JsonDocument.Parse(File.ReadAllText(kernelPath + ".json"));
    var build = buildDocument.RootElement;
    if (build.GetProperty("Kind").GetString() != "XfeIndependentKernelBuild" || !build.GetProperty("Passed").GetBoolean() ||
        build.GetProperty("KernelSha256").GetString() != Convert.ToHexStringLower(SHA256.HashData(kernel)) ||
        build.GetProperty("ContainsCosmos").GetBoolean() ||
        build.GetProperty("Target").GetString() != "x86_64-none")
        throw new InvalidDataException("The independent kernel needs its exact, Cosmos-free build receipt.");
    foreach (var source in build.GetProperty("SourceSha256").EnumerateObject())
    {
        var sourcePath = Path.GetFullPath(Path.Combine(rootPath, source.Name));
        if (!sourcePath.StartsWith(expectedPrefix, StringComparison.OrdinalIgnoreCase) ||
            Convert.ToHexStringLower(SHA256.HashData(File.ReadAllBytes(sourcePath))) != source.Value.GetString())
            throw new InvalidDataException("Independent kernel source changed: " + source.Name);
    }
    independentBuild = build.Clone();
}
var loaderPath = Path.Combine(rootPath, ".tools/limine/BOOTX64.EFI");
var loader = File.ReadAllBytes(loaderPath);
var shimReceiptPath = Path.Combine(rootPath, "build/CSharpBoot/BootShimBuild/BootShimReceipt.json");
if (!File.Exists(shimReceiptPath))
    throw new InvalidOperationException("Build the EFI display initializer first: Tools/CSharpBoot/BuildBootShim.ps1");
var shimReceiptBytes = File.ReadAllBytes(shimReceiptPath);
using var shimReceiptDocument = JsonDocument.Parse(shimReceiptBytes.AsMemory(shimReceiptBytes.AsSpan().StartsWith(new byte[] { 0xEF, 0xBB, 0xBF }) ? 3 : 0));
var shimReceipt = shimReceiptDocument.RootElement;
if (shimReceipt.GetProperty("Kind").GetString() != "XfeEfiDisplayInitializerBuild" ||
    shimReceipt.GetProperty("Target").GetString() != "x86_64-uefi" ||
    shimReceipt.GetProperty("Language").GetString() != "C# UEFI")
    throw new InvalidDataException("Unexpected EFI display initializer build receipt.");
var shimPath = Path.GetFullPath(shimReceipt.GetProperty("ArtifactPath").GetString()!);
if (!shimPath.StartsWith(expectedPrefix, StringComparison.OrdinalIgnoreCase))
    throw new InvalidDataException("The EFI display initializer must be a repository build artifact.");
var shim = File.ReadAllBytes(shimPath);
var shimSha256 = Convert.ToHexStringLower(SHA256.HashData(shim));
if (shimSha256 != shimReceipt.GetProperty("ArtifactSha256").GetString())
    throw new InvalidDataException("EFI display initializer hash differs from its build receipt.");
var shimSources = new SortedDictionary<string, string>(StringComparer.Ordinal);
foreach (var source in shimReceipt.GetProperty("SourceSha256").EnumerateObject())
{
    var sourcePath = Path.GetFullPath(Path.Combine(rootPath, source.Name));
    if (!sourcePath.StartsWith(expectedPrefix, StringComparison.OrdinalIgnoreCase) ||
        Convert.ToHexStringLower(SHA256.HashData(File.ReadAllBytes(sourcePath))) != source.Value.GetString())
        throw new InvalidDataException("EFI display initializer source changed; rebuild it: " + source.Name);
    shimSources.Add(source.Name, source.Value.GetString()!);
}
if (!shimSources.ContainsKey("Source/XFESystem.Boot/EfiEntry.cs") || !shimSources.ContainsKey("Source/XFESystem.Boot/DisplayContracts.cs"))
    throw new InvalidDataException("EFI display initializer receipt has incomplete source provenance.");
var shimMetadata = JsonSerializer.SerializeToUtf8Bytes(new
{
    Kind = "XfeEfiDisplayInitializer", BootPath = "EFI/BOOT/BOOTX64.EFI", ChainloadPath = "EFI/XFE/Limine.efi",
    ArtifactSha256 = shimSha256, SourceSha256 = shimSources, KernelLanguage = "C# NativeAOT",
    HelperLanguage = "C# UEFI", GopPolicy = "Maximum EDID/GOP intersection, firmware maximum without EDID",
    RefreshRateMeaning = "Monitor/EDID target; no GPU VSync guarantee"
}, new JsonSerializerOptions { WriteIndented = true });
const int sectorSize = 512;
bool usbWorkspace = independentBuild is { } storageBuild && storageBuild.GetProperty("Configuration").GetProperty("StorageEnabled").GetBoolean();
int diskBytes = (usbWorkspace ? 256 : 128) * 1024 * 1024;
Guid diskId = usbWorkspace ? Guid.NewGuid() : new("AA9FCB82-5F9D-4D4C-97D8-584645535953");
Guid partitionId = usbWorkspace ? Guid.NewGuid() : new("EE0D4800-3CC1-4ED0-8ED0-584645535953");
const int first = 2048;
int last = diskBytes / sectorSize - 1;
int partitionLast = last - 33;
const int reserved = 32;
int sectors = partitionLast - first + 1;
var fatSectors = 1;
while (true)
{
    var next = ((sectors - reserved - 2 * fatSectors + 2) * 4 + sectorSize - 1) / sectorSize;
    if (next == fatSectors) break;
    fatSectors = next;
}
var image = new byte[diskBytes];
var table = new byte[128 * 128];
new Guid("C12A7328-F81F-11D2-BA4B-00A0C93EC93B").TryWriteBytes(table);
partitionId.TryWriteBytes(table.AsSpan(16));
U64(table, 32, first); U64(table, 40, partitionLast);
Encoding.Unicode.GetBytes("XFESystem EFI").CopyTo(table, 56);
image[446 + 4] = 0xEE;
U32(image, 446 + 8, 1); U32(image, 446 + 12, last);
image[510] = 0x55; image[511] = 0xAA;
table.CopyTo(image, 2 * sectorSize);
table.CopyTo(image, (last - 32) * sectorSize);
Header(1, last, 2).CopyTo(image, sectorSize);
Header(last, 1, last - 32).CopyTo(image, last * sectorSize);
var partition = first * sectorSize;
var boot = new byte[sectorSize];
boot[0] = 0xEB; boot[1] = 0x58; boot[2] = 0x90;
Encoding.ASCII.GetBytes("XFESYST ").CopyTo(boot, 3);
U16(boot, 11, sectorSize); boot[13] = 1; U16(boot, 14, reserved); boot[16] = 2;
boot[21] = 0xF8; U16(boot, 24, 63); U16(boot, 26, 255);
U32(boot, 28, first); U32(boot, 32, sectors); U32(boot, 36, fatSectors);
U32(boot, 44, 2); U16(boot, 48, 1); U16(boot, 50, 6);
boot[64] = 0x80; boot[66] = 0x29; U32(boot, 67, 0x58464503);
Encoding.ASCII.GetBytes("XFESYSTEM  ").CopyTo(boot, 71);
Encoding.ASCII.GetBytes("FAT32   ").CopyTo(boot, 82);
boot[510] = 0x55; boot[511] = 0xAA;
boot.CopyTo(image, partition); boot.CopyTo(image, partition + 6 * sectorSize);
var info = new byte[sectorSize];
U32(info, 0, 0x41615252); U32(info, 484, 0x61417272);
U32(info, 488, uint.MaxValue); U32(info, 492, uint.MaxValue); U32(info, 508, 0xAA550000);
info.CopyTo(image, partition + sectorSize); info.CopyTo(image, partition + 7 * sectorSize);
var fat = new byte[fatSectors * sectorSize];
U32(fat, 0, 0x0FFFFFF8); U32(fat, 4, 0x0FFFFFFF);
var nextCluster = 2;
var rootCluster = Allocate(1); var efiCluster = Allocate(1);
var efiBootCluster = Allocate(1); var efiXfeCluster = Allocate(1); var systemCluster = Allocate(1);
var shimCluster = Store(shim); var loaderCluster = Store(loader); var kernelCluster = Store(kernel);
var shimMetadataCluster = Store(shimMetadata);
int workspaceCluster = 0, userFilesCluster = 0, usbProgramCluster = 0;
byte[]? usbProgram = null;
var storageSources = new SortedDictionary<string,string>(StringComparer.Ordinal);
if (usbWorkspace)
{
    userFilesCluster = Allocate(1);
    workspaceCluster = Allocate(UsbWorkspaceFormat.ContainerBytes / sectorSize);
    UsbWorkspaceFormat.CreateDescriptor().CopyTo(image, ClusterOffset(workspaceCluster));
    string programPath = Path.Combine(rootPath,"Samples/XFESystem.UsbHello/bin/Release/net10.0/XFESystem.UsbHello.dll");
    usbProgram = File.ReadAllBytes(programPath); usbProgramCluster = Store(usbProgram);
    foreach (string source in new[]{"Samples/XFESystem.UsbHello/Program.cs","Samples/XFESystem.UsbHello/XFESystem.UsbHello.csproj","Samples/XFESystem.UsbHello/packages.lock.json","Source/XFESystem.Core/UsbWorkspaceFormat.cs"})
        storageSources.Add(source,Convert.ToHexStringLower(SHA256.HashData(File.ReadAllBytes(Path.Combine(rootPath,source)))));
}
var edition = independentBuild is { } appBuild && appBuild.TryGetProperty("Edition", out var builtEdition) ? builtEdition.GetString() : "Desktop";
if (edition is not ("Console" or "Desktop")) throw new InvalidDataException("An image must contain exactly one Console or Desktop system.");
var configText = $"timeout: 0\nserial: yes\nverbose: no\nterse: yes\nquiet: yes\nfirmware_logo: no\n\n/XFESystem {(independentKernel ? edition : "CSharp")}\n    protocol: limine\n    resolution: {resolution}x32\n    path: boot():/System/Kernel.elf\n";
var config = Encoding.UTF8.GetBytes(configText);
var configCluster = Store(config);
var rootEntries = new List<byte[]> {
    Entry("EFI", "EFI", true, efiCluster, 0),
    Entry("System", "SYSTEM", true, systemCluster, 0),
    Entry("Limine.conf", "LIMINE  CON", false, configCluster, config.Length) };
if (usbWorkspace)
{
    rootEntries.Add(Entry("Workspace.bin","WORKSPACBIN",false,workspaceCluster,UsbWorkspaceFormat.ContainerBytes));
    rootEntries.Add(Entry("UserFiles","USERFILE",true,userFilesCluster,0));
    Directory(userFilesCluster,Dot(userFilesCluster),DotDot(rootCluster),Entry("UsbHello.dll","USBHELLODLL",false,usbProgramCluster,usbProgram!.Length));
}
Directory(rootCluster,rootEntries.ToArray());
Directory(efiCluster, Dot(efiCluster), DotDot(0), Entry("BOOT", "BOOT", true, efiBootCluster, 0), Entry("XFE", "XFE", true, efiXfeCluster, 0));
Directory(efiBootCluster, Dot(efiBootCluster), DotDot(efiCluster),
    Entry("BOOTX64.EFI", "BOOTX64 EFI", false, shimCluster, shim.Length));
Directory(efiXfeCluster, Dot(efiXfeCluster), DotDot(efiCluster),
    Entry("Limine.efi", "LIMINE  EFI", false, loaderCluster, loader.Length));
var systemEntries = new List<byte[]> { Dot(systemCluster), DotDot(0),
    Entry("Kernel.elf", "KERNEL  ELF", false, kernelCluster, kernel.Length),
    Entry("BootShim.json", "BOOTSHIMJSN", false, shimMetadataCluster, shimMetadata.Length) };
Directory(systemCluster, systemEntries.ToArray());
fat.CopyTo(image, partition + reserved * sectorSize);
fat.CopyTo(image, partition + (reserved + fatSectors) * sectorSize);
if (nextCluster >= sectors - reserved - 2 * fatSectors + 2)
    throw new InvalidDataException("The kernel does not fit this image.");
System.IO.Directory.CreateDirectory(Path.GetDirectoryName(outputPath)!);
using (var output = new FileStream(outputPath, FileMode.CreateNew, FileAccess.Write, FileShare.None))
{
    output.Write(image);
    output.Flush(true);
}
var readback = File.ReadAllBytes(outputPath);
if (!readback.AsSpan().SequenceEqual(image)) throw new IOException("Full image verification failed.");
var receipt = new
{
    Kind = "CSharpGptFat32Image", Edition = edition, KernelPath = kernelPath,
    KernelSha256 = Convert.ToHexStringLower(SHA256.HashData(kernel)),
    BootloaderSha256 = shimSha256, BootloaderPath = "EFI/BOOT/BOOTX64.EFI",
    ChainloaderSha256 = Convert.ToHexStringLower(SHA256.HashData(loader)), ChainloaderPath = "EFI/XFE/Limine.efi",
    BootShimSourceSha256 = shimSources, BootShimBuildReceipt = shimReceiptPath,
    ImagePath = outputPath, ImageSha256 = Convert.ToHexStringLower(SHA256.HashData(image)),
    ImageBytes = diskBytes, FatSectors = fatSectors, AllocatedClusters = nextCluster - 2, Resolution = resolution,
    FullImageReadback = true, KernelLanguage = "C# NativeAOT", FullDotNet10 = false,
    DotNetFramework48 = false, ManagedConsoleIlProfile = "XFESystem.Net10.ConsoleIL.v1",
    UnrestrictedManagedDllLoading = false,
    KernelImplementation = independentKernel ? "XFESystem independent kernel" : "Cosmos 3.0.94 integration",
    IndependentKernelBuild = independentBuild,
    DeploymentPolicy = "SingleEditionOnly", BootEntries = new[] { edition },
    DefaultBootEntry = edition, BootMenuTimeoutSeconds = 0,
    BootDiskGuid = diskId, BootPartitionGuid = partitionId,
    UsbWorkspace = usbWorkspace, WorkspaceFile = usbWorkspace ? "Workspace.bin" : null,
    WorkspaceBytes = usbWorkspace ? UsbWorkspaceFormat.ContainerBytes : 0,
    WorkspaceFirstLba = usbWorkspace ? ClusterOffset(workspaceCluster)/sectorSize : 0,
    StorageSourceSha256 = storageSources,
    BootVisualPolicy = independentBuild is { } visualBuild && visualBuild.TryGetProperty("BootVisualPolicy", out var visualPolicy) ? visualPolicy.GetString() : "Legacy"
};
File.WriteAllText(outputPath + ".json", JsonSerializer.Serialize(receipt, new JsonSerializerOptions { WriteIndented = true }));
Console.WriteLine(JsonSerializer.Serialize(receipt));

byte[] Header(int current, int backup, int entries)
{
    var header = new byte[sectorSize];
    "EFI PART"u8.CopyTo(header); U32(header, 8, 0x10000); U32(header, 12, 92);
    U64(header, 24, current); U64(header, 32, backup); U64(header, 40, 34); U64(header, 48, last - 33);
    diskId.TryWriteBytes(header.AsSpan(56));
    U64(header, 72, entries); U32(header, 80, 128); U32(header, 84, 128);
    U32(header, 88, Crc(table)); U32(header, 16, Crc(header.AsSpan(0, 92)));
    return header;
}
int Allocate(int count)
{
    if (count < 1 || nextCluster + (long)count > sectors - reserved - 2 * fatSectors + 2)
        throw new InvalidDataException("The kernels do not fit this image.");
    var start = nextCluster;
    for (var index = 0; index < count; index++)
    {
        U32(fat, nextCluster * 4, index == count - 1 ? 0x0FFFFFFF : nextCluster + 1);
        nextCluster++;
    }
    return start;
}
int ClusterOffset(int cluster) => partition + (reserved + 2 * fatSectors + cluster - 2) * sectorSize;
int Store(byte[] content)
{
    var cluster = Allocate(Math.Max(1, (content.Length + sectorSize - 1) / sectorSize));
    content.CopyTo(image, ClusterOffset(cluster));
    return cluster;
}
void Directory(int cluster, params byte[][] entries)
{
    var offset = ClusterOffset(cluster);
    foreach (var entry in entries)
    {
        if (offset + entry.Length > ClusterOffset(cluster) + sectorSize)
            throw new InvalidDataException("Directory exceeds the fixed one-cluster limit.");
        entry.CopyTo(image, offset); offset += entry.Length;
    }
}
byte[] Dot(int cluster) => Entry(".", ".", true, cluster, 0);
byte[] DotDot(int cluster) => Entry("..", "..", true, cluster, 0);
byte[] Entry(string name, string shortName, bool directory, int cluster, int size)
{
    var shortEntry = new byte[32];
    Encoding.ASCII.GetBytes(shortName.PadRight(11)).CopyTo(shortEntry, 0);
    shortEntry[11] = directory ? (byte)0x10 : (byte)0x20;
    U16(shortEntry, 20, cluster >> 16); U16(shortEntry, 26, cluster & 0xFFFF); U32(shortEntry, 28, size);
    if (name is "." or ".." || name == shortName.Trim()) return shortEntry;
    byte checksum = 0;
    for (var index = 0; index < 11; index++) checksum = (byte)(((checksum & 1) << 7) + (checksum >> 1) + shortEntry[index]);
    var pieces = (name.Length + 12) / 13;
    var result = new byte[(pieces + 1) * 32];
    int[] positions = [1, 3, 5, 7, 9, 14, 16, 18, 20, 22, 24, 28, 30];
    for (var piece = pieces; piece > 0; piece--)
    {
        var offset = (pieces - piece) * 32;
        result[offset] = (byte)(piece | (piece == pieces ? 0x40 : 0));
        result[offset + 11] = 0x0F; result[offset + 13] = checksum;
        for (var index = 0; index < 13; index++)
        {
            var at = (piece - 1) * 13 + index;
            U16(result, offset + positions[index], at < name.Length ? name[at] : at == name.Length ? 0 : 0xFFFF);
        }
    }
    shortEntry.CopyTo(result, pieces * 32);
    return result;
}
static void U16(byte[] bytes, int at, int value) => BinaryPrimitives.WriteUInt16LittleEndian(bytes.AsSpan(at), checked((ushort)value));
static void U32(byte[] bytes, int at, long value) => BinaryPrimitives.WriteUInt32LittleEndian(bytes.AsSpan(at), checked((uint)value));
static void U64(byte[] bytes, int at, long value) => BinaryPrimitives.WriteUInt64LittleEndian(bytes.AsSpan(at), checked((ulong)value));
static uint Crc(ReadOnlySpan<byte> bytes)
{
    var crc = uint.MaxValue;
    foreach (var value in bytes)
    {
        crc ^= value;
        for (var bit = 0; bit < 8; bit++) crc = crc >> 1 ^ ((crc & 1) == 0 ? 0 : 0xEDB88320);
    }
    return ~crc;
}