#include "arch/x86_64/disk/ATA.h"
#include "Log.h"
#include "arch/x86_64/IO.h"
#include <luna/Vector.h>

SharedPtr<ATA::Controller> g_controller;

namespace ATA
{
    Result<void> Controller::scan()
    {
        // FIXME: Propagate errors.
        PCI::scan(
            [](const PCI::Device& device) {
                if (!g_controller)
                {
                    auto controller = adopt_shared_if_nonnull(new (std::nothrow) Controller(device)).release_value();
                    kinfoln("ata: Found ATA controller on PCI bus (%x:%x:%x)", device.address.bus,
                            device.address.function, device.address.slot);

                    if (controller->initialize()) g_controller = controller;
                }
            },
            { .klass = 1, .subclass = 1 });

        if (!g_controller) kwarnln("ata: No ATA controller found.");

        return {};
    }

    bool Controller::initialize()
    {
        if (!m_primary_channel.initialize()) return false;

        return m_secondary_channel.initialize();
    }

    Controller::Controller(const PCI::Device& device)
        : m_device(device), m_primary_channel(this, 0, {}), m_secondary_channel(this, 1, {})
    {
    }

    Channel::Channel(Controller* controller, u8 channel_index, Badge<Controller>)
        : m_controller(controller), m_channel_index(channel_index)
    {
    }

    u8 Channel::read_register(Register reg)
    {
        return IO::inb(m_io_base + (u16)reg);
    }

    void Channel::write_register(Register reg, u8 value)
    {
        IO::outb(m_io_base + (u16)reg, value);
    }

    u8 Channel::read_control(ControlRegister reg)
    {
        return IO::inb(m_control_base + (u16)reg);
    }

    void Channel::write_control(ControlRegister reg, u8 value)
    {
        IO::outb(m_control_base + (u16)reg, value);
    }

    void Channel::delay_400ns()
    {
        read_control(ControlRegister::AltStatus);
        read_control(ControlRegister::AltStatus);
        read_control(ControlRegister::AltStatus);
        read_control(ControlRegister::AltStatus);
    }

    void Channel::select(u8 drive)
    {
        if (drive == m_current_drive) return;

        u8 value = (drive << 4) | 0xa0;
        write_register(Register::DriveSelect, value);

        delay_400ns();

        m_current_drive = drive;
    }

    bool Channel::initialize()
    {
        int offset = m_channel_index ? 2 : 0;
        m_is_pci_native_mode = m_controller->device().type.prog_if & (1 << offset);

        u32 control_port_base_address;
        u32 io_base_address;

        if (m_is_pci_native_mode)
        {
            // FIXME: Properly decode BARs.
            io_base_address = PCI::read32(m_controller->device().address, m_channel_index ? PCI::BAR2 : PCI::BAR0);
            control_port_base_address =
                PCI::read32(m_controller->device().address, m_channel_index ? PCI::BAR3 : PCI::BAR1);
        }
        else
        {
            io_base_address = m_channel_index ? 0x170 : 0x1f0;
            control_port_base_address = m_channel_index ? 0x376 : 0x3f6;
        }

        m_io_base = (u16)io_base_address;
        m_control_base = (u16)control_port_base_address + 2;

        if (m_is_pci_native_mode) m_interrupt_line = PCI::read8(m_controller->device().address, PCI::InterruptLine);
        else
            m_interrupt_line = m_channel_index ? 15 : 14;

        for (u8 drive = 0; drive < 2; drive++)
        {
            ScopedKMutexLock<100> lock(m_lock);

            select(drive);

            if (read_register(Register::Status) == 0)
            {
                // No drive on this slot.
                continue;
            }

            kinfoln("ata: Channel %d has a drive on slot %d!", m_channel_index, drive);
        }

        return true;
    }
}