#include "nvme.h"
#include "kmem.h"
#include "platform.h"
/* PCI helpers */
static inline void _pci_sel(unsigned b,unsigned d,unsigned f,unsigned r){
    unsigned a=0x80000000u|(b<<16)|(d<<11)|(f<<8)|(r&0xFC);
    __asm__ volatile("outl %0,%1"::"a"(a),"Nd"((unsigned short)PCI_CONFIG_ADDR));}
static unsigned _pci32(unsigned b,unsigned d,unsigned f,unsigned r){
    _pci_sel(b,d,f,r); unsigned v;
    __asm__ volatile("inl %1,%0":"=a"(v):"Nd"((unsigned short)PCI_CONFIG_DATA)); return v;}
static void _pciw32(unsigned b,unsigned d,unsigned f,unsigned r,unsigned v){
    _pci_sel(b,d,f,r);
    __asm__ volatile("outl %0,%1"::"a"(v),"Nd"((unsigned short)PCI_CONFIG_DATA));}
/* NVMe regs */
#define NVME_CAP  0x00
#define NVME_CC   0x14
#define NVME_CSTS 0x1C
#define NVME_AQA  0x24
#define NVME_ASQ  0x28
#define NVME_ACQ  0x30
#define NVME_CC_EN      (1u<<0)
#define NVME_CC_CFG     ((4u<<20)|(6u<<16))
#define NVME_CSTS_RDY   (1u<<0)
#define NVME_CSTS_CFS   (1u<<1)
#define ASQ_D 4
#define ACQ_D 4
#define IOSQ_D 16
#define IOCQ_D 16
struct NvmeSqe{unsigned int cdw0,nsid;unsigned long long r,m,p1,p2;
    unsigned int c10,c11,c12,c13,c14,c15;} __attribute__((packed));
struct NvmeCqe{unsigned int d0,r;unsigned short sh,si,cid,st;} __attribute__((packed));
static volatile unsigned char *g_b0=0;
static unsigned g_str=0,g_nsid=1;
static int g_ok=0;
static NvmeSqe *g_asq,*g_iosq;
static NvmeCqe *g_acq,*g_iocq;
static unsigned char *g_dma;
static unsigned g_at=0,g_ah=0,g_it=0,g_ih=0,g_ap=1,g_ip=1;
static unsigned short g_cid=0;
static unsigned _r32(unsigned o){return *((volatile unsigned*)(g_b0+o));}
static void _w32(unsigned o,unsigned v){*((volatile unsigned*)(g_b0+o))=v;}
static void _w64(unsigned o,unsigned long long v){*((volatile unsigned long long*)(g_b0+o))=v;}
static void _mfence(void){__asm__ volatile("mfence":::"memory");}
static int _wait(unsigned exp,int ms){
    for(int i=0;i<ms*100;i++){if((_r32(NVME_CSTS)&NVME_CSTS_RDY)==exp)return 1;
    for(volatile int d=0;d<10000;d++);}return 0;}
static void _db(unsigned q){
    unsigned o=0x1000+g_str*2*q;
    _w32(o,q==0?g_at:g_it);}
static void _dbc(unsigned q){
    unsigned o=0x1000+g_str*(2*q+1);
    _w32(o,q==0?g_ah:g_ih);}
static int _admin(NvmeSqe *cmd){
    cmd->cdw0=(cmd->cdw0&0xFFFF0000u)|((unsigned)(++g_cid)<<16)|(cmd->cdw0&0xFFFFu);
    g_asq[g_at]=*cmd; _mfence();
    g_at=(g_at+1)%ASQ_D; _db(0);
    for(int i=0;i<50000;i++){_mfence();
        if((g_acq[g_ah].st&1)==(unsigned short)g_ap){
            unsigned short s=(g_acq[g_ah].st>>1)&0xFF;
            g_ah=(g_ah+1)%ACQ_D; if(!g_ah)g_ap^=1;
            _dbc(0); return s?-1:0;}
        for(volatile int d=0;d<100;d++);}return -2;}
