using System.Buffers.Binary;
using System.Text;
using StellaOps.Scanner.CallGraph.Binary;
using StellaOps.Scanner.CallGraph.Binary.Analysis;
using StellaOps.Scanner.CallGraph.Binary.Disassembly;
using Xunit;

using StellaOps.TestKit;

// Avoid ambiguity by using fully qualified names where needed
using DisassemblyTextSection = StellaOps.Scanner.CallGraph.Binary.Disassembly.BinaryTextSection;
using DisassemblyArchitecture = StellaOps.Scanner.CallGraph.Binary.Disassembly.BinaryArchitecture;
using DisassemblyTextSectionReader = StellaOps.Scanner.CallGraph.Binary.Disassembly.BinaryTextSectionReader;

namespace StellaOps.Scanner.CallGraph.Tests;

public class BinaryTextSectionReaderTests
{
    [Trait("Category", TestCategories.Unit)]
        [Fact]
    public async Task ReadsElfTextSection()
    {
        var textBytes = new byte[] { 0x90, 0x90, 0xC3, 0x90 };
        var rodataBytes = Encoding.ASCII.GetBytes("libfoo.so\0");
        var data = BuildElf64Fixture(textBytes, rodataBytes);

        var path = WriteTempFile(data);
        try
        {
            var section = await DisassemblyTextSectionReader.TryReadAsync(path, BinaryFormat.Elf, TestContext.Current.CancellationToken);
            Assert.NotNull(section);
            Assert.Equal(".text", section!.SectionName);
            Assert.Equal(DisassemblyArchitecture.X64, section.Architecture);
            Assert.Equal(0x1000UL, section.VirtualAddress);
            Assert.Equal(textBytes, section.Bytes);
        }
        finally
        {
            File.Delete(path);
        }
    }

    [Trait("Category", TestCategories.Unit)]
        [Fact]
    public async Task ReadsPeTextSection()
    {
        var textBytes = new byte[] { 0x90, 0x90, 0xC3, 0x90 };
        var data = BuildPe64Fixture(textBytes);

        var path = WriteTempFile(data);
        try
        {
            var section = await DisassemblyTextSectionReader.TryReadAsync(path, BinaryFormat.Pe, TestContext.Current.CancellationToken);
            Assert.NotNull(section);
            Assert.Equal(".text", section!.SectionName);
            Assert.Equal(DisassemblyArchitecture.X64, section.Architecture);
            Assert.Equal(0x1000UL, section.VirtualAddress);
            Assert.Equal(textBytes, section.Bytes);
        }
        finally
        {
            File.Delete(path);
        }
    }

    [Trait("Category", TestCategories.Unit)]
        [Fact]
    public async Task ReadsMachOTextSection()
    {
        var textBytes = new byte[] { 0x1F, 0x20, 0x03, 0xD5 };
        var data = BuildMachO64Fixture(textBytes);

        var path = WriteTempFile(data);
        try
        {
            var section = await DisassemblyTextSectionReader.TryReadAsync(path, BinaryFormat.MachO, TestContext.Current.CancellationToken);
            Assert.NotNull(section);
            Assert.Equal("__text", section!.SectionName);
            Assert.Equal(DisassemblyArchitecture.Arm64, section.Architecture);
            Assert.Equal(0x1000UL, section.VirtualAddress);
            Assert.Equal(textBytes, section.Bytes);
        }
        finally
        {
            File.Delete(path);
        }
    }

    [Trait("Category", TestCategories.Unit)]
        [Fact]
    public async Task StringScannerExtractsLibraryCandidates()
    {
        var textBytes = new byte[] { 0x90, 0x90, 0xC3, 0x90 };
        var rodataBytes = Encoding.ASCII.GetBytes("libfoo.so\0libbar.so.1\0");
        var data = BuildElf64Fixture(textBytes, rodataBytes);

        var path = WriteTempFile(data);
        try
        {
            var scanner = new BinaryStringLiteralScanner();
            var candidates = await scanner.ExtractLibraryCandidatesAsync(path, BinaryFormat.Elf, TestContext.Current.CancellationToken);

            Assert.Contains("libfoo.so", candidates);
            Assert.Contains("libbar.so.1", candidates);
        }
        finally
        {
            File.Delete(path);
        }
    }

