#include "alloc.h"
#include "ata.h"
#include "io.h"
#include "irq.h"
#include "kernel.h"
#include "klibc.h"
#include "list.h"
#include "thread.h"

#define PART_TABLE_OFFSET 0x1BE

static struct ata_partition *partitions;
static int nextPartitionId = 0;

// from
// https://docs.microsoft.com/en-us/windows-hardware/drivers/ddi/ata/ns-ata-_identify_device_data
struct _IDENTIFY_DEVICE_DATA {
	struct {
		uint16_t Reserved1 : 1;
		uint16_t Retired3 : 1;
		uint16_t ResponseIncomplete : 1;
		uint16_t Retired2 : 3;
		uint16_t FixedDevice : 1;
		uint16_t RemovableMedia : 1;
		uint16_t Retired1 : 7;
		uint16_t DeviceType : 1;
	} GeneralConfiguration;
	uint16_t NumCylinders;
	uint16_t SpecificConfiguration;
	uint16_t NumHeads;
	uint16_t Retired1[2];
	uint16_t NumSectorsPerTrack;
	uint16_t VendorUnique1[3];
	uint8_t SerialNumber[20];
	uint16_t Retired2[2];
	uint16_t Obsolete1;
	uint8_t FirmwareRevision[8];
	uint8_t ModelNumber[40];
	uint8_t MaximumBlockTransfer;
	uint8_t VendorUnique2;
	struct {
		uint16_t FeatureSupported : 1;
		uint16_t Reserved : 15;
	} TrustedComputing;
	struct {
		uint8_t CurrentLongPhysicalSectorAlignment : 2;
		uint8_t ReservedByte49 : 6;
		uint8_t DmaSupported : 1;
		uint8_t LbaSupported : 1;
		uint8_t IordyDisable : 1;
		uint8_t IordySupported : 1;
		uint8_t Reserved1 : 1;
		uint8_t StandybyTimerSupport : 1;
		uint8_t Reserved2 : 2;
		uint16_t ReservedWord50;
	} Capabilities;
	uint16_t ObsoleteWords51[2];
	uint16_t TranslationFieldsValid : 3;
	uint16_t Reserved3 : 5;
	uint16_t FreeFallControlSensitivity : 8;
	uint16_t NumberOfCurrentCylinders;
	uint16_t NumberOfCurrentHeads;
	uint16_t CurrentSectorsPerTrack;
	uint32_t CurrentSectorCapacity;
	uint8_t CurrentMultiSectorSetting;
	uint8_t MultiSectorSettingValid : 1;
	uint8_t ReservedByte59 : 3;
	uint8_t SanitizeFeatureSupported : 1;
	uint8_t CryptoScrambleExtCommandSupported : 1;
	uint8_t OverwriteExtCommandSupported : 1;
	uint8_t BlockEraseExtCommandSupported : 1;
	uint32_t UserAddressableSectors;
	uint16_t ObsoleteWord62;
	uint16_t MultiWordDMASupport : 8;
	uint16_t MultiWordDMAActive : 8;
	uint16_t AdvancedPIOModes : 8;
	uint16_t ReservedByte64 : 8;
	uint16_t MinimumMWXferCycleTime;
	uint16_t RecommendedMWXferCycleTime;
	uint16_t MinimumPIOCycleTime;
	uint16_t MinimumPIOCycleTimeIORDY;
	struct {
		uint16_t ZonedCapabilities : 2;
		uint16_t NonVolatileWriteCache : 1;
		uint16_t ExtendedUserAddressableSectorsSupported : 1;
		uint16_t DeviceEncryptsAllUserData : 1;
		uint16_t ReadZeroAfterTrimSupported : 1;
		uint16_t Optional28BitCommandsSupported : 1;
		uint16_t IEEE1667 : 1;
		uint16_t DownloadMicrocodeDmaSupported : 1;
		uint16_t SetMaxSetPasswordUnlockDmaSupported : 1;
		uint16_t WriteBufferDmaSupported : 1;
		uint16_t ReadBufferDmaSupported : 1;
		uint16_t DeviceConfigIdentifySetDmaSupported : 1;
		uint16_t LPSAERCSupported : 1;
		uint16_t DeterministicReadAfterTrimSupported : 1;
		uint16_t CFastSpecSupported : 1;
	} AdditionalSupported;
	uint16_t ReservedWords70[5];
	uint16_t QueueDepth : 5;
	uint16_t ReservedWord75 : 11;
	struct {
		uint16_t Reserved0 : 1;
		uint16_t SataGen1 : 1;
		uint16_t SataGen2 : 1;
		uint16_t SataGen3 : 1;
		uint16_t Reserved1 : 4;
		uint16_t NCQ : 1;
		uint16_t HIPM : 1;
		uint16_t PhyEvents : 1;
		uint16_t NcqUnload : 1;
		uint16_t NcqPriority : 1;
		uint16_t HostAutoPS : 1;
		uint16_t DeviceAutoPS : 1;
		uint16_t ReadLogDMA : 1;
		uint16_t Reserved2 : 1;
		uint16_t CurrentSpeed : 3;
		uint16_t NcqStreaming : 1;
		uint16_t NcqQueueMgmt : 1;
		uint16_t NcqReceiveSend : 1;
		uint16_t DEVSLPtoReducedPwrState : 1;
		uint16_t Reserved3 : 8;
	} SerialAtaCapabilities;
	struct {
		uint16_t Reserved0 : 1;
		uint16_t NonZeroOffsets : 1;
		uint16_t DmaSetupAutoActivate : 1;
		uint16_t DIPM : 1;
		uint16_t InOrderData : 1;
		uint16_t HardwareFeatureControl : 1;
		uint16_t SoftwareSettingsPreservation : 1;
		uint16_t NCQAutosense : 1;
		uint16_t DEVSLP : 1;
		uint16_t HybridInformation : 1;
		uint16_t Reserved1 : 6;
	} SerialAtaFeaturesSupported;
	struct {
		uint16_t Reserved0 : 1;
		uint16_t NonZeroOffsets : 1;
		uint16_t DmaSetupAutoActivate : 1;
		uint16_t DIPM : 1;
		uint16_t InOrderData : 1;
		uint16_t HardwareFeatureControl : 1;
		uint16_t SoftwareSettingsPreservation : 1;
		uint16_t DeviceAutoPS : 1;
		uint16_t DEVSLP : 1;
		uint16_t HybridInformation : 1;
		uint16_t Reserved1 : 6;
	} SerialAtaFeaturesEnabled;
	uint16_t MajorRevision;
	uint16_t MinorRevision;
	struct {
		uint16_t SmartCommands : 1;
		uint16_t SecurityMode : 1;
		uint16_t RemovableMediaFeature : 1;
		uint16_t PowerManagement : 1;
		uint16_t Reserved1 : 1;
		uint16_t WriteCache : 1;
		uint16_t LookAhead : 1;
		uint16_t ReleaseInterrupt : 1;
		uint16_t ServiceInterrupt : 1;
		uint16_t DeviceReset : 1;
		uint16_t HostProtectedArea : 1;
		uint16_t Obsolete1 : 1;
		uint16_t WriteBuffer : 1;
		uint16_t ReadBuffer : 1;
		uint16_t Nop : 1;
		uint16_t Obsolete2 : 1;
		uint16_t DownloadMicrocode : 1;
		uint16_t DmaQueued : 1;
		uint16_t Cfa : 1;
		uint16_t AdvancedPm : 1;
		uint16_t Msn : 1;
		uint16_t PowerUpInStandby : 1;
		uint16_t ManualPowerUp : 1;
		uint16_t Reserved2 : 1;
		uint16_t SetMax : 1;
		uint16_t Acoustics : 1;
		uint16_t BigLba : 1;
		uint16_t DeviceConfigOverlay : 1;
		uint16_t FlushCache : 1;
		uint16_t FlushCacheExt : 1;
		uint16_t WordValid83 : 2;
		uint16_t SmartErrorLog : 1;
		uint16_t SmartSelfTest : 1;
		uint16_t MediaSerialNumber : 1;
		uint16_t MediaCardPassThrough : 1;
		uint16_t StreamingFeature : 1;
		uint16_t GpLogging : 1;
		uint16_t WriteFua : 1;
		uint16_t WriteQueuedFua : 1;
		uint16_t WWN64Bit : 1;
		uint16_t URGReadStream : 1;
		uint16_t URGWriteStream : 1;
		uint16_t ReservedForTechReport : 2;
		uint16_t IdleWithUnloadFeature : 1;
		uint16_t WordValid : 2;
	} CommandSetSupport;
	struct {
		uint16_t SmartCommands : 1;
		uint16_t SecurityMode : 1;
		uint16_t RemovableMediaFeature : 1;
		uint16_t PowerManagement : 1;
		uint16_t Reserved1 : 1;
		uint16_t WriteCache : 1;
		uint16_t LookAhead : 1;
		uint16_t ReleaseInterrupt : 1;
		uint16_t ServiceInterrupt : 1;
		uint16_t DeviceReset : 1;
		uint16_t HostProtectedArea : 1;
		uint16_t Obsolete1 : 1;
		uint16_t WriteBuffer : 1;
		uint16_t ReadBuffer : 1;
		uint16_t Nop : 1;
		uint16_t Obsolete2 : 1;
		uint16_t DownloadMicrocode : 1;
		uint16_t DmaQueued : 1;
		uint16_t Cfa : 1;
		uint16_t AdvancedPm : 1;
		uint16_t Msn : 1;
		uint16_t PowerUpInStandby : 1;
		uint16_t ManualPowerUp : 1;
		uint16_t Reserved2 : 1;
		uint16_t SetMax : 1;
		uint16_t Acoustics : 1;
		uint16_t BigLba : 1;
		uint16_t DeviceConfigOverlay : 1;
		uint16_t FlushCache : 1;
		uint16_t FlushCacheExt : 1;
		uint16_t Resrved3 : 1;
		uint16_t Words119_120Valid : 1;
		uint16_t SmartErrorLog : 1;
		uint16_t SmartSelfTest : 1;
		uint16_t MediaSerialNumber : 1;
		uint16_t MediaCardPassThrough : 1;
		uint16_t StreamingFeature : 1;
		uint16_t GpLogging : 1;
		uint16_t WriteFua : 1;
		uint16_t WriteQueuedFua : 1;
		uint16_t WWN64Bit : 1;
		uint16_t URGReadStream : 1;
		uint16_t URGWriteStream : 1;
		uint16_t ReservedForTechReport : 2;
		uint16_t IdleWithUnloadFeature : 1;
		uint16_t Reserved4 : 2;
	} CommandSetActive;
	uint16_t UltraDMASupport : 8;
	uint16_t UltraDMAActive : 8;
	struct {
		uint16_t TimeRequired : 15;
		uint16_t ExtendedTimeReported : 1;
	} NormalSecurityEraseUnit;
	struct {
		uint16_t TimeRequired : 15;
		uint16_t ExtendedTimeReported : 1;
	} EnhancedSecurityEraseUnit;
	uint16_t CurrentAPMLevel : 8;
	uint16_t ReservedWord91 : 8;
	uint16_t MasterPasswordID;
	uint16_t HardwareResetResult;
	uint16_t CurrentAcousticValue : 8;
	uint16_t RecommendedAcousticValue : 8;
	uint16_t StreamMinRequestSize;
	uint16_t StreamingTransferTimeDMA;
	uint16_t StreamingAccessLatencyDMAPIO;
	uint32_t StreamingPerfGranularity;
	uint32_t Max48BitLBA[2];
	uint16_t StreamingTransferTime;
	uint16_t DsmCap;
	struct {
		uint16_t LogicalSectorsPerPhysicalSector : 4;
		uint16_t Reserved0 : 8;
		uint16_t LogicalSectorLongerThan256Words : 1;
		uint16_t MultipleLogicalSectorsPerPhysicalSector : 1;
		uint16_t Reserved1 : 2;
	} PhysicalLogicalSectorSize;
	uint16_t InterSeekDelay;
	uint16_t WorldWideName[4];
	uint16_t ReservedForWorldWideName128[4];
	uint16_t ReservedForTlcTechnicalReport;
	uint16_t WordsPerLogicalSector[2];
	struct {
		uint16_t ReservedForDrqTechnicalReport : 1;
		uint16_t WriteReadVerify : 1;
		uint16_t WriteUncorrectableExt : 1;
		uint16_t ReadWriteLogDmaExt : 1;
		uint16_t DownloadMicrocodeMode3 : 1;
		uint16_t FreefallControl : 1;
		uint16_t SenseDataReporting : 1;
		uint16_t ExtendedPowerConditions : 1;
		uint16_t Reserved0 : 6;
		uint16_t WordValid : 2;
	} CommandSetSupportExt;
	struct {
		uint16_t ReservedForDrqTechnicalReport : 1;
		uint16_t WriteReadVerify : 1;
		uint16_t WriteUncorrectableExt : 1;
		uint16_t ReadWriteLogDmaExt : 1;
		uint16_t DownloadMicrocodeMode3 : 1;
		uint16_t FreefallControl : 1;
		uint16_t SenseDataReporting : 1;
		uint16_t ExtendedPowerConditions : 1;
		uint16_t Reserved0 : 6;
		uint16_t Reserved1 : 2;
	} CommandSetActiveExt;
	uint16_t ReservedForExpandedSupportandActive[6];
	uint16_t MsnSupport : 2;
	uint16_t ReservedWord127 : 14;
	struct {
		uint16_t SecuritySupported : 1;
		uint16_t SecurityEnabled : 1;
		uint16_t SecurityLocked : 1;
		uint16_t SecurityFrozen : 1;
		uint16_t SecurityCountExpired : 1;
		uint16_t EnhancedSecurityEraseSupported : 1;
		uint16_t Reserved0 : 2;
		uint16_t SecurityLevel : 1;
		uint16_t Reserved1 : 7;
	} SecurityStatus;
	uint16_t ReservedWord129[31];
	struct {
		uint16_t MaximumCurrentInMA : 12;
		uint16_t CfaPowerMode1Disabled : 1;
		uint16_t CfaPowerMode1Required : 1;
		uint16_t Reserved0 : 1;
		uint16_t Word160Supported : 1;
	} CfaPowerMode1;
	uint16_t ReservedForCfaWord161[7];
	uint16_t NominalFormFactor : 4;
	uint16_t ReservedWord168 : 12;
	struct {
		uint16_t SupportsTrim : 1;
		uint16_t Reserved0 : 15;
	} DataSetManagementFeature;
	uint16_t AdditionalProductID[4];
	uint16_t ReservedForCfaWord174[2];
	uint16_t CurrentMediaSerialNumber[30];
	struct {
		uint16_t Supported : 1;
		uint16_t Reserved0 : 1;
		uint16_t WriteSameSuported : 1;
		uint16_t ErrorRecoveryControlSupported : 1;
		uint16_t FeatureControlSuported : 1;
		uint16_t DataTablesSuported : 1;
		uint16_t Reserved1 : 6;
		uint16_t VendorSpecific : 4;
	} SCTCommandTransport;
	uint16_t ReservedWord207[2];
	struct {
		uint16_t AlignmentOfLogicalWithinPhysical : 14;
		uint16_t Word209Supported : 1;
		uint16_t Reserved0 : 1;
	} BlockAlignment;
	uint16_t WriteReadVerifySectorCountMode3Only[2];
	uint16_t WriteReadVerifySectorCountMode2Only[2];
	struct {
		uint16_t NVCachePowerModeEnabled : 1;
		uint16_t Reserved0 : 3;
		uint16_t NVCacheFeatureSetEnabled : 1;
		uint16_t Reserved1 : 3;
		uint16_t NVCachePowerModeVersion : 4;
		uint16_t NVCacheFeatureSetVersion : 4;
	} NVCacheCapabilities;
	uint16_t NVCacheSizeLSW;
	uint16_t NVCacheSizeMSW;
	uint16_t NominalMediaRotationRate;
	uint16_t ReservedWord218;
	struct {
		uint8_t NVCacheEstimatedTimeToSpinUpInSeconds;
		uint8_t Reserved;
	} NVCacheOptions;
	uint16_t WriteReadVerifySectorCountMode : 8;
	uint16_t ReservedWord220 : 8;
	uint16_t ReservedWord221;
	struct {
		uint16_t MajorVersion : 12;
		uint16_t TransportType : 4;
	} TransportMajorVersion;
	uint16_t TransportMinorVersion;
	uint16_t ReservedWord224[6];
	uint32_t ExtendedNumberOfUserAddressableSectors[2];
	uint16_t MinBlocksPerDownloadMicrocodeMode03;
	uint16_t MaxBlocksPerDownloadMicrocodeMode03;
	uint16_t ReservedWord236[19];
	uint16_t Signature : 8;
	uint16_t CheckSum : 8;
} __attribute__((packed));

