// ata.v: ATA driver.
// Code is governed by the GPL-2.0 license.
// Copyright (C) 2021-2022 The Vinix authors.

module ata

import pci
import stat
import klock
import memory
import resource
import fs
import x86.kio
import lib
import errno
import event.eventstruct
import katomic

const pci_class = 0x1

const pci_subclass = 0x1

const pci_progif = 0x80

const ata_ports = [0x1f0, 0x1f0, 0x170, 0x170]

const ata_bytes_per_sector = 512

const ata_sectors_per_prdt = 16

const ata_bytes_per_prdt = ata_bytes_per_sector * ata_sectors_per_prdt

@[packed]
struct PRDT {
pub mut:
	buffer_phys   u32
	transfer_size u16
	mark_end      u16
}

struct ATADrive {
pub mut:
	stat     stat.Stat
	refcount int
	l        klock.Lock
	event    eventstruct.Event
	status   int
	can_mmap bool
	// ATA specific stuff.
	is_master         bool
	data_port         u16
	error_port        u16
	sector_count_port u16
	lba_low_port      u16
	lba_mid_port      u16
	lba_hi_port       u16
	device_port       u16
	cmd_port          u16
	bar4              u16
	bmr_command       u16
	bmr_status        u16
	bmr_prdt          u16
	prdt              &PRDT
	prdt_phys         u32
	prdt_cache        &u8
}

pub fn initialise() {
	// Fetch and setup the PCI device.
	mut dev := pci.get_device_by_class(ata.pci_class, ata.pci_subclass, ata.pci_progif,
		0) or {
		print('ata: No suitable PCI devices found\n')
		return
	}
	dev.enable_bus_mastering()

	// Probe ports and add devices.
	mut index := 0
	for i in 0 .. ata.ata_ports.len {
		drive := init_ata_drive(i, mut dev) or { continue }
		name := 'ata$index'
		fs.devtmpfs_add_device(drive, name)
		print('ata: Port $i initialised as $name ($drive.stat.size bytes)\n')
		index += 1
	}
}

fn init_ata_drive(port_index int, mut pci_device pci.PCIDevice) ?&ATADrive {
	assert port_index < ata.ata_ports.len

	// Fetch the IO ports and comm addresses of the drive.
	port := u16(ata.ata_ports[port_index])
	mut bar4 := u16(pci_device.read<u32>(0x20))
	if bar4 & 1 != 0 {
		bar4 &= 0xfffffffc
	}
	mut dev := &ATADrive{
		is_master: port_index % 2 != 0
		data_port: port
		error_port: port + 1
		sector_count_port: port + 2
		lba_low_port: port + 3
		lba_mid_port: port + 4
		lba_hi_port: port + 5
		device_port: port + 6
		cmd_port: port + 7
		bar4: bar4
		bmr_command: bar4
		bmr_status: bar4 + 2
		bmr_prdt: bar4 + 4
		prdt: &PRDT(unsafe { nil })
		// To be filled later.
		prdt_phys: 0
		// To be filled later.
		prdt_cache: &u8(0)
		// To be filled later.
	}

	// Identify the drive.
	val := if dev.is_master { 0xa0 } else { 0xb0 }
	kio.port_out<u8>(dev.device_port, u8(val))
	kio.port_out<u8>(dev.sector_count_port, 0)
	kio.port_out<u8>(dev.lba_low_port, 0)
	kio.port_out<u8>(dev.lba_mid_port, 0)
	kio.port_out<u8>(dev.lba_hi_port, 0)
	kio.port_out<u8>(dev.cmd_port, 0xec)

	if kio.port_in<u8>(dev.cmd_port) == 0 {
		print('ata: Port $port_index is not connected\n')
		return none
	} else {
		mut timeout := 0
		for kio.port_in<u8>(dev.cmd_port) & 0b10000000 != 0 {
			if timeout == 100000 {
				print('ata: Port $port_index is not answering\n')
				return none
			}
			timeout += 1
		}
	}

	// Check for non-standard ATAPI.
	if kio.port_in<u8>(dev.lba_mid_port) != 0 || kio.port_in<u8>(dev.lba_hi_port) != 0 {
		print('ata: Port $port_index is non-standard ATAPI\n')
		return none
	}

	mut timeout := 0
	for {
		status := kio.port_in<u8>(dev.cmd_port)
		if status & 0b00000001 != 0 {
			print('ata: Port $port_index errored out\n')
			return none
		}
		if status & 0b00001000 != 0 {
			break
		}
		if timeout == 100000 {
			print('ata: Port $port_index hanged\n')
			return none
		}
	}

	// Storing identify info.
	mut identify := [256]u16{}
	for i in 0 .. identify.len {
		identify[i] = kio.port_in<u16>(dev.data_port)
	}

	// Wacky PCI things
	mut cmd_register := pci_device.read<u32>(0x4)
	if cmd_register & (1 << 2) == 0 {
		cmd_register |= 1 << 2
		pci_device.write(0x4, cmd_register)
	}

	// Final touches to the dev structure and return.
	dev.stat.blocks = unsafe { (&u64(u64(&identify[100])))[0] }
	dev.stat.blksize = 512
	dev.stat.size = dev.stat.blocks * dev.stat.blksize
	dev.stat.rdev = resource.create_dev_id()
	dev.stat.mode = 0o644 | stat.ifblk
	dev.prdt_phys = u32(u64(memory.pmm_alloc(1)))
	dev.prdt = &PRDT(dev.prdt_phys + higher_half)
	dev.prdt.buffer_phys = u32(u64(memory.pmm_alloc(lib.div_roundup<u64>(ata.ata_bytes_per_prdt,
		page_size))))
	dev.prdt.transfer_size = ata.ata_bytes_per_prdt
	dev.prdt.mark_end = 0x8000
	dev.prdt_cache = &u8(u64(dev.prdt.buffer_phys) + higher_half)
	return dev
}

