#include <ntddk.h>

#define WINAPI __stdcall

#define MAX_PCI_BUS 		255
#define MAX_PCI_SLOT 		32
#define SVGA_INDEX_PORT		0x0
#define SVGA_VALUE_PORT		0x1
#define PCI_ADDRESS_IO_MASK 	0xFFFFFFFC

typedef struct BARs{
	unsigned long BAR;
	unsigned long szBAR;
}BARs, *pBARs;

BARs bar_regs[5];

typedef unsigned long DWORD;
typedef DWORD *PDWORD;
typedef unsigned short WORD;
typedef WORD *PWORD;
typedef unsigned char BYTE;
typedef BYTE *PBYTE;

enum {
	SVGA_REG_ID = 0,
	SVGA_REG_ENABLE = 1,
	SVGA_REG_WIDTH = 2,
	SVGA_REG_HEIGHT = 3,
	SVGA_REG_MAX_WIDTH = 4,
	SVGA_REG_MAX_HEIGHT = 5,
	SVGA_REG_DEPTH = 6,
	SVGA_REG_BITS_PER_PIXEL = 7,     /* Current bpp in the guest */
	SVGA_REG_PSEUDOCOLOR = 8,
	SVGA_REG_RED_MASK = 9,
	SVGA_REG_GREEN_MASK = 10,
	SVGA_REG_BLUE_MASK = 11,
	SVGA_REG_BYTES_PER_LINE = 12,
	SVGA_REG_FB_START = 13,
	SVGA_REG_FB_OFFSET = 14,
	SVGA_REG_VRAM_SIZE = 15,
	SVGA_REG_FB_SIZE = 16,
   /* ID 0 implementation only had the above registers, then the palette */
	SVGA_REG_CAPABILITIES = 17,
	SVGA_REG_MEM_START = 18,	   /* Memory for command FIFO and bitmaps */
	SVGA_REG_MEM_SIZE = 19,
	SVGA_REG_CONFIG_DONE = 20,      /* Set when memory area configured */
	SVGA_REG_SYNC = 21,             /* See "FIFO Synchronization Registers" */
	SVGA_REG_BUSY = 22,             /* See "FIFO Synchronization Registers" */
	SVGA_REG_GUEST_ID = 23,	   /* Set guest OS identifier */
	SVGA_REG_CURSOR_ID = 24,	   /* ID of cursor */
	SVGA_REG_CURSOR_X = 25,	   /* Set cursor X position */
	SVGA_REG_CURSOR_Y = 26,	   /* Set cursor Y position */
	SVGA_REG_CURSOR_ON = 27,	   /* Turn cursor on/off */
	SVGA_REG_HOST_BITS_PER_PIXEL = 28, /* Current bpp in the host */
	SVGA_REG_SCRATCH_SIZE = 29,     /* Number of scratch registers */
	SVGA_REG_MEM_REGS = 30,         /* Number of FIFO registers */
	SVGA_REG_NUM_DISPLAYS = 31,     /* Number of guest displays */
	SVGA_REG_PITCHLOCK = 32,        /* Fixed pitch for all modes */
	SVGA_REG_IRQMASK = 33,          /* Interrupt mask */
	SVGA_REG_NUM_GUEST_DISPLAYS = 34,/* Number of guest displays in X/Y direction */
	SVGA_REG_DISPLAY_ID = 35,        /* The display ID for the following display attributes */
	SVGA_REG_DISPLAY_IS_PRIMARY = 36,/* Whether this is a primary display */
	SVGA_REG_DISPLAY_POSITION_X = 37,/* The display position x */
	SVGA_REG_DISPLAY_POSITION_Y = 38,/* The display position y */
	SVGA_REG_DISPLAY_WIDTH = 39,     /* The display's width */
	SVGA_REG_DISPLAY_HEIGHT = 40,    /* The display's height */
	SVGA_REG_TOP = 41,               /* Must be 1 more than the last register */
	SVGA_PALETTE_BASE = 1024	    /* Base of SVGA color map */
};

enum {
	SVGA_FIFO_MIN = 0,
	SVGA_FIFO_MAX,       /* The distance from MIN to MAX must be at least 10K */
	SVGA_FIFO_NEXT_CMD,
	SVGA_FIFO_STOP
};