    private static string WriteTempFile(byte[] data)
    {
        var path = Path.GetTempFileName();
        File.WriteAllBytes(path, data);
        return path;
    }

    private static byte[] BuildElf64Fixture(byte[] textBytes, byte[] rodataBytes)
    {
        const int textOffset = 0x100;
        var rodataOffset = textOffset + textBytes.Length;
        var shstrOffset = rodataOffset + rodataBytes.Length;
        const int sectionHeaderOffset = 0x200;
        const ushort sectionHeaderSize = 64;
        const ushort sectionCount = 4;

        var shstrtab = Encoding.ASCII.GetBytes("\0.text\0.rodata\0.shstrtab\0");

        var fileSize = sectionHeaderOffset + sectionHeaderSize * sectionCount;
        if (fileSize < shstrOffset + shstrtab.Length)
        {
            fileSize = shstrOffset + shstrtab.Length;
        }

        var data = new byte[fileSize];

        data[0] = 0x7F;
        data[1] = (byte)'E';
        data[2] = (byte)'L';
        data[3] = (byte)'F';
        data[4] = 2; // 64-bit
        data[5] = 1; // little endian

        WriteUInt16(data, 16, 2); // e_type
        WriteUInt16(data, 18, 62); // e_machine x86_64
        WriteInt64(data, 40, sectionHeaderOffset); // e_shoff
        WriteUInt16(data, 58, sectionHeaderSize); // e_shentsize
        WriteUInt16(data, 60, sectionCount); // e_shnum
        WriteUInt16(data, 62, 3); // e_shstrndx

        Array.Copy(textBytes, 0, data, textOffset, textBytes.Length);
        Array.Copy(rodataBytes, 0, data, rodataOffset, rodataBytes.Length);
        Array.Copy(shstrtab, 0, data, shstrOffset, shstrtab.Length);

        WriteElfSectionHeader(data, sectionHeaderOffset + sectionHeaderSize * 1, 1, textOffset, textBytes.Length, 0x1000);
        WriteElfSectionHeader(data, sectionHeaderOffset + sectionHeaderSize * 2, 7, rodataOffset, rodataBytes.Length, 0x2000);
        WriteElfSectionHeader(data, sectionHeaderOffset + sectionHeaderSize * 3, 15, shstrOffset, shstrtab.Length, 0);

        return data;
    }

    private static void WriteElfSectionHeader(
        byte[] data,
        int offset,
        uint nameOffset,
        int sectionOffset,
        int sectionSize,
        ulong address)
    {
        WriteUInt32(data, offset + 0, nameOffset);
        WriteUInt32(data, offset + 4, 1); // SHT_PROGBITS
        WriteUInt64(data, offset + 8, 0); // flags
        WriteUInt64(data, offset + 16, address);
        WriteUInt64(data, offset + 24, (ulong)sectionOffset);
        WriteUInt64(data, offset + 32, (ulong)sectionSize);
    }

    private static byte[] BuildPe64Fixture(byte[] textBytes)
    {
        const int peOffset = 0x80;
        const int optionalHeaderSize = 0xF0;
        const int sectionHeaderStart = peOffset + 4 + 20 + optionalHeaderSize;
        const int textOffset = 0x200;

        var fileSize = textOffset + textBytes.Length;
        if (fileSize < sectionHeaderStart + 40)
        {
            fileSize = sectionHeaderStart + 40;
        }

        var data = new byte[fileSize];
        WriteInt32(data, 0x3C, peOffset);
        WriteUInt32(data, peOffset, 0x00004550); // PE\0\0
        WriteUInt16(data, peOffset + 4, 0x8664); // machine
        WriteUInt16(data, peOffset + 6, 1); // sections
        WriteUInt16(data, peOffset + 20, optionalHeaderSize);

        WriteUInt16(data, peOffset + 24, 0x20b); // optional header magic

        WriteAscii(data, sectionHeaderStart + 0, ".text", 8);
        WriteUInt32(data, sectionHeaderStart + 8, (uint)textBytes.Length);
        WriteUInt32(data, sectionHeaderStart + 12, 0x1000);
        WriteUInt32(data, sectionHeaderStart + 16, (uint)textBytes.Length);
        WriteUInt32(data, sectionHeaderStart + 20, (uint)textOffset);

        Array.Copy(textBytes, 0, data, textOffset, textBytes.Length);
        return data;
    }

