;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;;                                                              ;;
;; Copyright (C) KolibriOS team 2004-2024. All rights reserved. ;;
;; Distributed under terms of the GNU General Public License    ;;
;;                                                              ;;
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

;driver sceletone

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"

proc START c, reason:dword
        cmp     [reason], DRV_ENTRY
        jne     .exit
.entry:
        push    esi
		DEBUGF DBG_INFO,"Detecting NVMe hardware...\n"
        call    detect
        pop     esi
        test    eax, eax
        jz      .exit

        invoke  RegService, my_service, service_proc
        ret
.exit:
        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     .fail
        mov     dword [eax], API_VERSION
        xor     eax, eax
        ret
@@:
.fail:
        or      eax, -1
        ret
endp

proc memset stdcall, p_data:dword, val:byte, sz:dword
	mov bh, byte [val]
	xor eax, eax
@@:
	mov byte [p_data + eax], bh
	inc eax
	cmp eax, dword [sz]
	jne @b
	ret
endp

; Submit a Command in the Admin Submission Queue
proc submit_asq stdcall, p_sq:dword
	mov ebx, dword [bar0]
	add ebx, NVME_REG_ASQ

	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 [esp + SQ_ENTRY.nsid], eax
	mov eax, dword [dptr]
	mov dword [esp + SQ_ENTRY.dptr], eax
	; TODO: setting CID to 1 for now but later on keep a list of unique list of identifiers
	mov dword [esp + SQ_ENTRY.cdw0], ADM_CMD_IDENTIFY or (1 shl 16)
	mov ah, byte [cns]
	mov byte [esp + SQ_ENTRY.cdw10], ah
	stdcall submit_asq, esp
	add esp, sizeof.SQ_ENTRY
	xor eax, eax
	ret
endp

proc detect
        push    ebx
        invoke  GetPCIList
		mov 	edx, eax
.check_dev:
		mov 	ecx, [eax+PCIDEV.class]
		and		ecx, 0x00ffff00 ; retrieve class/subclass code only
		cmp 	ecx, 0x00010800 ; Mass Storage Controller - Non-Volatile Memory Controller
		je		.check_cap
.next_dev:
		mov		eax, [eax + PCIDEV.fd]
		cmp		eax, edx
		jne 	.check_dev

		; no more PCI devices to enumerate?
		xor     eax, eax
		pop     ebx
        ret
.check_cap:
		push eax
		movzx ebx, [eax + PCIDEV.bus]
		mov [pcidev_bus], ebx
		movzx ebx, [eax + PCIDEV.devfn]
		mov [pcidev_devfn], ebx
		invoke PciRead16, [pcidev_bus], [pcidev_devfn], PCI_header00.status
		test ax, 0x10 ; check capabilities list bit
		jnz .got_cap
		pop eax
.got_cap:
		DEBUGF DBG_INFO,"Found NVMe device with capabilities\n"
		invoke PciRead8, [pcidev_bus], [pcidev_devfn], PCI_header00.cap_ptr
		and eax, 11111100b
		mov edi, eax
@@:
		invoke PciRead32, [pcidev_bus], [pcidev_devfn], edi
		mov ecx, eax
		;and ecx, 0xff
		mov eax, ecx
		movzx edi, ah
		test edi, edi
		jnz @b

		; read BAR0
		invoke PciRead32, [pcidev_bus], [pcidev_devfn], PCI_header00.base_addr_0
		mov [bar0], eax

		invoke Kmalloc, sizeof.NVME_IDENT_CONTROLLER
		test eax, eax
		jz .alloc_ident_controller_fail
		DEBUGF DBG_INFO,"Successfully allocated %u bytes for 'NVME_IDENT_CONTROLLER'\n",sizeof.NVME_IDENT_CONTROLLER
		mov [p_ident], eax

		stdcall nvme_identify, 0, dword [p_ident], CNS_IDCS
		mov ecx, dword [p_ident]
		mov edx, dword [ecx + NVME_IDENT_CONTROLLER.tnvmcap]
		DEBUGF DBG_INFO,"Total NVMe SSD capacity: %ub\n", edx

		; return successfully
		call nvme_cleanup
		pop 	eax
        pop     ebx
        xor     eax, eax
        inc     eax
        ret

.alloc_ident_controller_fail:
		DEBUGF DBG_INFO,"ERROR: failed to allocate %u bytes for 'NVME_IDENT_CONTROLLER'\n", sizeof.NVME_IDENT_CONTROLLER
		pop eax
		pop ebx
		xor eax, eax
		ret
endp

proc nvme_cleanup
	mov eax, [p_ident]
	test eax, eax
	jz @f
	invoke Kfree, eax
@@:
	ret
endp

; uninitialized data
align 4
pcidev_bus			dd ?
pcidev_devfn		dd ?
bar0                dd ?
p_ident             dd ?

;all initialized data place here
align 4
my_service   db "NVMe Service",0  ;max 16 chars include zero
include_debug_strings

align 4
data fixups
end data