;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;;                                                                 ;;
;; Copyright (C) KolibriOS team 2004-2024. All rights reserved.    ;;
;; Distributed under terms of the GNU General Public License       ;;
;;                                                                 ;;
;;          GNU GENERAL PUBLIC LICENSE                             ;;
;;             Version 2, June 1991                                ;;
;;                                                                 ;;
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

format PE DLL native
entry START

API_VERSION     equ 0  ;debug
SRV_GETVERSION  equ 0
__DEBUG__ = 1
__DEBUG_LEVEL__ = 1
DRIVER_VERSION = 1
DBG_INFO = 1
NULLPTR = 0

section ".flat" code readable writable executable
include "../proc32.inc"
include "../struct.inc"
include "../macros.inc"
include "../fdo.inc"
include "../pci.inc"
include "../peimport.inc"
include "nvme.inc"
include "macros.inc"
include "lib.asm"

proc START c, reason:dword

      	cmp     [reason], DRV_ENTRY
      	jne     .err

.entry:
		DEBUGF  DBG_INFO, "Detecting NVMe hardware...\n"
		call    detect_nvme
		test 	eax, eax
		jz      .err
		mov 	eax, dword [p_nvme_devices]
		test 	eax, eax
		jz 		.err
		xor 	ecx, ecx

.loop:
		mov 	ebx, dword [p_nvme_devices]
		stdcall device_is_compat, ebx
		test 	eax, eax
		jz 		@f
		stdcall nvme_init, ebx
		test 	eax, eax
		jz 		.err

;@@:
		;inc 	ecx
		;cmp 	ecx, dword [pcidevs_len]
		;jne 	.loop
		invoke  RegService, my_service, service_proc
		ret

.err:
		call    nvme_cleanup
      	xor     eax, eax
      	ret

endp

proc service_proc stdcall, ioctl:dword

      	mov     ebx, [ioctl]
      	mov     eax, [ebx+IOCTL.io_code]
      	cmp     eax, SRV_GETVERSION
      	jne     @F

      	mov     eax, [ebx+IOCTL.output]
      	cmp     [ebx+IOCTL.out_size], 4
      	jne     @F
      	mov     dword [eax], API_VERSION
      	xor     eax, eax
      	ret

@@:
        or      eax, -1
        ret

endp

proc set_cdw0 stdcall, pci:dword, y:dword, opcode:byte

		stdcall get_new_cid, [pci], [y]
		shl 	eax, 16
		or 		al, [opcode]
		ret

endp

; See pages 161-205 of the NVMe 1.4 specification for reference
proc nvme_identify stdcall, pci:dword, nsid:dword, dptr:dword, cns:byte

		sub 	esp, sizeof.SQ_ENTRY
		; It's important to check if CNS is a valid value here. In revision 1.0
		; CNS is a 1 bit field and a two bit field in revision 1.1, using invalid
		; values results in undefined behavior (see page 162 of NVMe 1.4 spec)
		if __DEBUG__
			push 	esi
			mov 	esi, [pci]
			mov 	esi, dword [esi + pcidev.io_addr]
			mov 	eax, dword [esi + NVME_MMIO.VS]
			cmp 	eax, VS110
			jne 	@f
			cmp 	[cns], 11b
			jle		.ok
			DEBUGF  DBG_INFO, "(NVMe) FATAL ERROR: INVALID CNS VALUE ON v1.1.0 CONTROLLERS\n"
			jmp 	.err
		
			@@:
			cmp 	eax, VS100
			jne 	.ok
			cmp 	[cns], 1b
			jle 	.ok
			DEBUGF  DBG_INFO, "(NVMe) FATAL ERROR: INVALID CNS VALUE ON v1.0.0 CONTROLLERS\n"
			jmp 	.err

			.err:
			jmp 	@b
			
			.ok:
			pop 	esi
			
		end if

		stdcall memset, esp, 0, sizeof.SQ_ENTRY

		mov     eax, [nsid]
		mov     dword [esp + SQ_ENTRY.nsid], eax
		mov     eax, [dptr]
		mov     dword [esp + SQ_ENTRY.dptr], eax
		stdcall set_cdw0, [pci], ADMIN_QUEUE, ADM_CMD_IDENTIFY
		mov 	dword [esp + SQ_ENTRY.cdw0], eax
		mov     al, [cns]
		mov     byte [esp + SQ_ENTRY.cdw10], al
		stdcall sqytdbl_write, [pci], ADMIN_QUEUE, esp

		add 	esp, sizeof.SQ_ENTRY
		ret

endp

; See pages 348-349 of the NVMe 1.4 specification for information on creating namespaces
proc create_namespace stdcall, pci:dword, cid:word
		
		push 	esi
		invoke  AllocPage
		test 	eax, eax
		jz 		.fail
		invoke  GetPhysAddr
		stdcall nvme_identify, [pci], 0xffffffff, eax, CNS_IDNS

.fail:
		pop 	esi
		ret

endp

