// AHCI driver state and register definitions shared across controller and port code.

#pragma once

#include <stdint.h>
#include <stdbool.h>
#include <string.h>
#include <ntlite/ddi/io_driver.h>
#include <ntlite/ddi/pci.h>
#include <kernel/mm_phys.h>
#include <kernel/mm_pfn.h>
#include <kernel/mm_virt.h>
#include <kernel/ke_debug.h>
#include <kernel/ke_event.h>
#include <kernel/ke_semaphore.h>
#include <kernel/ke_thread.h>
#include <kernel/ke_irq.h>
#include <kernel/ke_irp.h>
#include <kernel/ke_dpc.h>

typedef enum {
    FIS_TYPE_REG_H2D   = 0x27,
    FIS_TYPE_REG_D2H   = 0x34,
    FIS_TYPE_DMA_ACT   = 0x39,
    FIS_TYPE_DMA_SETUP = 0x41,
    FIS_TYPE_DATA      = 0x46,
    FIS_TYPE_BIST      = 0x58,
    FIS_TYPE_PIO_SETUP = 0x5F,
    FIS_TYPE_DEV_BITS  = 0xA1,
} FIS_TYPE;

#pragma pack(push, 1)
// Register FIS - Host to Device (H2D).
typedef struct {
    uint8_t  fis_type;          // FIS_TYPE_REG_H2D
    uint8_t  pmport     : 4;
    uint8_t  rsv0       : 3;
    uint8_t  c          : 1;    // 1 = Command, 0 = Control
    uint8_t  command;
    uint8_t  featurel;
    // DW1.
    uint8_t  lba0, lba1, lba2;
    uint8_t  device;
    // DW2.
    uint8_t  lba3, lba4, lba5;
    uint8_t  featureh;
    // DW3.
    uint8_t  countl, counth;
    uint8_t  icc;
    uint8_t  control;
    // DW4.
    uint8_t  rsv1[4];
} FIS_REG_H2D;

// Register FIS - Device to Host (D2H).
typedef struct {
    uint8_t  fis_type;
    uint8_t  pmport : 4;
    uint8_t  rsv0   : 2;
    uint8_t  i      : 1;
    uint8_t  rsv1   : 1;
    uint8_t  status;
    uint8_t  error;
    uint8_t  lba0, lba1, lba2;
    uint8_t  device;
    uint8_t  lba3, lba4, lba5;
    uint8_t  rsv2;
    uint8_t  countl, counth;
    uint8_t  rsv3[2];
    uint8_t  rsv4[4];
} FIS_REG_D2H;

// DMA Setup FIS.
typedef struct {
    uint8_t  fis_type;
    uint8_t  pmport : 4;
    uint8_t  rsv0   : 1;
    uint8_t  d      : 1;
    uint8_t  i      : 1;
    uint8_t  a      : 1;
    uint8_t  rsved[2];
    uint64_t DMAbufferID;
    uint32_t rsvd;
    uint32_t DMAbufOffset;
    uint32_t TransferCount;
    uint32_t resvd;
} FIS_DMA_SETUP;

// PIO Setup FIS.
typedef struct {
    uint8_t  fis_type;
    uint8_t  pmport : 4;
    uint8_t  rsv0   : 1;
    uint8_t  d      : 1;
    uint8_t  i      : 1;
    uint8_t  rsv1   : 1;
    uint8_t  status;
    uint8_t  error;
    uint8_t  lba0, lba1, lba2;
    uint8_t  device;
    uint8_t  lba3, lba4, lba5;
    uint8_t  rsv2;
    uint8_t  countl, counth;
    uint8_t  rsv3;
    uint8_t  e_status;
    uint16_t tc;
    uint8_t  rsv4[2];
} FIS_PIO_SETUP;

// Set Device Bits FIS.
typedef struct {
    uint8_t  fis_type;
    uint8_t  pmport : 4;
    uint8_t  rsv0   : 2;
    uint8_t  i      : 1;
    uint8_t  n      : 1;
    uint8_t  statusl : 4;
    uint8_t  rsv1    : 4;
    uint8_t  statush : 4;
    uint8_t  rsv2    : 4;
    uint8_t  error;
    uint32_t protocol;
} FIS_DEV_BITS;

