;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;;                                                                 ;;
;; 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
		mov 	ebx, eax

.loop:
		stdcall device_is_compat, dword [ebx + ecx * sizeof.pcidev]
		test 	eax, eax
		jz 		@f
		stdcall nvme_init, dword [ebx + ecx * sizeof.pcidev]

@@:
		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 edx
		mov 	edx, [sz]
		mov     bl, [val]
		xor     ecx, ecx 
@@:
		mov     byte [p_data + ecx], bl
		inc     ecx
		cmp     ecx, edx
		jne     @b
		pop     edx ecx ebx
		ret

endp

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

		xor 	eax, eax
		ret

endp

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

		sub     esp, sizeof.SQ_ENTRY
		stdcall memset, esp, 0, sizeof.SQ_ENTRY

		mov     eax, dword [nsid]
		mov     dword [esp + SQ_ENTRY.nsid], eax
		mov     eax, dword [dptr]
		mov     dword [esp + SQ_ENTRY.dptr], eax
		mov     dword [esp + SQ_ENTRY.cdw0], ADM_CMD_IDENTIFY ; TODO: setting CID to 0 for now but later on keep a unique list of identifiers
		mov     al, byte [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, TOTAL_PCIDEVS_MALLOC_SZ
		test 	eax, eax
		jz 		.err_no_mem
		mov     dword [p_nvme_devices], eax 

@@:
		mov 	ebx, dword [p_nvme_devices]
		mov 	ecx, dword [pcidevs_len]
		dec 	ecx
		pop		eax
		movzx 	edx, byte [eax + PCIDEV.bus]
		mov 	byte [ebx + ecx * sizeof.pcidev + pcidev.bus], dl
		movzx 	edx, byte [eax + PCIDEV.devfn]
		mov 	byte [ebx + ecx * sizeof.pcidev + 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 	 ebx
		invoke   PciRead32, dword [pci + pcidev.bus], dword [pci + pcidev.devfn], PCI_header00.base_addr_0
		and      eax, 0xfffffff0
		test     eax, eax
		jz       .failure

		invoke   MapIoMem, eax, sizeof.NVME_REG_MAP, PG_SW+PG_NOCACHE
		test     eax, eax
		jz       .failure
		mov 	 dword [pci + pcidev.mmio_ptr], eax
		mov 	 ebx, dword [eax + NVME_REG_MAP.VS]
		if __DEBUG__
			push 	 ecx
			mov  	 ecx, ebx
			shr  	 ecx, 16
			PDEBUGF  DBG_INFO, "PCI(%u.%u.%u): NVMe controller version: %u.%u.%u\n", byte [pci + pcidev.bus], byte [pci + pcidev.devfn], ecx, bh, bl
			pop 	 ecx
		end if
		cmp 	 ebx, NVM_SUPPORTED_CONTROLLER_VERSION
		jne 	 .unsupported
		PDEBUGF  DBG_INFO, "PCI(%u.%u.%u): NVMe controller version is supported\n", byte [pci + pcidev.bus], byte [pci + pcidev.devfn]
		xor 	 eax, eax
		inc 	 eax
		pop 	 ebx
		ret

.unsupported:
		PDEBUGF  DBG_INFO, "PCI(%u.%u.%u): unsupported controller version\n", byte [pci + pcidev.bus], byte [pci + pcidev.devfn]

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

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

		push 	 ebx ecx
		mov      eax, dword [pci + pcidev.mmio_ptr]
		mov      bx, word [eax + NVME_REG_MAP.CAP]
		PDEBUGF  DBG_INFO, "PCI(%u.%u.%u): NVMe maximum queue entries supported: %u\n", byte [pci + pcidev.bus], byte [pci + pcidev.devfn], bx
		test     ebx, CAP_CQR
		jz       .cqr_not_req
		PDEBUGF  DBG_INFO, "PCI(%u.%u.%u): NVMe contiguous queues required\n", byte [pci + pcidev.bus], byte [pci + pcidev.devfn]

.cqr_not_req:
		mov      ebx, dword [eax + NVME_REG_MAP.CAP + 4]
		mov      ecx, ebx
		test     ebx, CAP_CSS_NVM_CMDSET
		jz       .exit_fail

		if __DEBUG__
			and      ebx, CAP_MPSMIN
			and      ecx, CAP_MPSMAX
			shr      ebx, 16
			shr      ecx, 16
			PDEBUGF  DBG_INFO, "PCI(%u.%u.%u): NVMe memory page size minimum: %u\n", byte [pci + pcidev.bus], byte [pci + pcidev.devfn], ebx
			PDEBUGF  DBG_INFO, "PCI(%u.%u.%u): NVMe memory page size maximum: %u\n", byte [pci + pcidev.bus], byte [pci + pcidev.devfn], ecx

			mov      ebx, dword [eax + NVME_REG_MAP.CC]
			mov      ecx, ebx
			and      ebx, CC_IOSQES
			and      ecx, CC_IOCQES
			shl      ebx, 16
			shl      ecx, 16

			; TODO: Change entry sizes to their appropriate values
			PDEBUGF  DBG_INFO, "PCI(%u.%u.%u): NVMe I/O submission queue entry size: %u\n", byte [pci + pcidev.bus], byte [pci + pcidev.devfn], ebx
			PDEBUGF  DBG_INFO, "PCI(%u.%u.%u): NVMe I/O completion queue entry size: %u\n", byte [pci + pcidev.bus], byte [pci + pcidev.devfn], ecx
		end if

        xor     eax, eax
        inc     eax
		pop     ecx ebx
        ret

.exit_fail:
		PDEBUGF DBG_INFO, "PCI(%u.%u.%u): failed to initialize NVMe controller\n", byte [pci + pcidev.bus], byte [pci + pcidev.devfn]
		xor 	eax, eax
		pop     ecx ebx
		ret

endp

proc nvme_cleanup

		DEBUGF  DBG_INFO, "(NVMe): Cleaning up...\n"
		push 	ecx
		xor 	ecx, ecx
		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]
		inc 	ecx
		cmp 	ecx, dword [pcidevs_len]
		jne 	.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