struct part_table_entry{
	uint32_t status:8;
	uint32_t head_start:8;
	uint32_t sector_start:6;
	uint32_t cyl_start:10;
	uint32_t type:8;
	uint32_t head_end:8;
	uint32_t sector_end:6;
	uint32_t cyl_end:10;
	uint32_t lba;
	uint32_t size;
} __attribute__((packed));

static struct ata_controller controllers[MAX_ATA_CONTROLLER] = {{
								    .id               = 0,
								    .base             = 0x1F0,
								    .dev_ctl          = 0x3F6,
								    .present          = 0,
								    .last_device_used = -1,
								},
								{
								    .id               = 1,
								    .base             = 0x170,
								    .dev_ctl          = 0x376,
								    .present          = 0,
								    .last_device_used = -1,
								}};
int ATADetectDevice(struct ata_device *dev)
{
	struct ata_controller *ctl = dev->ctl;
	outb(ctl->base + ATA_PIO_DRIVE,
	     ATA_PIO_DRIVE_IBM | (dev->isSlave ? ATA_PIO_DRIVE_SLAVE : ATA_PIO_DRIVE_MASTER));
	unsigned st = inb(ctl->base + ATA_PIO_STATUS);
	if (st & ATA_PIO_STATUS_DRIVE_BUSY)
		goto no_disk;

	unsigned cl = inb(ctl->base + ATA_PIO_LBAMID); /* get the "signature bytes" */
	unsigned ch = inb(ctl->base + ATA_PIO_LBAHI);

	/* differentiate ATA, ATAPI, SATA and SATAPI */
	if (cl == 0x14 && ch == 0xEB) {
		return ATA_DEV_PATAPI;
	}
	if (cl == 0x69 && ch == 0x96) {
		return ATA_DEV_SATAPI;
	}
	if (cl == 0 && ch == 0) {
		return ATA_DEV_PATA;
	}
	if (cl == 0x3c && ch == 0xc3) {
		return ATA_DEV_SATA;
	}
no_disk:
	return ATA_DEV_UNKNOWN;
}