fn (mut this ATADrive) mmap(page u64, flags int) voidptr {
	return 0
}

fn (mut dev ATADrive) read(handle voidptr, buffer voidptr, loc u64, count u64) ?i64 {
	// Check alignment to the sector boundary.
	if loc % ata.ata_bytes_per_sector != 0 || count % ata.ata_bytes_per_sector != 0 {
		errno.set(errno.eio)
		return none
	}

	// Actually read.
	sector_start := loc / ata.ata_bytes_per_sector
	sector_count := count / ata.ata_bytes_per_sector

	dev.l.acquire()
	defer {
		dev.l.release()
	}

	for i := u64(0); i < sector_count; i += ata.ata_sectors_per_prdt {
		sector_loc := sector_start + i
		kio.port_out<u8>(dev.bmr_command, 0)
		kio.port_out<u32>(dev.bmr_prdt, dev.prdt_phys)
		bmr_status := kio.port_in<u8>(dev.bmr_status)
		kio.port_out<u8>(dev.bmr_status, bmr_status | 0x4 | 0x2)

		val := if dev.is_master { 0x40 } else { 0x50 }
		kio.port_out<u8>(dev.device_port, u8(val))

		actual_count := if i + ata.ata_sectors_per_prdt > sector_count {
			sector_count % ata.ata_sectors_per_prdt
		} else {
			ata.ata_sectors_per_prdt
		}

		kio.port_out<u8>(dev.sector_count_port, u8(actual_count >> 8))
		kio.port_out<u8>(dev.lba_low_port, u8((sector_loc & 0x000000FF000000) >> 24))
		kio.port_out<u8>(dev.lba_mid_port, u8((sector_loc & 0x0000FF00000000) >> 32))
		kio.port_out<u8>(dev.lba_hi_port, u8((sector_loc & 0x00FF0000000000) >> 40))
		kio.port_out<u8>(dev.sector_count_port, u8(actual_count & 0xff))
		kio.port_out<u8>(dev.lba_low_port, u8(sector_loc & 0x000000000000FF))
		kio.port_out<u8>(dev.lba_mid_port, u8((sector_loc & 0x0000000000FF00) >> 8))
		kio.port_out<u8>(dev.lba_hi_port, u8((sector_loc & 0x00000000FF0000) >> 16))

		kio.port_out<u8>(dev.cmd_port, 0x25) // READ_DMA command
		kio.port_out<u8>(dev.bmr_command, 0x8 | 0x1)

		for {
			status := kio.port_in<u8>(dev.cmd_port)
			if status & 0x80 == 0 {
				break
			}
			if status & 0x01 != 0 {
				print('ata: Error reading sector $sector_loc on drive')
				return none
			}
		}
		kio.port_out<u8>(dev.bmr_command, 0)

		buffer_final := voidptr(u64(buffer) + i * ata.ata_bytes_per_sector)
		unsafe { C.memcpy(buffer_final, dev.prdt_cache, actual_count * ata.ata_bytes_per_sector) }
	}
	return i64(count)
}

