#pragma once
#include "arch/PCI.h"
#include "fs/devices/DeviceRegistry.h"
#include "lib/KMutex.h"
#include <luna/Atomic.h>
#include <luna/SharedPtr.h>
#include <luna/StaticString.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 = 7,
    };

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

    enum StatusRegister
    {
        SR_Busy = 0x80,
        SR_DriveReady = 0x40,
        SR_WriteFault = 0x20,
        SR_SeekComplete = 0x10,
        SR_DataRequestReady = 0x08,
        SR_CorrectedData = 0x04,
        SR_Index = 0x02,
        SR_Error = 0x01
    };

    enum CommandRegister
    {
        CMD_Identify = 0xec,
        CMD_Identify_Packet = 0xa1
    };

    class Controller;
    class Channel;

    class Drive
    {
      public:
        Drive(Channel* channel, u8 drive_index, Badge<Channel>);

        bool initialize();

        void irq_handler();

      private:
        bool identify_ata();

        Channel* m_channel;

        u8 m_drive_index;
        union {
            u16 m_identify_words[256];
            u8 m_identify_data[512];
        };

        bool m_is_atapi { false };

        constexpr static usize SERIAL_LEN = 20;
        constexpr static usize REVISION_LEN = 8;
        constexpr static usize MODEL_LEN = 40;

        StaticString<SERIAL_LEN> m_serial;
        StaticString<REVISION_LEN> m_revision;
        StaticString<MODEL_LEN> m_model;
    };

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

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

        bool wait_for_reg_set(Register reg, u8 value, u64 timeout);
        bool wait_for_reg_clear(Register reg, u8 value, u64 timeout);

        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;

        SharedPtr<Drive> m_drives[2];
    };

    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;
};