int ATAGetDeviceInfo(struct ata_device *dev)
{
	int status;
	int16_t buffer[256];
	struct ata_controller *ctl = dev->ctl;
	struct _IDENTIFY_DEVICE_DATA *deviceInfo;

	outb(ctl->base + ATA_PIO_DRIVE,
	     ATA_PIO_DRIVE_IBM | (dev->isSlave ? ATA_PIO_DRIVE_SLAVE : ATA_PIO_DRIVE_MASTER));
	outb(ctl->base + ATA_PIO_CMD, ATA_PIO_CMD_IDENTIFY);

	/* Wait for command completion (wait while busy bit is set) */
	for (int timeout = 0; timeout < 30000; timeout++) {
		status = inb(ctl->base + ATA_PIO_STATUS);
		if (!(status & ATA_PIO_STATUS_DRIVE_BUSY))
			break;
	}

	/* DRQ bit indicates that data is ready to be read. If it is not set
	   after an IDENTIFY command, there is a problem */
	if (!(status & ATA_PIO_STATUS_DRQ)) {
		return -1;
	}
	/* Read data from the controller buffer to a temporary buffer */
	for (int i = 0; i < 256; i++)
		buffer[i] = inw(ctl->base + ATA_PIO_DATA);

	deviceInfo = (struct _IDENTIFY_DEVICE_DATA *)buffer;

	dev->heads   = deviceInfo->NumHeads;
	dev->cyls    = deviceInfo->NumCylinders;
	dev->sectors = deviceInfo->NumSectorsPerTrack;
	printf("Disk Detected[%d][%d]: %d Ko\n", ctl->id, dev->id,
	       (dev->heads * dev->cyls * dev->sectors * DISK_SECTOR_SIZE) / 1024);
	ATAReadPartition(dev);

	return 0;
}
const char *partTypeToStr(unsigned int type)
{
	switch (type) {
		case PART_TYPE_EXTENDED:
			return "Extended";
		case PART_TYPE_FAT16:
		case PART_TYPE_FAT16_LBA:
			return "FAT16";
		case PART_TYPE_FAT32:
		case PART_TYPE_FAT32_LBA:
			return "FAT32";
		case PART_TYPE_LINUX_SWAP:
			return "Linux Swap";
		case PART_TYPE_LINUX:
			return "Linux";
		default:
			return "Unknown";
	}
}