static KSPIN_LOCK PciLock;
static ULONG BusType = 0;
static DWORD IOports = 0;

VOID OnUnload(IN PDRIVER_OBJECT pDriverObject);
NTSTATUS OnStubDispatch(IN PDEVICE_OBJECT DeviceObject, IN PIRP Irp);
VOID InitPciBus(VOID);
VOID GetAddress(PDWORD framebuffer, PDWORD VgaFifo, PULONG framebufferSize, PULONG VgaFifoSize);
ULONG GetBusConfigType(VOID);
ULONG PciConfigUlong(UCHAR bus, UCHAR slot, UCHAR func, ULONG offset);
ULONG SizeBar(UCHAR bus, UCHAR slot, UCHAR func, ULONG offset);
ULONG ConvertBARsz(ULONG size);
VOID vmwareWriteWordToFIFO(PDWORD vmwareFIFO, DWORD value);

NTSTATUS DriverEntry(IN PDRIVER_OBJECT pDriverObject, IN PUNICODE_STRING RegistryPath) {
	NTSTATUS    	Status = 0;
	KIRQL 		oldIrql;
	ULONG		framebufferSize=0, VgaFifoSize=0, i;
	DWORD 		framebuffer = 0, VgaFifo = 0;
		
	pDriverObject->DriverUnload = OnUnload;

	for(i=0; i<IRP_MJ_MAXIMUM_FUNCTION; i++) pDriverObject->MajorFunction[i] = OnStubDispatch;

	GetAddress(&framebuffer, &VgaFifo, &framebufferSize, &VgaFifoSize);
	if(framebuffer == 0 || VgaFifo == 0 || IOports == 0) return -1;
	
	DbgPrint("\nMmMapIoSpace() Alloc...\n");
	DbgPrint("Frame buffer: [0x%x - 0x%x]\n", framebuffer, framebuffer+framebufferSize);

	DbgPrint("VgaFifo: [0x%x - 0x%x]\n", VgaFifo, VgaFifo+VgaFifoSize);
	DbgPrint("VgaFifoCode: [0x%x - 0x%x]\n", VgaFifo+*(PDWORD)(VgaFifo), VgaFifo+*((PDWORD)VgaFifo+1));

	//Bloquage
	KeAcquireSpinLock(&PciLock, &oldIrql);
	WRITE_PORT_ULONG((PULONG)(IOports+SVGA_INDEX_PORT), SVGA_REG_CONFIG_DONE);
	WRITE_PORT_ULONG((PULONG)(IOports+SVGA_VALUE_PORT), 0);
	
	vmwareWriteWordToFIFO((PDWORD)VgaFifo, 0x03); // SURF_COPY
	vmwareWriteWordToFIFO((PDWORD)VgaFifo, 0x80000000 - 0x4); // Source X
	vmwareWriteWordToFIFO((PDWORD)VgaFifo, 0); // Source Y 
	vmwareWriteWordToFIFO((PDWORD)VgaFifo, 0); // Dest X
	vmwareWriteWordToFIFO((PDWORD)VgaFifo, 0); // Dest Y
	vmwareWriteWordToFIFO((PDWORD)VgaFifo, 0x4); // Width
	vmwareWriteWordToFIFO((PDWORD)VgaFifo, 1); // Height

	WRITE_PORT_ULONG((PULONG)(IOports+SVGA_INDEX_PORT), SVGA_REG_CONFIG_DONE);
	WRITE_PORT_ULONG((PULONG)(IOports+SVGA_VALUE_PORT), 1);

	KeReleaseSpinLock(&PciLock, oldIrql);
	
	// Dump qui devrai contenir la memoire de l'host...
	DbgPrint("Dump:\n%x %x %x %x\n", *((PDWORD)framebuffer), *((PDWORD)framebuffer+1), *((PDWORD)framebuffer+2), *((PDWORD)framebuffer+3));

	return Status;
}

VOID OnUnload(IN PDRIVER_OBJECT pDriverObject) {
    DbgPrint("Unload...");
}