    private static byte[] BuildMachO64Fixture(byte[] textBytes)
    {
        const uint magic = 0xFEEDFACF;
        const int headerSize = 32;
        const int cmdSize = 152;
        const int commandStart = headerSize;
        const int textOffset = 0x200;

        var fileSize = textOffset + textBytes.Length;
        if (fileSize < commandStart + cmdSize)
        {
            fileSize = commandStart + cmdSize;
        }

        var data = new byte[fileSize];
        WriteUInt32(data, 0, magic);
        WriteInt32(data, 4, unchecked((int)0x0100000C)); // CPU_TYPE_ARM64
        WriteInt32(data, 8, 0);
        WriteUInt32(data, 12, 2); // filetype
        WriteUInt32(data, 16, 1); // ncmds
        WriteUInt32(data, 20, (uint)cmdSize);
        WriteUInt32(data, 24, 0); // flags
        WriteUInt32(data, 28, 0); // reserved

        WriteUInt32(data, commandStart + 0, 0x19); // LC_SEGMENT_64
        WriteUInt32(data, commandStart + 4, (uint)cmdSize);
        WriteAscii(data, commandStart + 8, "__TEXT", 16);
        WriteUInt64(data, commandStart + 24, 0x1000);
        WriteUInt64(data, commandStart + 32, 0x1000);
        WriteUInt64(data, commandStart + 40, (ulong)textOffset);
        WriteUInt64(data, commandStart + 48, (ulong)textBytes.Length);
        WriteInt32(data, commandStart + 56, 7);
        WriteInt32(data, commandStart + 60, 5);
        WriteUInt32(data, commandStart + 64, 1); // nsects
        WriteUInt32(data, commandStart + 68, 0); // flags

        var sectionStart = commandStart + 72;
        WriteAscii(data, sectionStart + 0, "__text", 16);
        WriteAscii(data, sectionStart + 16, "__TEXT", 16);
        WriteUInt64(data, sectionStart + 32, 0x1000);
        WriteUInt64(data, sectionStart + 40, (ulong)textBytes.Length);
        WriteUInt32(data, sectionStart + 48, (uint)textOffset);

        Array.Copy(textBytes, 0, data, textOffset, textBytes.Length);
        return data;
    }

    private static void WriteAscii(byte[] buffer, int offset, string value, int length)
    {
        var bytes = Encoding.ASCII.GetBytes(value);
        Array.Copy(bytes, 0, buffer, offset, Math.Min(length, bytes.Length));
    }

    private static void WriteUInt16(byte[] buffer, int offset, ushort value)
        => BinaryPrimitives.WriteUInt16LittleEndian(buffer.AsSpan(offset, 2), value);

    private static void WriteUInt32(byte[] buffer, int offset, uint value)
        => BinaryPrimitives.WriteUInt32LittleEndian(buffer.AsSpan(offset, 4), value);

    private static void WriteInt32(byte[] buffer, int offset, int value)
        => BinaryPrimitives.WriteInt32LittleEndian(buffer.AsSpan(offset, 4), value);

    private static void WriteUInt64(byte[] buffer, int offset, ulong value)
        => BinaryPrimitives.WriteUInt64LittleEndian(buffer.AsSpan(offset, 8), value);

    private static void WriteInt64(byte[] buffer, int offset, long value)
        => BinaryPrimitives.WriteInt64LittleEndian(buffer.AsSpan(offset, 8), value);
}