; See page 248 of the NVMe 1.4 specification for reference
proc determine_active_nsids stdcall, pci:dword

		push 	esi edi ebx
		xor 	esi, esi
		xor 	edi, edi
		invoke  KernelAlloc, 0x1000
		test 	eax, eax
		jz 		.fail
		mov 	esi, eax
		invoke  GetPhysAddr
		stdcall nvme_identify, [pci], 0, eax, CNS_IDCS
		mov 	eax, dword [esi + IDENTC.nn]
		test 	eax, eax
		jz 		.fail
		invoke  KernelAlloc, 0x1000
		test 	eax, eax
		jz 		.fail
		mov 	edi, eax
		invoke  GetPhysAddr
		mov 	ebx, eax
		xor 	ecx, ecx

.loop:
		cmp 	ecx, dword [esi + IDENTC.nn]
		jg 		.success
		stdcall nvme_identify, [pci], ecx, ebx, CNS_IDNS
		push 	ecx
		xor 	ecx, ecx

.inner_loop:
		mov 	eax, dword [edi + ecx * 4]
		test 	eax, eax
		jnz 	.is_active_nsid
		inc 	ecx
		cmp 	ecx, 0x1000 / 4
		jne 	.inner_loop
		pop 	ecx
		inc 	ecx
		jmp 	.loop

.is_active_nsid:
		pop 	ecx ; pop active NSID value into ecx
		mov 	eax, [pci]
		mov 	eax, dword [eax + pcidev.nsids]
		mov 	dword [eax + ecx * 4], ecx
		inc 	ecx
		jmp 	.loop

.fail:
		invoke  KernelFree, esi ; assumes KernelFree can check if its nullptr (hopefully)
		invoke  KernelFree, edi
		pop 	ebx edi esi
		xor 	eax, eax
		ret

.success:
		invoke  KernelFree, esi
		invoke  KernelFree, edi
		pop 	ebx edi esi
		xor 	eax, eax
		inc 	eax
		ret
		
endp

; See page 101 of the NVMe 1.4 specification for reference
proc create_io_completion_queue stdcall, pci:dword, prp1:dword, qid:dword, ien:byte
	
		sub 	esp, sizeof.SQ_ENTRY
		stdcall memset, esp, 0, sizeof.SQ_ENTRY
		stdcall set_cdw0, [pci], ADMIN_QUEUE, ADM_CMD_CRE_IO_COMPLETION_QUEUE
		mov 	dword [esp + SQ_ENTRY.cdw0], eax
		mov 	eax, [prp1]
		mov 	dword [esp + SQ_ENTRY.dptr], eax
		mov 	eax, sizeof.CQ_ENTRY shl 16 ; CDW10.QSIZE
		or 		eax, [qid] ; CDW10.QID
		mov 	dword [esp + SQ_ENTRY.cdw10], eax
		movzx 	eax, [ien] ; CDW11.IEN 
		or 		eax, 0x1 ; CDW11.PC
		; Don't set CDW11.IV since we're not using MSI-X or MSI vector
		mov 	dword [esp + SQ_ENTRY.cdw11], eax
		stdcall sqytdbl_write, [pci], ADMIN_QUEUE, esp
		add 	esp, sizeof.SQ_ENTRY
		ret

endp

; See page 103-104 of the NVMe 1.4 specification for reference
proc create_io_submission_queue stdcall, pci:dword, prp1:dword, qid:dword, cqid:word

		sub 	esp, sizeof.SQ_ENTRY
		stdcall memset, esp, 0, sizeof.SQ_ENTRY
		stdcall set_cdw0, [pci], ADMIN_QUEUE, ADM_CMD_CRE_IO_SUBMISSION_QUEUE
		mov 	dword [esp + SQ_ENTRY.cdw0], eax
		mov 	eax, [prp1]
		mov 	dword [esp + SQ_ENTRY.dptr], eax
		mov 	eax, sizeof.SQ_ENTRY shl 16 ; CDW10.QSIZE
		or 		eax, [qid]
		mov 	dword [esp + SQ_ENTRY.cdw10], eax
		movzx 	eax, [cqid]
		shl 	eax, 16 ; CDW11.CQID
		or 		eax, 0x1 ; CDW11.PC (always set this to 1 as some devices may not support non-contiguous pages)
		; TODO: Set CDW10.QPRIO 
		mov 	dword [esp + SQ_ENTRY.cdw11], eax
		stdcall sqytdbl_write, [pci], ADMIN_QUEUE, esp
		add 	esp, sizeof.SQ_ENTRY
		ret

endp

; See page 95-96 of the NVMe 1.4 specification for reference
proc abort stdcall, pci:dword, cid:word, sqid:word

		sub 	esp, sizeof.SQ_ENTRY
		stdcall memset, esp, 0, sizeof.SQ_ENTRY
		stdcall set_cdw0, [pci], ADMIN_QUEUE, ADM_CMD_ABORT
		mov 	dword [esp + SQ_ENTRY.cdw0], eax
		movzx 	eax, [cid]
		shl 	eax, 16
		or 		eax, word [sqid]
		mov 	dword [esp + SQ_ENTRY.cdw10], eax
		stdcall sqytdbl_write, [pci], ADMIN_QUEUE, esp
		add 	esp, sizeof.SQ_ENTRY
		ret