NTSTATUS OnStubDispatch(IN PDEVICE_OBJECT DeviceObject, IN PIRP Irp) {
    Irp->IoStatus.Status = STATUS_SUCCESS;
    IoCompleteRequest(Irp, IO_NO_INCREMENT);
    return STATUS_SUCCESS;
}

VOID InitPciBus(VOID) {
	KeInitializeSpinLock(&PciLock);
	BusType = GetBusConfigType();
	//DbgPrint("Bus type: %i\n", BusType);
}

VOID GetAddress(PDWORD framebuffer, PDWORD VgaFifo, PULONG framebufferSize, PULONG VgaFifoSize) {
	PHYSICAL_ADDRESS 	pAdrss;
	ULONG 				out, size;
	int 				bus, slot, i;

	InitPciBus();
	for(bus = 0; bus < MAX_PCI_BUS; bus++)
	{
		for(slot=0; slot < MAX_PCI_SLOT; slot++)
		{		
			out = PciConfigUlong((UCHAR)bus, (UCHAR)slot, 0, 0);
			if((0xFFFF & out) != 0xFFFF && (0xFFFF & out) > 0) {
				out = PciConfigUlong((UCHAR)bus, (UCHAR)slot, 0, 0);
				if((0xFFFF0000 & out)>>16 == 0x405 && (0xFFFF & out) == 0x15ad) {
					DbgPrint("Device ID: %x\tVendor ID: %x\n", (0xFFFF0000 & out)>>16, (0xFFFF & out));
					for(i = 0; i<3; i++) {
						out = PciConfigUlong((UCHAR)bus, (UCHAR)slot, 0, 0x10+i*4);
						if(out > 0) {
							size = SizeBar((UCHAR)bus, (UCHAR)slot, 0, 0x10+i*4);
							if(i == 0) {
								// DbgPrint("I/O Ports: [0x%x - 0x%x]\n", out, out+size);
								IOports = out & PCI_ADDRESS_IO_MASK;;
							} else {
								RtlZeroMemory(&pAdrss,sizeof(PHYSICAL_ADDRESS));
								pAdrss.LowPart = (DWORD)out;
								if(i == 1) {
									// DbgPrint("\tFrame buffer: [0x%x - 0x%x]\n", out, out+size);
									*framebufferSize = size;
									*framebuffer = (DWORD)MmMapIoSpace(pAdrss, size, FALSE);
								} else {
									// DbgPrint("\tSvga FIFO: [0x%x - 0x%x]\n", out, out+size);
									*VgaFifoSize = size;
									*VgaFifo = (DWORD)MmMapIoSpace(pAdrss, size, FALSE);
								}
							}									
						}
					}
				}
			} 
		}
	}
}

ULONG GetBusConfigType(VOID) {
	ULONG Value;
	KIRQL oldIrql;
	
	KeAcquireSpinLock(&PciLock, &oldIrql);
	
	WRITE_PORT_UCHAR((PUCHAR)0xcfb, 0x01);
	Value = READ_PORT_ULONG((PULONG)0xCF8);
	WRITE_PORT_ULONG((PULONG)0xcf8, 0x80000000);
	
	if (READ_PORT_ULONG((PULONG)0xcf8) == 0x80000000) {
		WRITE_PORT_ULONG((PULONG)0xcf8, Value);
		KeReleaseSpinLock(&PciLock, oldIrql);

		return 1;
	}
	
	WRITE_PORT_ULONG((PULONG)0xcf8, Value);
	WRITE_PORT_UCHAR((PUCHAR)0xcfb, 0x00);
	WRITE_PORT_UCHAR((PUCHAR)0xcf8, 0x00);
	WRITE_PORT_UCHAR((PUCHAR)0xcfa, 0x00);
	
	if (READ_PORT_UCHAR((PUCHAR)0xcf8) == 0x00 && READ_PORT_UCHAR((PUCHAR)0xcfB) == 0x00) {
		KeReleaseSpinLock(&PciLock, oldIrql);
		return 2;
	}
	KeReleaseSpinLock(&PciLock, oldIrql);

	return 0;
}