int nvme_init(void){
    unsigned long long b0=0;
    for(unsigned bus=0;bus<256&&!b0;bus++)for(unsigned dv=0;dv<32&&!b0;dv++)for(unsigned fn=0;fn<8&&!b0;fn++){
        unsigned id=_pci32(bus,dv,fn,0); if(id==0xFFFFFFFF||!id)continue;
        unsigned cr=_pci32(bus,dv,fn,8);
        if(((cr>>24)&0xFF)==0x01&&((cr>>16)&0xFF)==0x08){
            _pciw32(bus,dv,fn,4,_pci32(bus,dv,fn,4)|6);
            unsigned l=_pci32(bus,dv,fn,0x10)&~0xFu;
            unsigned h=_pci32(bus,dv,fn,0x14);
            b0=((unsigned long long)h<<32)|l;}}
    if(!b0)return 0;
    g_b0=(volatile unsigned char*)b0;
    unsigned long long cap=*((volatile unsigned long long*)(g_b0));
    g_str=4u<<(unsigned)((cap>>32)&0xF);
    _w32(NVME_CC,_r32(NVME_CC)&~NVME_CC_EN); if(!_wait(0,500))return 0;
    unsigned long long as=(unsigned long long)(unsigned long)(g_asq=(NvmeSqe*)kmem_alloc(ASQ_D*sizeof(NvmeSqe),"nvme-a"));
    unsigned long long ac=(unsigned long long)(unsigned long)(g_acq=(NvmeCqe*)kmem_alloc(ACQ_D*sizeof(NvmeCqe),"nvme-c"));
    g_iosq=(NvmeSqe*)kmem_alloc(IOSQ_D*sizeof(NvmeSqe),"nvme-is");
    g_iocq=(NvmeCqe*)kmem_alloc(IOCQ_D*sizeof(NvmeCqe),"nvme-ic");
    g_dma=(unsigned char*)kmem_alloc(131072,"nvme-d");
    if(!g_asq||!g_acq||!g_iosq||!g_iocq||!g_dma)return 0;
    auto z=[](void*p,unsigned long long n){unsigned char*b=(unsigned char*)p;for(unsigned long long i=0;i<n;i++)b[i]=0;};
    z(g_asq,ASQ_D*64);z(g_acq,ACQ_D*16);z(g_iosq,IOSQ_D*64);z(g_iocq,IOCQ_D*16);
    _w32(NVME_AQA,((ACQ_D-1)<<16)|(ASQ_D-1)); _w64(NVME_ASQ,as); _w64(NVME_ACQ,ac);
    _w32(NVME_CC,NVME_CC_EN|NVME_CC_CFG); if(!_wait(1,500))return 0;
    if(_r32(NVME_CSTS)&NVME_CSTS_CFS)return 0;
    {NvmeSqe cmd={};cmd.cdw0=0x05;cmd.p1=(unsigned long long)(unsigned long)g_iocq;
     cmd.c10=(1<<16)|1;cmd.c11=((IOCQ_D-1)<<16)|1;if(_admin(&cmd))return 0;}
    {NvmeSqe cmd={};cmd.cdw0=0x01;cmd.p1=(unsigned long long)(unsigned long)g_iosq;
     cmd.c10=(1<<16)|1;cmd.c11=(IOSQ_D-1)<<16;cmd.c12=1;if(_admin(&cmd))return 0;}
    g_ok=1;return 1;}
static int _io(int wr,unsigned long long lba,unsigned n,void*buf){
    if(!g_ok||!n||n>256)return 0; unsigned b=n*512; if(b>131072)return 0;
    if(wr){const unsigned char*s=(const unsigned char*)buf;for(unsigned i=0;i<b;i++)g_dma[i]=s[i];}
    NvmeSqe cmd={};cmd.cdw0=wr?0x01u:0x02u; cmd.nsid=g_nsid;
    cmd.p1=(unsigned long long)(unsigned long)g_dma;
    cmd.c10=(unsigned)(lba&0xFFFFFFFF);cmd.c11=(unsigned)(lba>>32);cmd.c12=n-1;
    unsigned short cv=++g_cid; cmd.cdw0|=(unsigned)cv<<16;
    g_iosq[g_it]=cmd; _mfence(); g_it=(g_it+1)%IOSQ_D; _db(1);
    for(int i=0;i<50000;i++){_mfence();
        if((g_iocq[g_ih].st&1)==(unsigned short)g_ip){
            unsigned short s=(g_iocq[g_ih].st>>1)&0xFF;
            g_ih=(g_ih+1)%IOCQ_D; if(!g_ih)g_ip^=1; _dbc(1);
            if(s)return 0;
            if(!wr){unsigned char*d=(unsigned char*)buf;for(unsigned j=0;j<b;j++)d[j]=g_dma[j];}
            return 1;}
        for(volatile int d=0;d<100;d++);}return 0;}
int nvme_read (unsigned long long l,unsigned n,void *b){return _io(0,l,n,b);}
int nvme_write(unsigned long long l,unsigned n,const void*b){return _io(1,l,n,(void*)b);}
int nvme_found(void){return g_ok;}