struct ata_partition *ATAGetPartition(int id)
{
	struct ata_partition *part;
	int count;
	list_foreach(partitions, part, count)
	{
		if (part->id == id)
			return part;
	}

	return NULL;
}

struct ata_partition *ATAPartitionCreate(uint type, uint size, uint32_t lba,
					 struct ata_device *dev)
{
	struct ata_partition *part =
	    (struct ata_partition *)malloc(sizeof(struct ata_partition));
	if (part == NULL)
		goto err;

	part->id = nextPartitionId++;
	part->type = type;
	part->size = size;
	part->lba = lba;
	part->device = dev;

	uint32_t flags;
	disable_IRQs(flags);

	list_add_tail(partitions, part);
	restore_IRQs(flags);

err:
	return part;
}


int ATAReadPartition(struct ata_device *dev)
{
	char buf[DISK_SECTOR_SIZE];
	int ret = ATAReadSector(dev, 0, 1, buf);

	if (ret)
		return ret;

	struct part_table_entry *parts = (struct part_table_entry *)(buf + PART_TABLE_OFFSET);
	for (int i = 0; i < 4; i++) {
		if (parts[i].size == 0)
			continue;
		if (parts[i].type == PART_TYPE_EXTENDED) {
			printf("Unsupported Partition extended\n");
			continue;
		}
		printf("Got partition %d type %s(0x%x) size %d kb lba 0x%x\n", i,
		       partTypeToStr(parts[i].type), parts[i].type, (parts[i].size * DISK_SECTOR_SIZE) >> 10, parts[i].lba);

		ATAPartitionCreate(parts[i].type, parts[i].size, parts[i].lba, dev);
		ret++;
	}

	return ret;
}

