using System.Runtime.InteropServices;

namespace XFESystem.NativeKernel;

public unsafe struct DiskController
{
    public ulong Base, Bytes, AdminSq, AdminCq, IoSq, IoCq, Data, Prps, MaxTransfer;
    public uint Pci, Kind, Active, Stride, AdminTail, AdminHead, AdminPhase, IoTail, IoHead, IoPhase, Tag;
}
public unsafe struct HardwareDisk
{
    public DiskController* Controller;
    public ulong Blocks, Port, Commands, Fis, Table, AllowedStart, AllowedCount, Reads, Writes, Flushes;
    public uint Active, Namespace, SectorShift, CanWrite, FlushCommand;
}

// Native scheduler worker owns command queues and DMA. No GC pointer survives
// an ABI call. Media commands are READ/WRITE/FLUSH only; no format/TRIM/security.
public static unsafe class DiskHardware
{
    const uint Limit=32, Controllers=8;
    static HardwareDisk* _disks;
    static DiskController* _controllers;
    static uint _count, _controllersCount, _state, _requestState, _disk, _operation, _sectors;
    static ulong _lba, _mailbox;
    [DllImport("xfe_native",EntryPoint="xfe_network_pci_read")] static extern uint PciRead(uint b,uint d,uint f,uint o);
    [DllImport("xfe_native",EntryPoint="xfe_network_pci_write")] static extern int PciWrite(uint b,uint d,uint f,uint o,uint v);
    static uint Read(ulong p)=>Native.MmioRead32(p);
    static void Write(ulong p,uint v)=>Native.MmioWrite32(p,v);
    static void Write64(ulong p,ulong v){Write(p,(uint)v);Write(p+4,(uint)(v>>32));}
    static ulong Now()=>MonotonicClock.Nanoseconds()/1000000;
    static byte* At(ulong p)=>p==0?null:(byte*)PhysicalMemory.ToVirtual(p);
    static ulong Allocate(uint pages=1){ulong irq=Native.IrqSave();ulong p=PhysicalMemory.Allocate(pages);Native.IrqRestore(irq);return p;}
    static bool Wait(ulong p,uint mask,uint wanted,uint ms)
    {ulong end=Now()+ms;do{if((Read(p)&mask)==wanted)return true;Threads.Sleep(1);}while(Now()<end);return false;}
    static void Log(string message){Serial.Write("XFE_DISK|");Serial.Line(message);}
    static bool MapAllowed(ulong physical,ulong bytes)
    {
        if(physical<0x100000||(physical&4095)!=0||bytes<4096||bytes>2097152||physical>~0UL-bytes)return false;
        for(ulong i=0;i<Kernel.Boot->MemoryEntryCount;i++)
        {MemoryMapEntry* e=Kernel.Boot->MemoryEntries[i];if(e->Length!=0&&physical<e->Base+e->Length&&e->Base<physical+bytes&&e->Kind!=1&&e->Kind!=8)return false;}
        return true;
    }
    static bool MapBar(DiskController* c,uint bar,bool size)
    {
        uint b=c->Pci>>16,d=c->Pci>>8&31,f=c->Pci&7;
        uint low=PciRead(b,d,f,bar),high=(low&6)==4?PciRead(b,d,f,bar+4):0;
        if((low&1)!=0||(low&6)is not(0 or 4))return false;
        ulong physical=(ulong)high<<32|(low&~15u),bytes=c->Kind==1?8192UL:4096UL;
        if(!MapAllowed(physical,bytes))return false;
        ulong irq=Native.IrqSave();
        if(size)
        {
            uint command=PciRead(b,d,f,4)&65535;
            PciWrite(b,d,f,4,command&~6u);PciWrite(b,d,f,bar,0xffffffff);if((low&6)==4)PciWrite(b,d,f,bar+4,0xffffffff);
            ulong mask=(PciRead(b,d,f,bar)&~15u)|((ulong)((low&6)==4?PciRead(b,d,f,bar+4):0xffffffff)<<32);
            PciWrite(b,d,f,bar,low);if((low&6)==4)PciWrite(b,d,f,bar+4,high);PciWrite(b,d,f,4,command);
            bytes=~mask+1;
            if((bytes&(bytes-1))!=0||!MapAllowed(physical,bytes)){Native.IrqRestore(irq);return false;}
            VirtualMemory.UnmapDevice(c->Base,c->Bytes);
        }
        c->Base=VirtualMemory.MapDevice(physical,bytes);c->Bytes=bytes;Native.IrqRestore(irq);return c->Base!=0;
    }
    static void EnablePci(DiskController* c)
    {uint b=c->Pci>>16,d=c->Pci>>8&31,f=c->Pci&7;uint command=PciRead(b,d,f,4)&65535;PciWrite(b,d,f,4,(command|0x406u)&~0x100u);}
    public static bool Start(){_state=1;return Threads.Create(&Worker,null,262144,0,false)!=0;}
    [UnmanagedCallersOnly(EntryPoint="xfe_disk_worker")]
    public static uint Worker(void* argument)
    {
        _disks=(HardwareDisk*)At(Allocate(2));_controllers=(DiskController*)At(Allocate());_mailbox=Allocate(16);
        if(_disks==null||_controllers==null||_mailbox==0){_state=4;return 1;}
        for(uint b=0;b<256;b++)for(uint d=0;d<32;d++)
        {
            if((PciRead(b,d,0,0)&65535)==65535)continue;
            uint functions=(PciRead(b,d,0,12)&0x800000)!=0?8u:1u;
            for(uint f=0;f<functions;f++)
            {
                if((PciRead(b,d,f,0)&65535)==65535)continue;
                uint type=PciRead(b,d,f,8)>>8;if(type!=0x010802&&type!=0x010601)continue;
                if(_controllersCount==Controllers)continue;
                DiskController* c=&_controllers[_controllersCount++];c->Pci=b<<16|d<<8|f;c->Kind=type==0x010802?2u:1u;
                bool ok=c->Kind==2?InitializeNvme(c):InitializeAhci(c);if(!ok)Log("CONTROLLER_UNAVAILABLE");
            }
        }
        Native.MemoryBarrier();_state=2;Log("ENUMERATION_COMPLETE");
        for(;;){ProcessRequest();Threads.Sleep(1);}
    }
    static void Publish(HardwareDisk* disk)
    {
        Serial.Write("XFE_DISK_READY|Index|");Serial.Number(_count);Serial.Write("|Kind|");Serial.Number(disk->Controller->Kind);
        Serial.Write("|Sectors512|");Serial.Number(disk->Blocks);Serial.Write("|SectorShift|");Serial.Number(disk->SectorShift);Serial.Line("");
        Native.MemoryBarrier();_count++;
    }
    static bool StopPort(ulong p)
    {Write(p+24,Read(p+24)&~1u);if(!Wait(p+24,1u<<15,0,1000))return false;Write(p+24,Read(p+24)&~16u);return Wait(p+24,1u<<14,0,1000);}
    static bool InitializeAhci(DiskController* c)
    {
        if(!MapBar(c,36,false))return false;
        if((Read(c->Base+36)&1)!=0)
        {Write(c->Base+40,Read(c->Base+40)|2);if(!Wait(c->Base+40,17,0,2500))return false;}
        Write(c->Base+4,(Read(c->Base+4)|0x80000000u)&~2u);
        uint ports=Read(c->Base+12);
        // The first mapping covers all 32 standard port-register blocks.
        for(uint i=0;i<32;i++)if((ports&(1u<<(int)i))!=0&&!StopPort(c->Base+256+i*128))return false;
        uint portCount=(Read(c->Base)&31)+1;
        if(!MapBar(c,36,true)||c->Bytes<256+portCount*128)return false;
        EnablePci(c);c->Data=Allocate(32);if(c->Data==0)return false;c->Active=1;
        bool address64=(Read(c->Base)&0x80000000)!=0;
        for(uint i=0;i<32&&_count<Limit;i++)
        {
            if((ports&(1u<<(int)i))==0)continue;ulong port=c->Base+256+i*128;
            uint status=Read(port+40);if((status&15)!=3||((status>>8)&15)!=1||Read(port+36)!=0x101)continue;
            HardwareDisk* disk=&_disks[_count];*disk=new HardwareDisk();disk->Controller=c;disk->Port=port;disk->SectorShift=9;
            disk->Commands=Allocate();disk->Fis=Allocate();disk->Table=Allocate();
            if(disk->Commands==0||disk->Fis==0||disk->Table==0||!address64&&(disk->Commands>=0x100000000||disk->Fis>=0x100000000||disk->Table>=0x100000000||c->Data+131072>0x100000000))continue;
            Write(port+20,0);Write64(port,disk->Commands);Write64(port+8,disk->Fis);Write(port+48,0xffffffff);Write(port+16,0xffffffff);
            Write(port+24,Read(port+24)|17);disk->Active=1;
            if(!AhciCommand(disk,0xec,0,0,512,false))continue;
            ushort* id=(ushort*)At(c->Data);
            if((id[83]&0x400)==0||((id[106]&0xd000)==0x5000&&((uint)id[117]|(uint)id[118]<<16)!=256)){disk->Active=0;continue;}
            disk->Blocks=(ulong)id[100]|(ulong)id[101]<<16|(ulong)id[102]<<32|(ulong)id[103]<<48;
            disk->FlushCommand=(id[83]&0x2000)!=0?0xeau:0xe7u;disk->CanWrite=1;
            if(disk->Blocks==0||disk->Blocks>0x1000000000000){disk->Active=0;continue;}Publish(disk);
        }
        return true;
    }
    static bool AhciCommand(HardwareDisk* d,uint command,ulong lba,uint count,uint bytes,bool write)
    {
        if(d->Active==0||(command!=0xec&&command!=0x25&&command!=0x35&&command!=0xea&&command!=0xe7)||bytes>65536)return false;
        ulong p=d->Port;
        if(!Wait(p+32,0x88,0,2000)||Read(p+56)!=0||Read(p+52)!=0)return false;
        byte* table=At(d->Table);for(uint i=0;i<160;i++)table[i]=0;
        table[0]=0x27;table[1]=128;table[2]=(byte)command;table[7]=64;
        table[4]=(byte)lba;table[5]=(byte)(lba>>8);table[6]=(byte)(lba>>16);table[8]=(byte)(lba>>24);table[9]=(byte)(lba>>32);table[10]=(byte)(lba>>40);
        table[12]=(byte)count;table[13]=(byte)(count>>8);
        uint* header=(uint*)At(d->Commands);for(uint i=0;i<8;i++)header[i]=0;
        header[0]=5u|(write?64u:0)|(bytes!=0?1u<<16:0);((ulong*)(header+2))[0]=d->Table;
        if(bytes!=0){*(ulong*)(table+128)=d->Controller->Data;*(uint*)(table+140)=bytes-1;}
        Write(p+16,0xffffffff);Native.MemoryBarrier();Write(p+56,1);ulong end=Now()+5000;
        do
        {
            if((Read(p+16)&0x40000000)!=0)break;
            if((Read(p+56)&1)==0){Native.MemoryBarrier();return (Read(p+32)&1)==0&&(bytes==0||header[1]==bytes);}
            Threads.Sleep(1);
        }while(Now()<end);
        StopPort(p);d->Active=0;Log("AHCI_COMMAND_FAILED_QUARANTINED");return false;
    }
    static bool NvmeCommand(DiskController* c,bool admin,uint opcode,uint ns,ulong prp,uint dw10,uint dw11=0,uint dw12=0,ulong dataBytes=4096)
    {
        if(c->Active==0)return false;
        uint tail=admin?c->AdminTail:c->IoTail,head=admin?c->AdminHead:c->IoHead,phase=admin?c->AdminPhase:c->IoPhase;
        uint* command=(uint*)At((admin?c->AdminSq:c->IoSq)+tail*64);for(uint i=0;i<16;i++)command[i]=0;
        uint tag=++c->Tag&65535;command[0]=opcode|tag<<16;command[1]=ns;
        ((ulong*)(command+6))[0]=prp;
        if(prp==c->Data)((ulong*)(command+8))[0]=dataBytes<=8192?c->Data+4096:c->Prps;
        command[10]=dw10;command[11]=dw11;command[12]=dw12;
        tail=(tail+1)%16;Native.MemoryBarrier();Write(c->Base+4096+(admin?0u:2u)*c->Stride,tail);
        uint* completion=(uint*)At((admin?c->AdminCq:c->IoCq)+head*16);ulong end=Now()+5000;
        do
        {
            Native.MemoryBarrier();uint status=Read((ulong)(completion+3));
            if(((status>>16)&1)==phase)
            {
                bool ok=(status>>17)==0&&(status&65535)==tag&&(completion[2]>>16)==(admin?0u:1u);
                if(++head==16){head=0;phase^=1;}
                if(admin){c->AdminTail=tail;c->AdminHead=head;c->AdminPhase=phase;}else{c->IoTail=tail;c->IoHead=head;c->IoPhase=phase;}
                Write(c->Base+4096+(admin?1u:3u)*c->Stride,head);return ok;
            }
            if((Read(c->Base+28)&2)!=0)break;Threads.Sleep(1);
        }while(Now()<end);
        c->Active=0;Write(c->Base+20,0);Wait(c->Base+28,1,0,1000);Log("NVME_COMMAND_FAILED_QUARANTINED");return false;
    }
    static bool InitializeNvme(DiskController* c)
    {
        if(!MapBar(c,16,false))return false;
        uint cap=Read(c->Base),high=Read(c->Base+4);if((cap&65535)<15||((high>>16)&15)!=0||(high&15)>8)return false;
        Write(c->Base+20,0);if(!Wait(c->Base+28,1,0,5000))return false;
        if(!MapBar(c,16,true))return false;EnablePci(c);c->Stride=4u<<(int)(high&15);
        if(4096+4*c->Stride>c->Bytes)return false;
        c->AdminSq=Allocate();c->AdminCq=Allocate();c->IoSq=Allocate();c->IoCq=Allocate();c->Data=Allocate(32);c->Prps=Allocate();
        if(c->AdminSq==0||c->AdminCq==0||c->IoSq==0||c->IoCq==0||c->Data==0||c->Prps==0)return false;
        for(uint i=0;i<31;i++)((ulong*)At(c->Prps))[i]=c->Data+(i+1)*4096;
        Write(c->Base+12,0xffffffff);Write(c->Base+36,15u|15u<<16);Write64(c->Base+40,c->AdminSq);Write64(c->Base+48,c->AdminCq);
        Write(c->Base+20,1u|6u<<16|4u<<20);if(!Wait(c->Base+28,1,1,5000))return false;
        c->Active=1;c->AdminPhase=c->IoPhase=1;
        if(!NvmeCommand(c,true,6,0,c->Data,1))return false;
        byte* data=At(c->Data);uint namespaces=*(uint*)(data+516),mdts=data[77];c->MaxTransfer=mdts==0||mdts>=5?131072UL:4096UL<<(int)mdts;
        if(!NvmeCommand(c,true,5,0,c->IoCq,1u|15u<<16,1)||!NvmeCommand(c,true,1,0,c->IoSq,1u|15u<<16,1u|1u<<16))return false;
        if(!NvmeCommand(c,true,6,0,c->Data,2))return false;
        uint* active=stackalloc uint[32];uint found=0;
        for(uint i=0;i<1024&&found<32;i++){uint id=((uint*)At(c->Data))[i];if(id==0)break;active[found++]=id;}
        if(found==0&&namespaces!=0){active[0]=1;found=1;}
        for(uint i=0;i<found&&_count<Limit;i++)
        {
            uint ns=active[i];if(!NvmeCommand(c,true,6,ns,c->Data,0))continue;data=At(c->Data);
            ulong blocks=*(ulong*)data;uint format=(uint)data[26]&15;if(format>data[25]||*(ushort*)(data+128+format*4)!=0||(data[29]&7)!=0)continue;
            uint shift=data[130+format*4];if((shift!=9&&shift!=12)||blocks==0||blocks>(~0UL>>(int)(shift-9)))continue;
            HardwareDisk* disk=&_disks[_count];*disk=new HardwareDisk {Controller=c,Namespace=ns,Blocks=blocks<<(int)(shift-9),SectorShift=shift,Active=1,CanWrite=(data[99]&1)==0?1u:0};Publish(disk);
        }
        return true;
    }
    static bool Transfer(HardwareDisk* d,uint operation,ulong lba,uint sectors)
    {
        if(d->Active==0||d->Controller->Active==0)return false;
        bool write=operation==2;DiskController* c=d->Controller;
        if(operation==3)return c->Kind==1?AhciCommand(d,d->FlushCommand,0,0,0,false):NvmeCommand(c,false,0,d->Namespace,0,0);
        if(sectors==0||sectors>128||lba>=d->Blocks||sectors>d->Blocks-lba)return false;
        if(write&&(d->CanWrite==0||d->AllowedCount==0||lba<d->AllowedStart||lba-d->AllowedStart>=d->AllowedCount||sectors>d->AllowedCount-(lba-d->AllowedStart)))return false;
        uint unit=1u<<(int)(d->SectorShift-9),skip=(uint)(lba%unit)*512;
        uint bytes=sectors*512,blocks=(skip+bytes+(unit*512-1))/(unit*512);ulong physicalLba=lba/unit;
        if(c->Kind==1)
        {
            if(write)for(uint i=0;i<bytes;i++)At(c->Data)[i]=At(_mailbox)[i];
            if(!AhciCommand(d,write?0x35u:0x25u,lba,sectors,bytes,write))return false;
        }
        else
        {
            if((ulong)blocks*unit*512>c->MaxTransfer)return false;
            // A 4 KiB namespace uses read/modify/write for unaligned 512-byte
            // filesystem metadata. Serialize the entire operation on this worker.
            if(!write||skip!=0||bytes%(unit*512)!=0)
                if(!NvmeCommand(c,false,2,d->Namespace,c->Data,(uint)physicalLba,(uint)(physicalLba>>32),blocks-1,(ulong)blocks*unit*512))return false;
            if(write)
            {
                for(uint i=0;i<bytes;i++)At(c->Data)[skip+i]=At(_mailbox)[i];
                if(!NvmeCommand(c,false,1,d->Namespace,c->Data,(uint)physicalLba,(uint)(physicalLba>>32),blocks-1,(ulong)blocks*unit*512))return false;
            }
        }
        if(write)d->Writes+=sectors;else {for(uint i=0;i<bytes;i++)At(_mailbox)[i]=At(c->Data)[skip+i];d->Reads+=sectors;}
        return true;
    }
    static void ProcessRequest()
    {
        Native.MemoryBarrier();if(_requestState!=1)return;_requestState=2;
        bool ok=_disk<_count&&Transfer(&_disks[_disk],_operation,_lba,_sectors);
        if(ok&&_operation==3)_disks[_disk].Flushes++;
        Native.MemoryBarrier();_requestState=_requestState==6?0u:ok?3u:4u;
    }
    static int Request(uint disk,uint operation,ulong lba,uint sectors,byte* buffer)
    {
        if(disk>=_count||_state!=2||operation!=3&&(buffer==null||sectors==0||sectors>128))return 0;
        ulong irq=Native.IrqSave();if(_requestState!=0){Native.IrqRestore(irq);return 0;}
        if(operation==2)for(uint i=0;i<sectors*512;i++)At(_mailbox)[i]=buffer[i];
        _disk=disk;_operation=operation;_lba=lba;_sectors=sectors;_requestState=1;Native.MemoryBarrier();Native.IrqRestore(irq);
        ulong end=Now()+15000;
        for(;;)
        {
            Native.MemoryBarrier();uint state=_requestState;
            if(state==3||state==4){if(state==3&&operation==1)for(uint i=0;i<sectors*512;i++)buffer[i]=At(_mailbox)[i];_requestState=0;Native.MemoryBarrier();return state==3?1:0;}
            if(Now()>=end){irq=Native.IrqSave();uint current=_requestState;if(current==1)_requestState=0;else if(current==2)_requestState=6;Native.IrqRestore(irq);if(current==3||current==4)continue;return 0;}
            Threads.Sleep(1);
        }
    }
    [UnmanagedCallersOnly(EntryPoint="xfe_disk_info")]
    public static ulong Info(uint index,uint field)
    {
        Native.MemoryBarrier();if(field==0)return _state;if(field==1)return _count;if(index>=_count)return 0;
        HardwareDisk* d=&_disks[index];return field switch {2=>d->Active*d->Controller->Active,3=>d->Blocks,4=>d->Controller->Kind,5=>d->SectorShift,6=>d->CanWrite,7=>d->Reads,8=>d->Writes,9=>d->Flushes,10=>d->Controller->Kind==1?65536:d->Controller->MaxTransfer-4096,_=>0};
    }
    [UnmanagedCallersOnly(EntryPoint="xfe_disk_read")]
    public static int ReadSectors(uint index,ulong lba,uint count,byte* buffer)=>Request(index,1,lba,count,buffer);
    [UnmanagedCallersOnly(EntryPoint="xfe_disk_write")]
    public static int WriteSectors(uint index,ulong lba,uint count,byte* buffer)=>Request(index,2,lba,count,buffer);
    [UnmanagedCallersOnly(EntryPoint="xfe_disk_flush")]
    public static int Flush(uint index)=>Request(index,3,0,0,null);
    [UnmanagedCallersOnly(EntryPoint="xfe_disk_allow_partition")]
    public static int AllowPartition(uint index,ulong start,ulong count)
    {
        if(index>=_count||start<1||count==0||start>=_disks[index].Blocks||count>_disks[index].Blocks-start||_disks[index].CanWrite==0)return 0;
        uint unit=1u<<(int)(_disks[index].SectorShift-9);
        if(start%unit!=0||count%unit!=0)return 0;
        // One mounted partition per disk request carries its exact extent;
        // managed calls serialize switching this window with all filesystem I/O.
        ulong irq=Native.IrqSave();if(_requestState!=0){Native.IrqRestore(irq);return 0;}
        _disks[index].AllowedStart=start;_disks[index].AllowedCount=count;Native.MemoryBarrier();Native.IrqRestore(irq);return 1;
    }
}