#include "msc_disk.h"

#include <string.h>

namespace esphome {
namespace grub_boot_switch {

MscDisk::MscDisk() {
  file_date_ = fat_date(2026, 1, 1);
}

void MscDisk::set_volume_label(const char *label) {
  // Pad to 11 chars with spaces, like a FAT volume label.
  memset(volume_label_, ' ', sizeof(volume_label_));
  for (size_t i = 0; i < sizeof(volume_label_) && label[i] != '\0'; ++i) {
    volume_label_[i] = label[i];
  }
}

void MscDisk::set_file_date(uint16_t year, uint8_t month, uint8_t day) {
  file_date_ = fat_date(year, month, day);
}

void MscDisk::add_file(const std::string &long_name, const std::string &short_name,
                       const std::string &short_ext, uint32_t size, FileReadCallback content_cb) {
  VirtualFile vf{};
  vf.long_name = long_name;
  vf.read = std::move(content_cb);

  memset(&vf.dir, 0, sizeof(vf.dir));
  memset(vf.dir.name, ' ', FAT_NAME_LEN);
  memset(vf.dir.ext, ' ', FAT_EXT_LEN);
  for (size_t i = 0; i < FAT_NAME_LEN && i < short_name.size(); ++i) {
    vf.dir.name[i] = short_name[i];
  }
  for (size_t i = 0; i < FAT_EXT_LEN && i < short_ext.size(); ++i) {
    vf.dir.ext[i] = short_ext[i];
  }
  vf.dir.size = size;
  files_.push_back(std::move(vf));
}

void MscDisk::write_boot_sector(uint8_t *out) {
  // Mirrors BootSector[] from reference/src/usb.c:126, plus a configurable
  // volume serial / label.
  static const uint8_t header[] = {
      0xEB, 0x3C, 0x90,
      'm', 'k', 'f', 's', '.', 'f', 'a', 't',
      WBVAL(SECTOR_SIZE),
      SECTORS_PER_CLUSTER,
      WBVAL(RESERVED_SECTORS),
      FAT_COPIES,
      WBVAL(ROOT_ENTRIES),
      WBVAL(SECTOR_COUNT),
      FAT_MEDIA_FIXED_DISK,
      WBVAL(1),    // sectors per FAT
      WBVAL(32),   // sectors per track
      WBVAL(64),   // heads
      QBVAL(0),    // hidden sectors
      QBVAL(0),    // large sector count
      0x00,        // drive number
      0x00,        // reserved
      0x29,        // extended boot signature
  };
  memcpy(out, header, sizeof(header));

  uint8_t *p = out + sizeof(header);
  *p++ = volume_serial_ & 0xFF;
  *p++ = (volume_serial_ >> 8) & 0xFF;
  *p++ = (volume_serial_ >> 16) & 0xFF;
  *p++ = (volume_serial_ >> 24) & 0xFF;

  memcpy(p, volume_label_, sizeof(volume_label_));
  p += sizeof(volume_label_);

  static const char fs_type[8] = {'F', 'A', 'T', '1', '2', ' ', ' ', ' '};
  memcpy(p, fs_type, sizeof(fs_type));

  // Bootable partition signature
  out[SECTOR_SIZE - 2] = 0x55;
  out[SECTOR_SIZE - 1] = 0xAA;
}

void MscDisk::write_fat_sector(uint8_t *out) {
  // First two FAT entries are the media descriptor + EOC marker.
  // Then one EOC entry per file (each file occupies a single cluster).
  out[0] = 0xF8;
  out[1] = 0xFF;
  out[2] = 0xFF;

  // Cluster N → 12-bit entry. Files start at FILEDATA_START_CLUSTER.
  // For single-cluster files we just write 0xFFF (EOC) at each entry.
  uint32_t entry = FILEDATA_START_CLUSTER;
  for (size_t i = 0; i < files_.size(); ++i, ++entry) {
    const uint32_t bit_offset = entry * 12;
    const uint32_t byte_offset = bit_offset / 8;
    if (byte_offset + 1 >= SECTOR_SIZE) {
      break;
    }
    if (entry & 1) {
      // Odd: high nibble of byte N + all of byte N+1
      out[byte_offset] |= 0xF0;
      out[byte_offset + 1] = 0xFF;
    } else {
      // Even: all of byte N + low nibble of byte N+1
      out[byte_offset] = 0xFF;
      out[byte_offset + 1] |= 0x0F;
    }
  }
}

void MscDisk::write_dir_sector(uint8_t *out) {
  uint8_t *p = out;
  for (size_t i = 0; i < files_.size(); ++i) {
    VirtualFile &vf = files_[i];
    vf.dir.start = FILEDATA_START_CLUSTER + i;
    vf.dir.cdate = file_date_;
    vf.dir.mdate = file_date_;
    vf.dir.adate = file_date_;

    p += fat_write_lfn(vf.long_name.c_str(), &vf.dir, p);
    p += fat_write_dir(&vf.dir, p);
  }
}

int MscDisk::read_lba(uint32_t lba, uint8_t *out) {
  memset(out, 0, SECTOR_SIZE);

  switch (lba) {
    case 0:
      write_boot_sector(out);
      return 0;
    case 1:
    case 2:
      write_fat_sector(out);
      return 0;
    case 3:
      write_dir_sector(out);
      return 0;
    default: {
      if (lba < FILEDATA_START_SECTOR) {
        return 0;  // gap between root dir and data region
      }
      const uint32_t file_index = (lba - FILEDATA_START_SECTOR) / SECTORS_PER_CLUSTER;
      if (file_index >= files_.size()) {
        return 0;  // beyond any file
      }
      if (files_[file_index].read) {
        files_[file_index].read(out);
      }
      return 0;
    }
  }
}

}  // namespace grub_boot_switch
}  // namespace esphome