endp


; See page 205 of the NVMe 1.4 specification for reference
proc set_features stdcall, pci:dword, dptr:dword, fid:byte, cdw11:dword
		
		sub 	esp, sizeof.SQ_ENTRY
		stdcall memset, esp, 0, sizeof.SQ_ENTRY
		stdcall set_cdw0, [pci], ADMIN_QUEUE, ADM_CMD_SET_FEATURES
		mov 	dword [esp + SQ_ENTRY.cdw0], eax
		mov 	eax, [dptr]
		mov 	dword [esp + SQ_ENTRY.dptr], eax
		movzx 	eax, [fid]
		;or 		eax, 1 shl 31 ; CDW10.SV
		mov 	dword [esp + SQ_ENTRY.cdw10], eax
		mov 	eax, [cdw11]
		mov 	dword [esp + SQ_ENTRY.cdw11], eax
		stdcall sqytdbl_write, [pci], ADMIN_QUEUE, esp
		add 	esp, sizeof.SQ_ENTRY
		ret

endp

; See page 105 of the NVMe 1.4 specification for reference
proc delete_io_completion_queue stdcall, pci:dword, qid:word

		sub 	esp, sizeof.SQ_ENTRY
		stdcall memset, esp, 0, sizeof.SQ_ENTRY
		stdcall set_cdw0, [pci], ADMIN_QUEUE, ADM_CMD_DEL_IO_COMPLETION_QUEUE
		mov 	dword [esp + SQ_ENTRY.cdw0], eax
		mov 	ax, [qid]
		mov 	word [esp + SQ_ENTRY.cdw10], ax
		stdcall sqytdbl_write, [pci], ADMIN_QUEUE, esp
		add 	esp, sizeof.SQ_ENTRY
		ret

endp

; See page 114-116 of the NVMe 1.4 specification for reference
proc get_features stdcall, pci:dword, dptr:dword, sel:byte, fid:byte
		
		sub 	esp, sizeof.SQ_ENTRY
		stdcall memset, esp, 0, sizeof.SQ_ENTRY
		stdcall set_cdw0, [pci], ADMIN_QUEUE, ADM_CMD_GET_FEATURES
		mov 	dword [esp + SQ_ENTRY.cdw0], eax
		movzx 	eax, [sel]
		and 	eax, 111b
		shl 	eax, 8 ; CDW10.SEL
		or 		eax, byte [fid] ; CDW10.FID
		mov 	dword [esp + SQ_ENTRY.cdw10], eax
		mov 	eax, [dptr]
		mov 	dword [esp + SQ_ENTRY.dptr], eax
		; TODO: Implement CDW14.UUID?
		stdcall sqytdbl_write, [pci], ADMIN_QUEUE, esp
		add 	esp, sizeof.SQ_ENTRY
		ret

endp

; See page 105-106 of the NVMe 1.4 specification for reference
proc delete_io_submission_queue stdcall, pci:dword, qid:word

		sub 	esp, sizeof.SQ_ENTRY
		stdcall memset, esp, 0, sizeof.SQ_ENTRY
		stdcall set_cdw0, [pci], ADMIN_QUEUE, ADM_CMD_DEL_IO_SUBMISSION_QUEUE
		mov 	dword [esp + SQ_ENTRY.cdw0], eax
		mov 	ax, [qid]
		mov 	word [esp + SQ_ENTRY.cdw10], ax
		stdcall sqytdbl_write, [pci], ADMIN_QUEUE, esp
		add 	esp, sizeof.SQ_ENTRY
		ret

endp

; See page 117-118 of the NVMe 1.4 specification for reference
; INCOMPLETE
proc get_log_page stdcall, pci:dword, dptr:dword, lid:byte

		sub 	esp, sizeof.SQ_ENTRY
		stdcall memset, esp, 0, sizeof.SQ_ENTRY
		stdcall set_cdw0, [pci], ADMIN_QUEUE, ADM_CMD_GET_LOG_PAGE
		mov 	dword [esp + SQ_ENTRY.cdw0], eax
		mov 	eax, [dptr]
		mov 	dword [esp + SQ_ENTRY.dptr], eax
		add 	esp, sizeof.SQ_ENTRY
		ret

endp

; See page 269-271 of the NVMe 1.4 specification for reference
proc nvme_write stdcall, pci:dword, qid:word, slba:qword, nlb:dword, dsm:byte
		
		; TODO: Use IDENTC.NOIOB to construct read/write commands that don't
		; cross the I/O boundary to achieve optimal performance
		; Also add DPTR/MPTR
		sub 	esp, sizeof.SQ_ENTRY
		stdcall memset, esp, 0 sizeof.SQ_ENTRY
		stdcall set_cdw0, [pci], [qid], NVM_CMD_WRITE
		mov 	dword [esp + SQ_ENTRY.cdw0], eax ; CDW0

		; Starting LBA (SLBA)
		mov 	eax, dword [slba]
		mov 	dword [esp + SQ_ENTRY.cdw10], eax
		mov 	eax, dword [slba + 4]
		mov 	dword [esp + SQ_ENTRY.cdw11], eax
		mov 	ax, [nlb]
		mov 	word [esp + SQ_ENTRY.cdw12], ax
		mov 	al, [dsm]
		mov 	byte [esp + SQ_ENTRY.cdw13], al
		stdcall sqytdbl_write, [pci], [qid], esp
		
		add 	esp, sizeof.SQ_ENTRY
		ret
		
