using System;
using System.Collections.Generic;
using System.IO.Enumeration;
using System.IO;
using System.Linq;
using System.Reflection.PortableExecutable;
using System.Runtime.ConstrainedExecution;
using System.Text;
using System.Threading.Tasks;
using skyscraper5.Skyscraper.IO;
using System.Runtime.InteropServices;
using skyscraper5.Skyscraper;

namespace skyscraper5.src.Skyscraper.IO.PortableExecutables
{
	public class PeFile : IDisposable
	{
		const ushort PortableExecutable = 17744;
		const ushort ZbykowskiExecutable = 23117;
		private PeHeader header;
		private PeSection[] sections;
		private bool ownsStream;
		private bool disposed;
		private Stream stream;
		public PeOptionalHeader OptionalHeader { get; private set; }
	

		public PeFile(Stream str, bool ownsStream = true)
		{
			long peOffset = FindPeOffset(str);
			str.Position = peOffset;
			header.mMagic = str.ReadUInt32LE();
			header.mMachine = str.ReadUInt16LE();
			header.mNumberOfSections = str.ReadUInt16LE();
			header.mTimeDateStamp = str.ReadUInt32LE();
			header.mPointerToSymbolTable = str.ReadUInt32LE();
			header.mNumberOfSymbols = str.ReadUInt32LE();
			header.mSizeOfOptionalHeader = str.ReadUInt16LE();
			header.mCharacteristics = str.ReadUInt16LE();
			if (header.mSizeOfOptionalHeader != 0)
			{
				byte[] optionalHeaderBytes = str.ReadBytes(header.mSizeOfOptionalHeader);
				OptionalHeader = new PeOptionalHeader(optionalHeaderBytes);
			}

			sections = new PeSection[header.mNumberOfSections];

			for (int i = 0; i < header.mNumberOfSections; i++)
			{
				sections[i].sectionName = str.ReadBytes(8);
				sections[i].virtualSize = str.ReadUInt32LE();
				sections[i].virtualAddress = str.ReadUInt32LE();
				sections[i].sizeOfRawData = str.ReadUInt32LE();
				sections[i].pointerToRawData = str.ReadUInt32LE();
				str.Position += 16;
			}
			this.stream = str;
			this.ownsStream = ownsStream;
		}

		private long FindPeOffset(Stream str)
		{
			ushort magic = str.ReadUInt16LE();
			switch (magic)
			{
				case ZbykowskiExecutable:
					str.Position = 0x3C;
					uint peAddress = str.ReadUInt32LE();
					str.Position = peAddress;
					magic = str.ReadUInt16LE();
					if (magic == PortableExecutable)
					{
						return peAddress;
					}
					else
					{
						throw new PeException("Found a MZ header, but not a PE header.");
					}
				case PortableExecutable:
					return 0;
				default:
					throw new PeException("Failed to find a PE header.");
			}
		}

		public void Dispose()
		{
			if (disposed)
				throw new ObjectDisposedException(this.ToString());
			if (ownsStream)
			{
				stream.Close();
				stream.Dispose();
			}
			disposed = true;
		}

		public PeSectionStream GetSectionStream(string sectionName)
		{
			if (sectionName.Length > 8)
				throw new ArgumentException(String.Format("{0} too long.", nameof(sectionName)));
			if (disposed)
				throw new ObjectDisposedException(this.ToString());
			foreach(PeSection section in sections)
			{
				if (section.ToString().Equals(sectionName))
					return GetSectionStream(section);
			}
			throw new PeSectionNotFoundException(sectionName);
		}

		public PeSectionStream GetSectionStream(PeSection section)
		{
			if (!sections.Contains(section, new PeSectionEqualityComparer()))
				throw new ArgumentException(String.Format("{0} must be a {1} in this {2}", nameof(section), nameof(PeSection), nameof(PeFile)));

			return new PeSectionStream(this, section, stream);
		}
	}

	struct PeHeader
	{
		public UInt32 mMagic; // PE\0\0 or 0x00004550
		public UInt16 mMachine;
		public UInt16 mNumberOfSections;
		public UInt32 mTimeDateStamp;
		public UInt32 mPointerToSymbolTable;
		public UInt32 mNumberOfSymbols;
		public UInt16 mSizeOfOptionalHeader;
		public UInt16 mCharacteristics;
	};

	[StructLayout(LayoutKind.Sequential, Pack = 1)]
	public struct PeSection
	{
		public byte[] sectionName;
		public uint virtualSize, virtualAddress, sizeOfRawData;
		internal uint pointerToRawData;