int ATADetectController(struct ata_controller *controller)
{

	mutexInit(&controller->mutex);
	controller->devices[0].id      = 0;
	controller->devices[0].isSlave = 0;
	controller->devices[0].ctl     = controller;
	controller->devices[0].type    = ATADetectDevice(&controllers->devices[0]);

	if (controller->devices[0].type == ATA_DEV_PATA ||
	    controller->devices[0].type == ATA_DEV_SATA) {
		ATAGetDeviceInfo(&controller->devices[0]);
	}

	controller->devices[1].id      = 1;
	controller->devices[1].isSlave = 1;
	controller->devices[1].ctl     = controller;
	controller->devices[1].type    = ATADetectDevice(&controllers->devices[1]);

	if (controller->devices[1].type == ATA_DEV_PATA ||
	    controller->devices[1].type == ATA_DEV_SATA) {
		ATAGetDeviceInfo(&controller->devices[1]);
	}

	return 0;
}

int ATAInit()
{
	list_init(partitions);
	for (uint i = 0; i < MAX_ATA_CONTROLLER; i++) {
		ATADetectController(&controllers[i]);
	}
	return 0;
}

int ATAReadPartitionSector(struct ata_partition *part, int offset, uint nbSector, void *buf)
{
	int lba = part->lba + offset;

	return ATAReadSector(part->device, lba, nbSector, buf);
}