endp


proc detect_nvme

      	invoke  GetPCIList
		mov 	edx, eax

.check_dev:
		mov 	ebx, dword [eax + PCIDEV.class]
		and		ebx, 0x00ffff00 ; retrieve class/subclass code only
		cmp 	ebx, 0x00010800 ; Mass Storage Controller - Non-Volatile Memory Controller
		je		.found_dev

.next_dev:
		mov		eax, dword [eax + PCIDEV.fd]
		cmp		eax, edx
		jne 	.check_dev
		jmp 	.exit_success

.found_dev:
		push    edx eax
		PDEBUGF DBG_INFO, "PCI(%u.%u.%u): Detected NVMe device...\n", byte [eax + PCIDEV.bus], byte [eax + PCIDEV.devfn]
		cmp 	dword [pcidevs_len], TOTAL_PCIDEVS
		jne 	@f
		pop 	eax edx
		jmp 	.exit_success

@@:
		inc 	dword [pcidevs_len]
		mov     ebx, dword [p_nvme_devices]
		test    ebx, ebx
		jnz		@f
		invoke  KernelAlloc, sizeof.pcidev
		test 	eax, eax
		jz 		.err_no_mem
		mov     dword [p_nvme_devices], eax 
		DEBUGF  DBG_INFO, "(NVMe) Allocated pcidev struct at 0x%x\n", [p_nvme_devices]

@@:
		mov 	ecx, dword [pcidevs_len]
		dec 	ecx
		pop		eax
		mov 	ebx, dword [p_nvme_devices]

		movzx 	edx, byte [eax + PCIDEV.bus]
		mov 	byte [ebx + pcidev.bus], dl
		movzx 	edx, byte [eax + PCIDEV.devfn]
		mov 	byte [ebx + pcidev.devfn], dl

		pop 	edx
		jmp		.next_dev

.err_no_mem:
		pop     eax edx
		xor 	eax, eax
		ret

.exit_success:
		xor 	eax, eax
		inc 	eax
		ret

endp

proc device_is_compat stdcall, pci:dword

		push  	 esi edx ecx
		mov 	 esi, [pci]
		invoke 	 PciRead8, dword [esi + pcidev.bus], dword [esi + pcidev.devfn], PCI_header00.interrupt_line
		mov 	 byte [esi + pcidev.iline], al
		invoke   PciRead32, dword [esi + pcidev.bus], dword [esi + pcidev.devfn], PCI_header00.base_addr_0
		and      eax, 0xfffffff0
		test     eax, eax
		jz       .failure
		mov 	 edx, eax
		push 	 edx

		invoke   MapIoMem, eax, sizeof.NVME_MMIO, PG_SW+PG_NOCACHE
		test     eax, eax
		jz       .failure
		;DEBUGF   DBG_INFO, "(NVMe) MMIO allocated at: 0x%x\n", eax
		mov 	 dword [esi + pcidev.io_addr], eax
		mov 	 eax, dword [eax + NVME_MMIO.CAP + 4]
		and 	 eax, CAP_DSTRD
		mov 	 byte [esi + pcidev.dstrd], al
		
		; 1003h + ((2y + 1) * (4 << CAP.DSTRD))
		mov 	 eax, 4
		shl 	 ax, cl
		mov 	 ecx, NVM_ASQS
		shl 	 ecx, 1
		inc 	 ecx
		imul 	 ecx, eax
		add 	 ecx, 0x1003

		pop 	 edx
		invoke 	 MapIoMem, edx, ecx, PG_SW+PG_NOCACHE
		mov 	 dword [esi + pcidev.io_addr], eax
		mov 	 eax, dword [eax + NVME_MMIO.VS]
		DEBUGF   DBG_INFO, "(NVMe) Controller version: 0x%x\n", eax
		mov 	 dword [esi + pcidev.version], eax
		pop 	 ecx edx esi
		xor 	 eax, eax
		inc 	 eax
		ret

.failure:
		PDEBUGF  DBG_INFO, "PCI(%u.%u.%u): something went wrong checking NVMe device compatibility\n", byte [esi + pcidev.bus], byte [esi + pcidev.devfn]
		pop 	 ecx edx esi
		xor      eax, eax
		ret
	
endp