		public override bool Equals(object obj)
		{
			return obj is PeSection section &&
				   EqualityComparer<byte[]>.Default.Equals(sectionName, section.sectionName) &&
				   virtualSize == section.virtualSize &&
				   virtualAddress == section.virtualAddress &&
				   sizeOfRawData == section.sizeOfRawData;
		}

		public override int GetHashCode()
		{
			return HashCode.Combine(sectionName, virtualSize, virtualAddress, sizeOfRawData);
		}

		public override string ToString()
		{
			return Encoding.ASCII.GetString(sectionName).Split('\0')[0];
		}


	}

	public class PeOptionalHeader : Validatable
	{
		internal PeOptionalHeader(byte[] buffer)
		{
			MemoryStream ms = new MemoryStream(buffer, false);
			ushort magic = ms.ReadUInt16LE();
			if (magic != 0x010b && magic != 0x020b)
			{
				Valid = false;
				return;
			}
			byte majorLinker = ms.ReadUInt8();
			byte minorLinker = ms.ReadUInt8();
			LinkerVersion = new Version(majorLinker, minorLinker);
			SizeOfCode = ms.ReadUInt32LE();
			SizeOfInitializedData = ms.ReadUInt32LE();
			SizeOfUninitializedData = ms.ReadUInt32LE();
			AddressOfEntryPoint = ms.ReadUInt32LE();
			BaseOfCode = ms.ReadUInt32LE();
			BaseOfData = ms.ReadUInt32LE();
			ImageBase = ms.ReadUInt32LE();
			SectionAlignment = ms.ReadUInt32LE();
			FileAlignment = ms.ReadUInt32LE();
			ushort majorOs = ms.ReadUInt16LE();
			ushort minorOs = ms.ReadUInt16LE();
			OperatingSystemVersion = new Version(majorOs, minorOs);
			ushort majorImage = ms.ReadUInt16LE();
			ushort minorImage = ms.ReadUInt16LE();
			ImageVersion = new Version(majorImage, minorImage);
			ushort majorSubsystem = ms.ReadUInt16LE();
			ushort minorSubsystem = ms.ReadUInt16LE();
			SubsystemVersion = new Version(majorSubsystem, minorSubsystem);
			Win32VersionValue = ms.ReadUInt32LE();
			SizeOfImage = ms.ReadUInt32LE();
			SizeOfHeaders = ms.ReadUInt32LE();
			Checksum = ms.ReadUInt32LE();
			this.Subsystem = ms.ReadUInt16LE();
			this.DllCharacteristics = (DllCharacteristics)ms.ReadUInt16LE();
			SizeOfStackReserve = ms.ReadUInt32LE();
			SizeOfStackCommit = ms.ReadUInt32LE();
			SizeOfHeapReserve = ms.ReadUInt32LE();
			SizeOfHeapCommit = ms.ReadUInt32LE();
			LoaderFlags = ms.ReadUInt32LE();
			uint numRva = ms.ReadUInt32LE();
			DataDirectories = new DataDirectory[numRva];
			for (uint i = 0; i < numRva; i++)
			{
				uint rva = ms.ReadUInt32LE();
				uint size = ms.ReadUInt32LE();
				DataDirectories[i] = new DataDirectory(rva, size);
			}
		}

		public Version LinkerVersion { get; }
		public uint SizeOfCode { get; }
		public uint SizeOfInitializedData { get; }
		public uint SizeOfUninitializedData { get; }
		public uint AddressOfEntryPoint { get; }
		public uint BaseOfCode { get; }
		public uint BaseOfData { get; }
		public uint ImageBase { get; }
		public uint SectionAlignment { get; }
		public uint FileAlignment { get; }
		public Version OperatingSystemVersion { get; }
		public Version ImageVersion { get; }
		public Version SubsystemVersion { get; }
		public uint Win32VersionValue { get; }
		public uint SizeOfImage { get; }
		public uint SizeOfHeaders { get; }
		public uint Checksum { get; }
		public ushort Subsystem { get; }
		public DllCharacteristics DllCharacteristics { get; }
		public uint SizeOfStackReserve { get; }
		public uint SizeOfStackCommit { get; }
		public uint SizeOfHeapReserve { get; }
		public uint SizeOfHeapCommit { get; }
		public uint LoaderFlags { get; }

		public DataDirectory[] DataDirectories { get; private set; }

		public class DataDirectory
		{
			public DataDirectory(uint relativeVirtualAddress, uint size)
			{
				RelativeVirtualAddress = relativeVirtualAddress;
				Size = size;
			}

			public uint RelativeVirtualAddress { get; }
			public uint Size { get; }
		}
	}

}
