#pragma once
#include "arch/PCI.h"
#include "fs/devices/DeviceRegistry.h"
#include "lib/KMutex.h"
#include <luna/Atomic.h>
#include <luna/SharedPtr.h>

namespace ATA
{
    enum class Register : u16
    {
        Data = 0,
        Error = 1,
        Features = 1,
        SectorCount = 2,
        SectorNumber = 3,
        LBALow = 3,
        CylinderLow = 4,
        LBAMiddle = 4,
        CylinderHigh = 5,
        LBAHigh = 5,
        DriveSelect = 6,
        Status = 7,
        Command = 8,
    };

    enum class ControlRegister : u16
    {
        AltStatus = 0,
        DeviceControl = 0,
        DriveAddress = 1,
    };

    class Controller;
    class Channel;

    class Drive
    {
      private:
        Channel* m_channel;
    };

    class Channel
    {
      public:
        Channel(Controller* controller, u8 channel_index, Badge<Controller>);

        u8 read_register(Register reg);
        void write_register(Register reg, u8 value);
        u8 read_control(ControlRegister reg);
        void write_control(ControlRegister reg, u8 value);

        void delay_400ns();

        void wait_for_irq();
        void irq_handler(Registers*);

        void select(u8 drive);

        bool initialize();

      private:
        Controller* m_controller;
        u8 m_channel_index;
        bool m_is_pci_native_mode;

        u8 m_interrupt_line;

        KMutex<100> m_lock {};

        Thread* m_thread;

        u16 m_io_base;
        u16 m_control_base;

        u8 m_current_drive = (u8)-1;
    };

    class Controller
    {
      public:
        static Result<void> scan();

        const PCI::Device& device() const
        {
            return m_device;
        }

        bool initialize();

      private:
        Controller(const PCI::Device& device);
        PCI::Device m_device;
        Channel m_primary_channel;
        Channel m_secondary_channel;
    };

}

class ATADevice : public Device
{
  public:
    // Initializer for DeviceRegistry.
    static Result<void> create(SharedPtr<ATA::Controller> controller, int channel, int drive);

    Result<usize> read(u8*, usize, usize) const override;

    Result<usize> write(const u8*, usize, usize) override;

    bool blocking() const override;

    Result<u64> ioctl(int request, void* arg) override;

    usize size() const override;

    virtual ~ATADevice() = default;

  private:
    ATADevice(SharedPtr<ATA::Controller> controller);
    SharedPtr<ATA::Controller> m_controller;
    int m_channel;
    int m_drive;
};