; nvme_init: Initializes the NVMe controller 
proc nvme_init stdcall, pci:dword

		push 	 ebx esi edi
		mov 	 esi, dword [pci]
		mov      edi, dword [esi + pcidev.io_addr]

		if 0
			mov    eax, dword [edi + NVME_MMIO.CAP]
			DEBUGF DBG_INFO, "(NVMe) CAP (0-31): 0x%x\n", eax
			mov    eax, dword [edi + NVME_MMIO.CAP + 4]
			DEBUGF DBG_INFO, "(NVMe) CAP (32-63): 0x%x\n", eax
			mov    eax, dword [edi + NVME_MMIO.CC]
			DEBUGF DBG_INFO, "(NVMe) CC: 0x%x\n", eax
			mov    eax, dword [edi + NVME_MMIO.CSTS]
			DEBUGF DBG_INFO, "(NVMe) CSTS: 0x%x\n", eax
		end if

		; For some reason, bit 7 (No I/O command set supported) is also set to 1 despite bit 0 (NVM command set)
	    ; being set to 1.. so I am not sure if bit 7 should be checked at all.. investigate later.
		mov      eax, dword [edi + NVME_MMIO.CAP + 4]
		test 	 eax, CAP_CSS_NVM_CMDSET
		jz       .exit_fail

		; Reset controller before we configure it
		stdcall nvme_controller_reset, edi
		if __DEBUG__
		stdcall nvme_wait, edi
		end if
		mov 	eax, dword [edi + NVME_MMIO.CAP + 4]
		and 	eax, CAP_MPSMIN
		shr 	eax, 16
		cmp 	eax, NVM_MPS
		jg 		.exit_fail
		mov 	eax, dword [edi + NVME_MMIO.CAP + 4]
		and 	eax, CAP_MPSMAX
		shr 	eax, 20
		cmp 	eax, NVM_MPS
		jl		.exit_fail

		; Configure IOSQES, IOCQES, AMS, MPS, CSS
		and 	dword [edi + NVME_MMIO.CC], not (CC_AMS or CC_MPS or CC_CSS or CC_IOSQES or CC_IOCQES)
		mov 	eax, dword [edi + NVME_MMIO.CC]
		; 		CSS = 0 (NVM Command Set)
		; 		AMS = 0 (Round Robin)
		; 		MPS = 0 (4KiB Pages)
		; 		IOSQES = 6 (64B)
		; 		IOCQES = 4 (16B)
		or 		eax, (4 shl 20) or (6 shl 16)
		mov 	dword [edi + NVME_MMIO.CC], eax

		; Configure Admin Queue Attributes
		mov 	eax, dword [edi + NVME_MMIO.AQA]
		and 	eax, not (AQA_ASQS or AQA_ACQS)
		or 		eax, NVM_ASQS or (NVM_ACQS shl 16)
		mov 	dword [edi + NVME_MMIO.AQA], eax

		; Allocate list of queues
		invoke  KernelAlloc, sizeof.NVM_QUEUE_ENTRY * NVM_ASQS
		test 	eax, eax
		jz 		.exit_fail
		mov 	dword [esi + pcidev.queue_entries], eax
		stdcall memset, eax, 0, sizeof.NVM_QUEUE_ENTRY * NVM_ASQS

		; Configure Admin Submission/Completion Queue Base Address
		push 	esi
		mov 	esi, dword [esi + pcidev.queue_entries]
		; TODO: Allocate ring buffer? (see page 8 of NVMe 1.4 spec)
		invoke  CreateRingBuffer, 0x1000, PG_SW
		test 	eax, eax
		jz 		.exit_fail
		mov 	dword [esi + NVM_QUEUE_ENTRY.sq_ptr], eax
		invoke  GetPhysAddr
		mov 	dword [edi + NVME_MMIO.ASQ], eax
		and 	dword [edi + NVME_MMIO.ASQ + 4], 0
		; TODO: Allocate ring buffer? (see page 8 of NVMe 1.4 spec)
		invoke  CreateRingBuffer, 0x1000, PG_SW
		test 	eax, eax
		jz 		.exit_fail
		mov 	dword [esi + NVM_QUEUE_ENTRY.cq_ptr], eax
		invoke  GetPhysAddr
		mov 	dword [edi + NVME_MMIO.ACQ], eax
		and 	dword [edi + NVME_MMIO.ACQ + 4], 0

		stdcall memset, dword [esi + NVM_QUEUE_ENTRY.sq_ptr], 0, sizeof.SQ_ENTRY * NVM_CMDS
		stdcall memset, dword [esi + NVM_QUEUE_ENTRY.cq_ptr], 0, sizeof.CQ_ENTRY * NVM_CMDS

		pop 	esi

		; we want to disable all interrupts for now, since the controller randomly
	    ; generates interrupts while starting up
		;mov 	dword [edi + NVME_MMIO.INTMS], 0xffffffff 

		; Attach interrupt handler
		movzx 	eax, byte [esi + pcidev.iline]
		DEBUGF  DBG_INFO, "(NVMe) Attaching interrupt handler to IRQ %u\n", eax
		invoke  AttachIntHandler, eax, irq_handler, 0
		test 	eax, eax
		jz 		.exit_fail
		DEBUGF  DBG_INFO, "(NVMe) Successfully attached interrupt handler\n"

		; Restart the controller
		stdcall nvme_controller_start, edi
		;mov 	dword [edi + NVME_MMIO.INTMC], 0xffffffff ; re-enable interrupts

		invoke  KernelAlloc, 0x1000
		test 	eax, eax
		jz 		.exit_fail
		mov 	ebx, eax
		invoke  GetPhysAddr
		;                      pci:dword, nsid:dword, dptr:dword, cns:byte
		stdcall nvme_identify, [pci],     0,          eax,        CNS_IDCS
		mov 	edi, ebx
		lea 	ebx, byte [edi + IDENTC.sn]
		lea 	eax, byte [esi + pcidev.serial]
		stdcall memcpy, eax, ebx, 20
		DEBUGF  DBG_INFO, "(NVMe) Serial Number: %s\n", eax
		add 	ebx, 20
		lea 	eax, byte [esi + pcidev.model]
		stdcall memcpy, eax, ebx, 40
		DEBUGF  DBG_INFO, "(NVMe) Model: %s\n", eax
		mov 	edx, dword [esi + pcidev.version]
		
		cmp 	edx, VS140
		jl 		@f
		; This is a reserved field in pre-1.4 controllers
		mov 	al, byte [edi + IDENTC.cntrltype]
		cmp 	al, CNTRLTYPE_IO_CONTROLLER
		jne 	.exit_fail 	
		DEBUGF  DBG_INFO, "(NVMe) I/O controller detected...\n"