// HBA structures (volatile - MMIO).

// Received FIS structure in host memory.
typedef volatile struct {
    FIS_DMA_SETUP  dsfis;       // 0x00 DMA Setup FIS
    uint8_t        pad0[4];
    FIS_PIO_SETUP  psfis;       // 0x20 PIO Setup FIS
    uint8_t        pad1[12];
    FIS_REG_D2H    rfis;        // 0x40 Register D2H FIS
    uint8_t        pad2[4];
    FIS_DEV_BITS   sdbfis;      // 0x58 Set Device Bits FIS
    uint8_t        ufis[64];    // 0x60 Unknown FIS
    uint8_t        rsv[0x100-0xA0];
} HBA_FIS;

// Command List header (one entry = one slot).
typedef struct {
    // DW0.
    uint8_t  cfl    : 5;        // Command FIS length in DWORDs (2..16)
    uint8_t  a      : 1;        // ATAPI
    uint8_t  w      : 1;        // Write (H2D=1, D2H=0)
    uint8_t  p      : 1;        // Prefetchable
    uint8_t  r      : 1;        // Reset
    uint8_t  b      : 1;        // BIST
    uint8_t  c      : 1;        // Clear busy on R_OK
    uint8_t  rsv0   : 1;
    uint8_t  pmp    : 4;        // Port multiplier port
    uint16_t prdtl;             // PRDT length in entries
    // DW1.
    volatile uint32_t prdbc;    // PRDT byte count transferred
    // DW2,3.
    uint32_t ctba;              // Command table base address (128-byte aligned)
    uint32_t ctbau;             // Upper 32 bits
    // DW4-7.
    uint32_t rsv1[4];
} HBA_CMD_HEADER;

// Physical Region Descriptor Table entry.
typedef struct {
    uint32_t dba;               // Data base address
    uint32_t dbau;              // Upper 32 bits
    uint32_t rsv0;
    uint32_t dbc   : 22;        // Byte count (0-based, so actual = dbc+1)
    uint32_t rsv1  : 9;
    uint32_t i     : 1;         // Interrupt on completion
} HBA_PRDT_ENTRY;

// Command Table.
typedef struct {
    uint8_t        cfis[64];    // Command FIS (up to 64 bytes)
    uint8_t        acmd[16];    // ATAPI command (12 or 16 bytes)
    uint8_t        rsv[48];     // Reserved
    HBA_PRDT_ENTRY prdt_entry[1]; // PRDT entries (variable length)
} HBA_CMD_TBL;
#pragma pack(pop)

// Port registers.
typedef volatile struct {
    uint32_t clb;               // 0x00 Command list base (1K-aligned)
    uint32_t clbu;              // 0x04 Command list base upper
    uint32_t fb;                // 0x08 FIS base (256-byte aligned)
    uint32_t fbu;               // 0x0C FIS base upper
    uint32_t is;                // 0x10 Interrupt status
    uint32_t ie;                // 0x14 Interrupt enable
    uint32_t cmd;               // 0x18 Command and status
    uint32_t rsv0;              // 0x1C Reserved
    uint32_t tfd;               // 0x20 Task file data
    uint32_t sig;               // 0x24 Signature
    uint32_t ssts;              // 0x28 SATA status
    uint32_t sctl;              // 0x2C SATA control
    uint32_t serr;              // 0x30 SATA error
    uint32_t sact;              // 0x34 SATA active
    uint32_t ci;                // 0x38 Command issue
    uint32_t sntf;              // 0x3C SATA notification
    uint32_t fbs;               // 0x40 FIS-based switch control
    uint32_t rsv1[11];
    uint32_t vendor[4];
} HBA_PORT;