fn (mut dev ATADrive) write(handle voidptr, buffer voidptr, loc u64, count u64) ?i64 {
	// Check alignment to the sector boundary.
	if loc % ata.ata_bytes_per_sector != 0 || count % ata.ata_bytes_per_sector != 0 {
		errno.set(errno.eio)
		return none
	}

	// Actually write.
	sector_start := loc / ata.ata_bytes_per_sector
	sector_count := count / ata.ata_bytes_per_sector

	dev.l.acquire()
	defer {
		dev.l.release()
	}

	for i := u64(0); i < sector_count; i += ata.ata_sectors_per_prdt {
		sector_loc := sector_start + i

		kio.port_out<u8>(dev.bmr_command, 0)
		kio.port_out<u32>(dev.bmr_prdt, dev.prdt_phys)
		bmr_status := kio.port_in<u8>(dev.bmr_status)
		kio.port_out<u8>(dev.bmr_status, bmr_status | 0x4 | 0x2)

		val := if dev.is_master { 0x40 } else { 0x50 }
		kio.port_out<u8>(dev.device_port, u8(val))

		actual_count := if i + ata.ata_sectors_per_prdt > sector_count {
			sector_count % ata.ata_sectors_per_prdt
		} else {
			ata.ata_sectors_per_prdt
		}

		// Copy buffer to DMA area.
		buffer_final := voidptr(u64(buffer) + i * ata.ata_bytes_per_sector)
		unsafe { C.memcpy(dev.prdt_cache, buffer_final, actual_count * ata.ata_bytes_per_sector) }

		kio.port_out<u8>(dev.sector_count_port, u8(actual_count >> 8))
		kio.port_out<u8>(dev.lba_low_port, u8((sector_loc & 0x000000FF000000) >> 24))
		kio.port_out<u8>(dev.lba_mid_port, u8((sector_loc & 0x0000FF00000000) >> 32))
		kio.port_out<u8>(dev.lba_hi_port, u8((sector_loc & 0x00FF0000000000) >> 40))
		kio.port_out<u8>(dev.sector_count_port, u8(actual_count & 0xff))
		kio.port_out<u8>(dev.lba_low_port, u8(sector_loc & 0x000000000000FF))
		kio.port_out<u8>(dev.lba_mid_port, u8((sector_loc & 0x0000000000FF00) >> 8))
		kio.port_out<u8>(dev.lba_hi_port, u8((sector_loc & 0x00000000FF0000) >> 16))

		kio.port_out<u8>(dev.cmd_port, 0x35) // WRITE_DMA command
		kio.port_out<u8>(dev.bmr_command, 0x1)

		for {
			status := kio.port_in<u8>(dev.cmd_port)
			if status & 0x80 == 0 {
				break
			}
			if status & 0x01 != 0 {
				print('ata: Error reading sector $sector_loc on drive')
				return none
			}
		}
		kio.port_out<u8>(dev.bmr_command, 0)

		kio.port_out<u8>(dev.device_port, u8(val))
		kio.port_out<u8>(dev.cmd_port, 0xea) // Cache flush EXT command.
		for {
			status := kio.port_in<u8>(dev.cmd_port)
			if status & 0x80 == 0 {
				break
			}
			if status & 0x01 != 0 {
				print('ata: Error reading sector $sector_loc on drive')
				return none
			}
		}
	}

	return i64(count)
}

fn (mut dev ATADrive) ioctl(handle voidptr, request u64, argp voidptr) ?int {
	return resource.default_ioctl(handle, request, argp)
}

fn (mut dev ATADrive) unref(handle voidptr) ? {
	katomic.dec(dev.refcount)
}

fn (mut dev ATADrive) link(handle voidptr) ? {
	katomic.inc(dev.stat.nlink)
}

fn (mut dev ATADrive) unlink(handle voidptr) ? {
	katomic.dec(dev.stat.nlink)
}

fn (mut dev ATADrive) grow(handle voidptr, new_size u64) ? {
	return none
}