@@:
		mov 	al, byte [edi + IDENTC.sqes]
		and 	al, 11110000b
		cmp 	al, 0x60 ; maximum submission queue size should at least be 64 bytes
		jl 		.exit_fail
		mov 	al, byte [edi + IDENTC.cqes]
		and 	al, 11110000b
		and 	al, 0x40 ; maximum completion queue entry size should at least be 16 bytes
		jl 		.exit_fail
		invoke  KernelFree, edi

		mov 	eax, (NVM_ASQS - 1) or ((NVM_ACQS - 1) shl 16) ; CDW11 (set the number of queues we want)
		stdcall set_features, [pci], NULLPTR, FID_NUMBER_OF_QUEUES, eax
		mov 	esi, dword [p_nvme_devices]
		mov 	esi, dword [esi + pcidev.queue_entries]
		mov 	esi, dword [esi + NVM_QUEUE_ENTRY.cq_ptr]
		mov 	eax, dword [esi + sizeof.CQ_ENTRY + CQ_ENTRY.cdw0]
		DEBUGF  DBG_INFO, "(NVMe) Set Features CDW0: 0x%x\n", eax
		test 	ax, ax ; Number of I/O Submission Queues allocated
		jz 		.exit_fail
		shl 	eax, 16
		test 	ax, ax ; Number of I/O Completion Queues allocated
		jnz		.exit_fail

		; Create I/O Queues
		; (TODO: create N queue pairs for N CPU cores, see page 8 of NVMe 1.4 spec for an explaination
		mov 	esi, [pci]
		mov 	esi, dword [esi + pcidev.queue_entries]
		lea 	esi, [esi + sizeof.NVM_QUEUE_ENTRY]
		invoke  CreateRingBuffer, 0x1000, PG_SW
		test 	eax, eax
		jz 		.exit_fail
		mov 	dword [esi + NVM_QUEUE_ENTRY.cq_ptr], eax
		invoke 	GetPhysAddr
		stdcall create_io_completion_queue, [pci], eax, 1, IEN_ON
		invoke  CreateRingBuffer, 0x1000, PG_SW
		test 	eax, eax
		jz 		.exit_fail
		mov 	dword [esi + NVM_QUEUE_ENTRY.sq_ptr], eax
		invoke 	GetPhysAddr
		stdcall create_io_submission_queue, [pci], eax, 1, 1
		
		DEBUGF  DBG_INFO, "(NVMe) Successfully initialized driver!\n"
      	xor     eax, eax
      	inc     eax
		pop 	edi esi ebx
      	ret

.exit_fail:
		DEBUGF  DBG_INFO, "(NVMe) failed to initialize controller\n"
		xor 	eax, eax
		pop 	edi esi ebx
		ret

endp

proc get_new_cid stdcall, pci:dword, y:dword
		
		push 	esi ebx
		mov 	ecx, [y]
		imul 	ecx, sizeof.NVM_QUEUE_ENTRY
		mov 	esi, [pci]
		mov 	esi, dword [esi + pcidev.queue_entries]
		lea 	esi, dword [esi + ecx]
		mov 	eax, dword [esi + NVM_QUEUE_ENTRY.cid_slots]
		xor 	edx, edx
		cmp 	eax, 0xffffffff
		jne		@f      	
		mov 	eax, dword [esi + NVM_QUEUE_ENTRY.cid_slots + 4]
		add 	edx, 4
		cmp 	eax, 0xffffffff
		jne 	@f
		pop 	ebx esi
		mov 	eax, -1
		ret

@@:
		mov 	ebx, eax

		; Equivalant to (~(-X) & ~(X)), retrieves first free bit
		not 	eax
		mov 	ecx, eax
		neg 	ecx
		and 	eax, ecx
		or 		ebx, eax

		mov 	ecx, [y]
		mov 	dword [esi + edx + NVM_QUEUE_ENTRY.cid_slots], ebx
		bsf 	eax, eax
		cmp 	edx, 4
		jne 	@f
		add 	eax, 32

@@:
		pop 	ebx esi
		ret

endp

proc nvme_controller_reset stdcall, mmio:dword

		DEBUGF  DBG_INFO, "(NVMe) Resetting Controller...\n"
		push 	edi
		mov 	edi, dword [mmio]
		and 	dword [edi + NVME_MMIO.CC], 0xfffffffe ; CC.EN = 0

; Wait for controller to be brought to idle state, CSTS.RDY should be cleared to 0 when this happens
.wait:
		test    dword [edi + NVME_MMIO.CSTS], CSTS_RDY
		jnz 	.wait
		DEBUGF  DBG_INFO, "(NVMe) Successfully reset controller...\n"
		pop 	edi
		ret

endp

proc nvme_controller_start stdcall, mmio:dword

		DEBUGF  DBG_INFO, "(NVMe) Starting Controller...\n"
		push 	edi
		mov 	edi, dword [mmio]
		or 		dword [edi + NVME_MMIO.CC], 1 ;; CC.EN = 1

; Wait for controller to be brought into active state, CSTS.RDY should be set to 1 when this happens
.wait:
		test 	dword [edi + NVME_MMIO.CSTS], CSTS_RDY
		jz 		.wait
		DEBUGF  DBG_INFO, "(NVMe) Successfully started controller...\n"
		pop 	edi
		ret

endp

; Should be called only after the value of CC.EN has changed
proc nvme_wait stdcall, mmio:dword

		push 	esi
		mov 	esi, [mmio]
		mov 	esi, dword [esi + NVME_MMIO.CAP]
		and 	esi, CAP_TO
		shr 	esi, 24
		imul 	esi, 150 ; TODO: bad time delay, set to appropriate value later
		invoke  Sleep
		pop 	esi
		ret

endp

; Writes to completion queue 'y' head doorbell
proc cqyhdbl_write stdcall, pci:dword, y:word, cqh:word

		push 	esi edi
		mov 	esi, [pci]

		; 1000h + ((2y + 1) * (4 << CAP.DSTRD))
		movzx 	eax, [y]
		shl 	al, 1
		inc 	al
		mov 	edx, 4
		mov 	cl, byte [esi + pcidev.dstrd]
		shl 	dx, cl
		imul 	dx, ax
		add 	dx, 0x1000
		movzx 	ecx, [y]
		imul 	ecx, sizeof.NVM_QUEUE_ENTRY
		mov 	edi, dword [esi + pcidev.queue_entries]
		lea 	edi, dword [edi + ecx]
		mov 	esi, dword [esi + pcidev.io_addr]
		mov 	ax, [cqh]
		DEBUGF  DBG_INFO, "(NVMe) Writing to completion queue doorbell register 0x%x: %u\n", dx, ax
		mov 	word [esi + edx], ax ; Write to CQyHDBL
		mov 	word [edi + NVM_QUEUE_ENTRY.head], ax
		pop 	edi esi
		ret

endp

; Writes to submission queue 'y' tail doorbell
proc sqytdbl_write stdcall, pci:dword, y:word, cmd:dword

		push 	ebx esi edi
		mov 	edi, [pci]
		mov 	edi, dword [edi + pcidev.queue_entries]
		movzx 	ecx, [y]
		imul 	ecx, sizeof.NVM_QUEUE_ENTRY
		mov 	edi, dword [edi + ecx + NVM_QUEUE_ENTRY.sq_ptr]
		mov 	esi, [cmd]
		mov 	ecx, dword [esi + SQ_ENTRY.cdw0]
		shr 	ecx, 16 ; Get CID
		imul 	ecx, sizeof.SQ_ENTRY
		lea 	edi, [edi + ecx]
		stdcall memcpy, edi, esi, sizeof.SQ_ENTRY

		mov 	edi, [pci]
		mov 	esi, dword [edi + pcidev.io_addr]
		mov 	edi, dword [edi + pcidev.queue_entries]
		mov 	ax, word [edi + NVM_QUEUE_ENTRY.tail]
		cmp 	ax, NVM_ASQS
		jl		@f 		
		xor 	ax, ax
		
@@:
		mov 	esi, [pci]
		inc 	ax
		; 1000h + ((2y + 1) * (4 << CAP.DSTRD))
		movzx 	ebx, [y]
		shl 	bl, 1
		mov 	edx, 4
		mov 	cl, byte [esi + pcidev.dstrd]
		shl 	dx, cl
		imul 	dx, bx
		add 	dx, 0x1000
		DEBUGF  DBG_INFO, "(NVMe) Writing to submission queue doorbell register 0x%x: %u\n", dx, ax
		mov 	esi, dword [esi + pcidev.io_addr]
		mov 	word [esi + edx], ax
		movzx 	ecx, [y]
		mov 	word [edi + NVM_QUEUE_ENTRY.tail], ax
		dec 	ax
		stdcall nvme_cmd_wait, [pci], ecx, eax
		pop 	edi esi ebx
		ret

endp

; Calculates 2^x
proc pow2 stdcall, x:byte 

		push 	ecx
		mov 	cl, [x]
		xor 	eax, eax
		inc 	eax
		test 	cl, cl
		jnz 	@f
		pop 	ecx
		ret

@@:
		shl 	eax, cl
		pop 	ecx
		ret
		
endp

proc nvme_cmd_wait stdcall, pci:dword, y:dword, cid:dword
		
		push 	esi
		mov 	esi, [pci]
		movzx 	ecx, word [cid]
		movzx 	edx, byte [y]
		imul 	edx, sizeof.NVM_QUEUE_ENTRY
		mov 	esi, dword [esi + pcidev.queue_entries]
		lea 	esi, [esi + edx]
		imul 	ecx, sizeof.CQ_ENTRY
		mov 	esi, dword [esi + NVM_QUEUE_ENTRY.cq_ptr]

; WARNING: status bit flips on each pass, count for this later...
@@:
		test	byte [esi + ecx + CQ_ENTRY.status], CQ_PHASE_TAG
		jz 		@b
		pop 	esi
		ret

endp

proc irq_handler

		push 	esi edi
		mov 	esi, dword [p_nvme_devices]

		; check if the NVMe device generated an interrupt
		invoke 	PciRead16, dword [esi + pcidev.bus], dword [esi + pcidev.devfn], PCI_header00.status
		test 	al, 1000b ; check interrupt status
		jz 		.not_our_irq
		
		mov 	edi, esi
		mov 	edi, dword [edi + pcidev.io_addr]
		mov 	dword [edi + NVME_MMIO.INTMS], 0x1
		mov 	esi, dword [esi + pcidev.queue_entries]
		movzx 	ecx, word [esi + NVM_QUEUE_ENTRY.head]
		mov 	dx, word [esi + NVM_QUEUE_ENTRY.tail]
		DEBUGF  DBG_INFO, "IRQ (head): 0x%x, (tail): 0x%x\n", cx, dx
		cmp 	cx, dx
		je 		.end
		mov 	edx, ecx
		imul 	edx, sizeof.CQ_ENTRY
		mov 	esi, dword [p_nvme_devices]
		mov 	esi, dword [esi + pcidev.queue_entries]
		mov 	esi, dword [esi + NVM_QUEUE_ENTRY.cq_ptr]
		mov 	ax, word [esi + edx + CQ_ENTRY.status]
		and 	ax, not CQ_PHASE_TAG ; ignore phase tag bit
		DEBUGF  DBG_INFO, "(NVMe) Status: 0x%x\n", ax
		test 	al, al ; check status code (0 on success)
		jz 		.ok

.error:
		; we have to initiate a controller reset if a admin command encounters
		; a fatal error or if a completion is not received for a deletion
	    ; of a submission or completion queue (section 10.1 - page 400 of NVMe 1.4 spec)
		mov 	esi, dword [p_nvme_devices]
		mov 	esi, dword [esi + pcidev.io_addr]
		stdcall nvme_controller_reset, esi
		stdcall nvme_controller_start, esi
		jmp 	.end

.ok:
		mov 	eax, dword [esi + edx + CQ_ENTRY.cdw0]
		inc 	ecx
		; TODO: Check how many commands were consumed later
		stdcall cqyhdbl_write, dword [p_nvme_devices], 0, ecx

.end:
		mov 	edi, dword [p_nvme_devices]
		mov 	edi, dword [edi + pcidev.io_addr]
		mov 	dword [edi + NVME_MMIO.INTMC], 0x1

		; Interrupt handled by driver, return 1
		xor 	eax, eax
		inc 	eax
		pop 	edi esi
		ret

.not_our_irq:
		; Interrupt not handled by driver, return 0
		xor 	eax, eax
		pop 	edi esi
		ret

endp

proc nvme_cleanup

		DEBUGF  DBG_INFO, "(NVMe): Cleaning up...\n"
		mov 	ecx, dword [pcidevs_len]
		mov 	eax, dword [p_nvme_devices]
		test 	eax, eax
		jnz 	.loop
		ret

.loop:
		;invoke  KernelFree, dword [p_nvme_devices + ecx * sizeof.pcidev + pcidev.ident_ptr]
		dec 	ecx
		test 	ecx, ecx
		jnz 	.loop
		invoke  KernelFree, dword [p_nvme_devices]

@@:
		ret

endp

;all initialized data place here
align 4
		p_nvme_devices dd 0
		pcidevs_len  dd 0
		my_service   db "NVMe",0  ;max 16 chars include zero
		if __DEBUG__
		include_debug_strings
		end if

align 4
data fixups
end data