// Generic Host Control registers (starts at ABAR).
typedef volatile struct {
    uint32_t cap;               // 0x00 Host capability
    uint32_t ghc;               // 0x04 Global host control
    uint32_t is;                // 0x08 Interrupt status
    uint32_t pi;                // 0x0C Ports implemented
    uint32_t vs;                // 0x10 Version
    uint32_t ccc_ctl;           // 0x14 CCC control
    uint32_t ccc_pts;           // 0x18 CCC ports
    uint32_t em_loc;            // 0x1C Enclosure mgmt location
    uint32_t em_ctl;            // 0x20 Enclosure mgmt control
    uint32_t cap2;              // 0x24 Extended capabilities
    uint32_t bohc;              // 0x28 BIOS/OS handoff control
    uint8_t  rsv[0xA0-0x2C];
    uint8_t  vendor[0x100-0xA0];
    HBA_PORT ports[32];         // 0x100 Port control registers
} HBA_MEM;

// Constants.

// PCI/AHCI binding.  PCI enumeration, BAR decoding, interrupt line
// discovery, and command-register policy are owned by pci.sys.
#define AHCI_PCI_BAR_ABAR     5u
#define AHCI_CONTROLLER_MAGIC  0x41484349u // 'AHCI'
#define AHCI_DRIVE_MAGIC       0x41484452u // 'AHDR'
#define MAX_AHCI_CONTROLLERS   8

// GHC register bits.
#define GHC_AE                (1u << 31)  // AHCI Enable
#define GHC_IE                (1u << 1)   // Interrupt Enable
#define GHC_HR                (1u << 0)   // HBA Reset

// BOHC register bits.
#define BOHC_BOS              (1u << 0)   // BIOS-owned semaphore
#define BOHC_OOS              (1u << 1)   // OS-owned semaphore
#define BOHC_BB               (1u << 4)   // BIOS Busy

// CAP2 bits.
#define CAP2_BOH              (1u << 0)   // BIOS/OS handoff capable

// Port CMD register bits.
#define PORT_CMD_ST           (1u << 0)   // Start
#define PORT_CMD_FRE          (1u << 4)   // FIS Receive Enable
#define PORT_CMD_FR           (1u << 14)  // FIS Receive Running
#define PORT_CMD_CR           (1u << 15)  // Command List Running
#define PORT_CMD_POD          (1u << 2)   // Power On Device
#define PORT_CMD_SUD          (1u << 1)   // Spin-Up Device
#define PORT_CMD_ICC_ACTIVE   (0x1u << 28)

// Port IS register bits (interrupt status).
#define PORT_IS_TFES          (1u << 30)  // Task File Error Status
#define PORT_IS_HBFS          (1u << 29)  // Host Bus Fatal Error
#define PORT_IS_HBDS          (1u << 28)  // Host Bus Data Error
#define PORT_IS_IFS           (1u << 27)  // Interface Fatal Error

// Port SSTS register.
#define PORT_DET_PRESENT      0x3
#define PORT_IPM_ACTIVE       0x1

// Port signatures.
#define SATA_SIG_ATA          0x00000101
#define SATA_SIG_ATAPI        0xEB140101
#define SATA_SIG_SEMB         0xC33C0101
#define SATA_SIG_PM           0x96690101

// ATA commands.
#define ATA_CMD_READ_DMA_EX   0x25
#define ATA_CMD_WRITE_DMA_EX  0x35
#define ATA_CMD_IDENTIFY      0xEC
#define ATA_CMD_FLUSH_EXT     0xEA

// ATA status bits.
#define ATA_SR_BUSY           0x80
#define ATA_SR_DRQ            0x08
#define ATA_SR_ERR            0x01

// IDENTIFY word offsets.
#define ATA_IDENT_W83         83   // bit 10 = 48-bit LBA supported
#define ATA_IDENT_W100        100  // 48-bit LBA sector count (4 words)
#define ATA_IDENT_W60         60   // 28-bit LBA sector count (2 words)

// Layout constants (per port, single contiguous physical allocation)
// Offset 0:    Command List   1024 bytes (32 × 32-byte headers)
// Offset 1024: Received FIS    256 bytes
// Offset 2048: Command Tables  32 × 256 bytes = 8192 bytes
// Total:                       10240 bytes → 3 pages (12288 bytes).
#define PORT_PAGES_STRUCT     3
#define PORT_OFF_CMDLIST      0
#define PORT_OFF_FIS          1024
#define PORT_OFF_CMDTBLS      2048
#define PORT_CMDTBL_SIZE      256   // 64 CFIS + 16 ACMD + 48 rsv + 8×16 PRDT
#define PORT_PRDT_COUNT       8     // PRDT entries per command table