int ATAReadSector(struct ata_device *dev, int lba, uint nb_sector, void *buf)
{
	struct ata_controller *ctl = dev->ctl;

	mutexLock(&ctl->mutex);

	if (ctl->last_device_used != dev->id) {
		outb(ctl->base + ATA_PIO_DRIVE, ATA_PIO_DRIVE_IBM | dev->isSlave << 4);
		while (inb(ctl->base + ATA_PIO_STATUS) & ATA_PIO_STATUS_DRIVE_BUSY) {
		}
		ctl->last_device_used = dev->id;
	}
	outb(ctl->base + ATA_PIO_DRIVE, lba >> 24 | ATA_PIO_DRIVE_IBM | ATA_PIO_DRIVE_LBA);
	outb(ctl->base + ATA_PIO_SEC_COUNT, nb_sector);
	outb(ctl->base + ATA_PIO_LBALO, lba & 0xff);
	outb(ctl->base + ATA_PIO_LBAMID, (lba >> 8) & 0xff);
	outb(ctl->base + ATA_PIO_LBAHI, (lba >> 16) & 0xff);
	outb(ctl->base + ATA_PIO_CMD, ATA_PIO_CMD_READ);

	uint16_t *ptr = (uint16_t *)buf;
	for (uint sector = 0; sector < nb_sector; sector++) {
		while (inb(ctl->base + ATA_PIO_STATUS) & ATA_PIO_STATUS_DRIVE_BUSY) {
		}
		if (inb(ctl->base + ATA_PIO_STATUS) & ATA_PIO_STATUS_REG_ERR) {
			mutexUnlock(&ctl->mutex);
			printf("ATA read error at sector %d\n", sector);
			return -1;
		}
		for (uint i = 0; i < (DISK_SECTOR_SIZE / sizeof(uint16_t)); i++) {
			*ptr = inw(ctl->base + ATA_PIO_DATA);
			ptr++;
		}
	}

	mutexUnlock(&ctl->mutex);

	return 0;
}