static ULONG PciConfigUlong(UCHAR bus, UCHAR slot, UCHAR func, ULONG offset)
{ 
	KIRQL oldIrql;
	ULONG tmp;
	ULONG address;
	ULONG lbus = (ULONG)bus;
	ULONG lslot = (ULONG)slot;
	ULONG lfunc = (ULONG)func;

	address = ((lbus << 16) | (lslot << 11) |(lfunc << 8 ) | (offset & 0xfc) | ((UINT32)0x80000000));

	if(BusType == 1)
	{
		KeAcquireSpinLock(&PciLock, &oldIrql);
		
		WRITE_PORT_ULONG((PULONG)0xcf8, address);
		tmp = READ_PORT_ULONG((PULONG)0xcfc);
		
		KeReleaseSpinLock(&PciLock, oldIrql);
		return tmp;
	}
	return STATUS_UNSUCCESSFUL;
}

static ULONG SizeBar(UCHAR bus, UCHAR slot, UCHAR func, ULONG offset)
{ 
	KIRQL oldIrql;
	ULONG tmp, address, bar;
	ULONG lbus = (ULONG)bus;
	ULONG lslot = (ULONG)slot;
	ULONG lfunc = (ULONG)func;
	address = ((lbus << 16) | (lslot << 11) |(lfunc << 8 ) | (offset & 0xfc) | ((UINT32)0x80000000));

	if(BusType == 1)
	{
		KeAcquireSpinLock(&PciLock, &oldIrql);

		WRITE_PORT_ULONG((PULONG)0xcf8, address);
		bar = READ_PORT_ULONG((PULONG)0xcfc);
		WRITE_PORT_ULONG((PULONG)0xcf8, address);
		WRITE_PORT_ULONG((PULONG)0xcfc, 0xffffffff);
		tmp = ConvertBARsz(READ_PORT_ULONG((PULONG)0xcfc));
		WRITE_PORT_ULONG((PULONG)0xcf8, address);
		WRITE_PORT_ULONG((PULONG)0xcfc, bar);
		
		KeReleaseSpinLock(&PciLock, oldIrql);

		return tmp;
	}
	return 0;
}

ULONG ConvertBARsz(ULONG size)
{
	ULONG realsize;

	realsize = (size&0xfffffff0);
	realsize = (1 + ~realsize);

	if(realsize & 1)
	{
		realsize = (realsize&0xffff);
	} else {

		realsize = (realsize&0xffffffff);
	}
	return realsize;
}

VOID vmwareWriteWordToFIFO(PDWORD vmwareFIFO, DWORD value)
{
	/* Need to sync? */
	if ((vmwareFIFO[SVGA_FIFO_NEXT_CMD] + sizeof(DWORD) == vmwareFIFO[SVGA_FIFO_STOP])
	|| (vmwareFIFO[SVGA_FIFO_NEXT_CMD] == vmwareFIFO[SVGA_FIFO_MAX] - sizeof(DWORD) &&
	vmwareFIFO[SVGA_FIFO_STOP] == vmwareFIFO[SVGA_FIFO_MIN])) {
		DbgPrint("Syncing because of full fifo\n");

		WRITE_PORT_ULONG((PULONG)(IOports+SVGA_INDEX_PORT), SVGA_REG_SYNC);
		WRITE_PORT_ULONG((PULONG)(IOports+SVGA_VALUE_PORT), 1);

		while (1) {
			WRITE_PORT_ULONG((PULONG)(IOports+SVGA_INDEX_PORT), SVGA_REG_BUSY);
			if(READ_PORT_ULONG((PULONG)(IOports+SVGA_VALUE_PORT)) == 0) break;
		}
	}

	vmwareFIFO[vmwareFIFO[SVGA_FIFO_NEXT_CMD] / sizeof(DWORD)] = value;
	if(vmwareFIFO[SVGA_FIFO_NEXT_CMD] == vmwareFIFO[SVGA_FIFO_MAX] - sizeof(DWORD)) {
		vmwareFIFO[SVGA_FIFO_NEXT_CMD] = vmwareFIFO[SVGA_FIFO_MIN];
	} else {
		vmwareFIFO[SVGA_FIFO_NEXT_CMD] += sizeof(DWORD);
	}
}