;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;;                                                                 ;;
;; 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

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"

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

;@@:
		;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 memset stdcall, p_data:dword, val:byte, sz:dword

		push    ebx ecx
		mov 	eax, [p_data]
		mov 	ecx, [sz]
		mov     bl, [val]

@@:
		mov     byte [eax + ecx], bl
		dec 	ecx
		test 	ecx, ecx
		jnz     @b
		pop     ecx ebx
		ret

endp

; Submit a Command in the Admin Submission Queue
proc submit_asq stdcall, sq_ptr:dword

		mov 	esi, [sq_ptr]
		xor 	eax, eax
		ret

endp

proc nvme_identify stdcall, nsid:dword, dptr:dword, cid:word, cns:byte

		sub     esp, sizeof.SQ_ENTRY
		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
		movzx 	eax, [cid]
		or 		dword [esp + SQ_ENTRY.cdw0], ADM_CMD_IDENTIFY
		or 		dword [esp + SQ_ENTRY.cdw0], eax
		mov     al, [cns]
		mov     byte [esp + SQ_ENTRY.cdw10], al
		stdcall submit_asq, esp

		add     esp, sizeof.SQ_ENTRY
		xor     eax, eax
		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
		mov 	 esi, [pci]
		invoke   PciRead32, dword [esi + pcidev.bus], dword [esi + pcidev.devfn], PCI_header00.base_addr_0
		and      eax, 0xfffffff0
		test     eax, eax
		jz       .failure

		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.mmio_ptr], eax
		mov 	 eax, dword [eax + NVME_MMIO.CAP + 4]
		and 	 eax, CAP_DSTRD

		; Stride is (2 ^ (2 + DSTRD)) bytes
		add 	 eax, 2
		stdcall  pow2, eax
		mov 	 dword [esi + pcidev.dstrd], eax
		
		; TODO: Remap I/O with new size

		mov 	 esi, dword [esi + pcidev.mmio_ptr]
		mov 	 eax, dword [esi + NVME_MMIO.VS]
		DEBUGF   DBG_INFO, "(NVMe) Controller version: 0x%x\n", eax
		pop 	 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 	 esi
		xor      eax, eax
		ret
	
endp

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

		push 	 eax
		mov 	 edi, dword [pci]
		mov      edi, dword [edi + pcidev.mmio_ptr]

		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

		mov      eax, dword [edi + NVME_MMIO.CAP]
		test     eax, CAP_CQR
		jz       .cqr_not_req

.cqr_not_req:
		; 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
		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 AMS, MPS, CSS
		mov 	eax, dword [edi + NVME_MMIO.CC]
		and 	eax, not (CC_AMS or CC_MPS or CC_CSS)
		or 		eax, 111b shl 4 ; Admin Command Set Only (temporary)
		and 	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

		; Configure Admin Submission/Completion Queue Base Address
		invoke  AllocPage
		test 	eax, eax
		jz 		.exit_fail
		;or 	eax, PG_SW or PG_NOCACHE
		mov 	dword [edi + NVME_MMIO.ASQ], eax
		and 	dword [edi + NVME_MMIO.ASQ + 4], 0
		;DEBUGF  DBG_INFO, "(NVMe) ASQ = 0x%x\n", eax
		invoke  AllocPage
		test 	eax, eax
		jz 		.exit_fail
		;or 	eax, PG_SW or PG_NOCACHE
		mov 	dword [edi + NVME_MMIO.ACQ], eax
		and 	dword [edi + NVME_MMIO.ACQ + 4], 0
		;DEBUGF  DBG_INFO, "(NVMe) ACQ = 0x%x\n", eax

		; Restart the controller
		stdcall nvme_controller_start, edi

		invoke  AllocPage
		test 	eax, eax
		jz 		.exit_fail
		;                      nsid:dword, dptr:dword, cid:word, cns:byte
		stdcall nvme_identify, 0,          eax,        0,        CNS_IDCS
		
      	xor     eax, eax
      	inc     eax
      	ret

.exit_fail:
		DEBUGF DBG_INFO, "(NVMe) failed to initialize NVMe controller\n"
		xor 	eax, eax
		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
		stdcall nvme_wait, [mmio]

; 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
		stdcall nvme_wait, [mmio]

.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

		mov 	esi, [mmio]
		mov 	esi, dword [esi + NVME_MMIO.CAP]
		and 	esi, CAP_TO
		shr 	esi, 24
		imul 	esi, 50
		invoke  Sleep
		ret

endp

proc cq0thbl_write stdcall, dstrd:dword

		; 1000h + (2 * (4 << CAP.DSTRD))
		mov 	eax, [dstrd]
		shl 	eax, 4
		imul 	eax, 2
		add 	eax, 0x1000
		ret

endp

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

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

@@:
		shl 	eax, ecx
		pop 	ecx
		ret
		
endp

proc nvme_cleanup

		DEBUGF  DBG_INFO, "(NVMe): Cleaning up...\n"
		push 	ecx
		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]

@@:
		pop 	ecx
		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