// Bounce buffer: 128 sectors × 512 bytes = 64 KB = 16 pages.
#define BOUNCE_SECTORS        128
#define BOUNCE_BYTES          (BOUNCE_SECTORS * 512)
#define BOUNCE_PAGES          (BOUNCE_BYTES / 4096)

// Polling timeout (spin iterations).
#define POLL_TIMEOUT          1000000

#define MAX_AHCI_DRIVES       8
#define NUM_CMD_SLOTS         32

// Map MMIO physical range into the physical window as uncacheable.
// Must be called before any pointer dereference via PhysToVirt().
// AHCI MMIO layout: 0x100 GHC + 32*0x80 ports = 0x1100 bytes → 2 pages minimum.
// We map 8 pages (32KB) without probing the BAR size.
#define AHCI_MMIO_PAGES 8

typedef struct AHCI_CONTROLLER AHCI_CONTROLLER;

// Per-drive context.
typedef struct {
    uint32_t        magic;          // AHCI_DRIVE_MAGIC
    AHCI_CONTROLLER* controller;    // owning AHCI controller FDO extension
    HBA_PORT       *port;           // MMIO port registers (virtual)
    uint64_t        sectorCount;    // Total LBA sectors on device
    uint8_t         portNo;         // Physical port index (for ISR dispatch)

    // Structures block (command list + FIS + cmd tables).
    uint64_t        structPhys;     // Physical base
    void*           structVirt;     // Virtual base

    // Bounce buffer for DMA data.
    uint64_t        bouncePhys;
    void*           bounceVirt;

    // Signaled by the completion DPC when the port raises an interrupt.
    // Auto-reset: one Wait consumes one completion.
    EVENT          *completionEvent;

    // Last IS snapshot captured by the ISR (for error checking post-wait).
    volatile uint32_t lastIs;

    // Bottom-half used to wake the worker outside hard IRQ/DIRQL context.
    KDPC            completionDpc;

    // IRP queue - async dispatch
    // AhciWorker blocks here waiting for work. AHCI dispatch enqueues an IRP
    // and signals the semaphore to wake the worker.
    IRP_QUEUE       irpQueue;
    IO_REMOVE_LOCK  removeLock;
    bool            surpriseRemoved;
    bool            removed;
    THREAD*         workerThread;
    DEVICE_OBJECT*  deviceObject;

    // Multi-slot NCQ needs per-slot ownership state in addition to the serialized queue:
    //   uint32_t         slotMask;              // bitmask of in-flight slots
    //   IRP*             slotIrp[NUM_CMD_SLOTS]; // slot → pending IRP
    //   uint64_t         slotBouncePhys[NUM_CMD_SLOTS];
    //   void*            slotBounceVirt[NUM_CMD_SLOTS];
    //   volatile uint32_t prevSact;             // previous port->sact snapshot.
} AHCI_DRIVE;

struct AHCI_CONTROLLER {
    uint32_t       magic;           // AHCI_CONTROLLER_MAGIC
    DEVICE_OBJECT* fdo;
    DEVICE_OBJECT* pdo;
    DEVICE_OBJECT* lowerDevice;

    PCI_DEVICE_INFO pciInfo;
    PCI_BAR         abar;
    HBA_MEM*        hba;
    uint64_t        abarLength;
    uint8_t         irq;
    bool            usingMsi;
    bool            pollingMode;

    AHCI_DRIVE      drives[MAX_AHCI_DRIVES];
    int             driveCount;

    bool            started;
    bool            surpriseRemoved;
    bool            removed;
};

extern const HAL_API *g_hal;
extern DRIVER_OBJECT g_driver;

KSTATUS ahci_map_mmio(uint64_t physBase, uint64_t length);
void port_stop(HBA_PORT *port);
void port_start(HBA_PORT *port);
KSTATUS hba_reset(HBA_MEM *hba);
bool port_link_up(HBA_PORT *port);
bool port_is_sata_disk(HBA_PORT *port);
void ahci_delay_ms(uint32_t ms);
KSTATUS port_comreset(AHCI_DRIVE *drv);