int ATAWriteSector(struct ata_device *dev, int lba, uint nb_sector, void *buf)
{
	struct ata_controller *ctl = dev->ctl;
	int ret = 0;

	mutexLock(&ctl->mutex);

	if (ctl->last_device_used != dev->id) {
		outb(ctl->base + ATA_PIO_DRIVE, ATA_PIO_DRIVE_IBM | dev->isSlave << 4);
		while (inb(ctl->base + ATA_PIO_STATUS) & ATA_PIO_STATUS_DRIVE_BUSY) {
		}
		ctl->last_device_used = dev->id;
	}
	outb(ctl->base + ATA_PIO_DRIVE, lba >> 24 | ATA_PIO_DRIVE_IBM | ATA_PIO_DRIVE_LBA);
	outb(ctl->base + ATA_PIO_SEC_COUNT, nb_sector);
	outb(ctl->base + ATA_PIO_LBALO, lba & 0xff);
	outb(ctl->base + ATA_PIO_LBAMID, (lba >> 8) & 0xff);
	outb(ctl->base + ATA_PIO_LBAHI, (lba >> 16) & 0xff);
	outb(ctl->base + ATA_PIO_CMD, ATA_PIO_CMD_WRITE);

	uint16_t *ptr = (uint16_t *)buf;
	for (uint sector = 0; sector < nb_sector; sector++) {
		while (inb(ctl->base + ATA_PIO_STATUS) & ATA_PIO_STATUS_DRIVE_BUSY) {
		}
		if (inb(ctl->base + ATA_PIO_STATUS) & ATA_PIO_STATUS_REG_ERR) {
			mutexUnlock(&ctl->mutex);
			printf("ATA write error\n");
			return -1;
		}
		int retry = 100;
		while (retry) {
			int status = inb(ctl->base + ATA_PIO_STATUS);
			if (status & ATA_PIO_STATUS_REG_ERR) {
				mutexUnlock(&ctl->mutex);
				printf("ATA write error\n");
				return -1;
			}
			if ((status & ATA_PIO_STATUS_DRQ) &&
			    !(status & ATA_PIO_STATUS_DRIVE_BUSY))
				break;
			retry--;
		}
		if(retry == 0){
			printf("ATA write timeout error\n");
			ret=-1;
			break;
		}
		for (uint i = 0; i < (DISK_SECTOR_SIZE / sizeof(uint16_t)); i++) {
			outw(ctl->base + ATA_PIO_DATA, *ptr++);
		}
		retry = 1000;
		// Wait for the device to receive the data
		while (retry) {
			int status = inb(ctl->base + ATA_PIO_STATUS);
			if (status & ATA_PIO_STATUS_REG_ERR) {
				mutexUnlock(&ctl->mutex);
				printf("ATA write error\n");
				return -1;
			}
			if (!(status & ATA_PIO_STATUS_DRQ) &&
			    !(status & ATA_PIO_STATUS_DRIVE_BUSY))
				break;
			retry --;
		}
		if(retry == 0){
			printf("ATA write data timeout error\n");
			ret=-1;
			break;
		}
	}

	mutexUnlock(&ctl->mutex);

	return ret;
}

struct ata_device *ATAGetDevice(int ctlId, int devId)
{
	if (ctlId >= MAX_ATA_CONTROLLER)
		return NULL;
	if (devId >= MAX_ATA_DEVICES)
		return NULL;
	return &controllers[ctlId].devices[devId];
}