// SPDX-License-Identifier: MPL-2.0
#include "vybuild.hpp"
#include <zlib.h>
using namespace std;
#pragma pack(1)
struct mbr_entry {
  uint8_t boot_indicator=0x00;
  uint8_t start_chs[3]={0,0x02,0};
  uint8_t partition_type=0xEE;
  uint8_t end_chs[3]={0xFF,0xFF,0xFF};
  uint32_t start_lba=1;
  uint32_t last_lba=0xFFFFFFFF;
};
#pragma pack()
#pragma pack(1)
struct mbr_protect {
  uint8_t bootstrap_code[440];
  uint32_t disk_signature;
  uint16_t reserved=0;
  mbr_entry first_entry;
  uint8_t other_entries[48]={0};
  uint16_t signature=0xAA55;
};
#pragma pack()
#pragma pack(1)
struct gpt_header {
  char signature[8]={'E','F','I',' ','P','A','R','T'};
  uint32_t revision=0x00010000;
  uint32_t header_size=92;
  uint32_t header_crc32;
  uint32_t reserved=0;
  uint64_t current_lba=1;
  uint64_t backup_lba;
  uint64_t first_usable_lba;
  uint64_t last_usable_lba;
  uint8_t disk_guid[16];
  uint64_t partition_entries_lba=2;
  uint32_t num_partitions_entries=128;
  uint32_t size_of_partitions_entries=128;
  uint32_t partitions_entries_crc32;
};
#pragma pack()
#pragma pack(1)
struct partition_entry {
  uint8_t type_guid[16];
  uint8_t unique_guid[16];
  uint64_t start_lba;
  uint64_t last_lba;
  uint64_t attributes;
  uint16_t name[36];
};
#pragma pack()
struct disk_container {
  mbr_protect mbr;
  gpt_header header;
  array<partition_entry,128> parts;
  vector<uint8_t> data;
  size_t part_count=0;
  size_t lba_bumb=0;
};
array<uint8_t,16> vybuild::disk::uuid4_to_bytes(const string& uuid) {
  array<uint8_t,16> raw{};
  int byte_index=0;
  for (size_t i=0;i<uuid.size();i++) {
    char c=uuid[i];
    if (c=='-') continue;
    auto hex_val=[](char x) -> uint8_t {
      if (x>='0' && x<='9') return x-'0';
      if (x>='a' && x<='f') return x-'a'+10;
      if (x>='A' && x<='F') return x-'A'+10;
      throw std::runtime_error("Invalid hex char in UUID.");
    };
    uint8_t high=hex_val(c);
    if (++i>=uuid.size()) throw std::runtime_error("Malformed UUID.");
    uint8_t low=hex_val(uuid[i]);
    if (byte_index>=16) throw std::runtime_error("UUID too long.");
    raw[byte_index++]=(high<<4)|low;
  }
  if (byte_index!=16) throw std::runtime_error("Invalid UUID size.");
  array<uint8_t,16> out{};
  out[0]=raw[3];
  out[1]=raw[2];
  out[2]=raw[1];
  out[3]=raw[0];
  out[4]=raw[5];
  out[5]=raw[4];
  out[6]=raw[7];
  out[7]=raw[6];
  for (size_t i=8;i<16;i++) out[i]=raw[i];
  return out;
}
map<string,disk_container> disk_map;
void vybuild::disk::create_disk(string key,size_t size,string disk_guid) {
  if (size%512!=0) throw std::runtime_error("Disk image size isn't a multiple of 512 bytes.");
  disk_container cont;
  size_t lba_count=size/512;
  size_t last_lba=lba_count-1;
  cont.header.backup_lba=last_lba;
  size_t first_usable_lba=2+((cont.header.size_of_partitions_entries*cont.header.num_partitions_entries)/512);
  cont.header.first_usable_lba=first_usable_lba;
  cont.header.last_usable_lba=last_lba-1;
  auto disk_id=uuid4_to_bytes(disk_guid);
  memcpy(cont.header.disk_guid,disk_id.data(),16);
  memset(cont.parts.data(),0,sizeof(cont.parts));
  cont.header.partitions_entries_crc32=crc32(0,(uint8_t*)cont.parts.data(),sizeof(cont.parts));
  cont.header.header_crc32=0;
  cont.header.header_crc32=crc32(0,(uint8_t*)&cont.header,sizeof(cont.header));
  cont.data.resize((cont.header.last_usable_lba-cont.header.first_usable_lba)*512);
  disk_map[key]=cont;
  return;
}
array<uint16_t,36> gpt_name_to_utf16le(const string& str) {
  std::array<uint16_t,36> out{};
  size_t len=str.size();
  if (len>35) len=35;
  for (size_t i=0;i<len;i++) out[i]=static_cast<uint16_t>(static_cast<unsigned char>(str[i]));
  out[len]=0;
  return out;
}
void add_partition(disk_container& disk,size_t partition_size_bytes,string type_guid,string unique_guid,uint64_t attributes,string name,uint8_t* in_buf,size_t in_size) {
  if (partition_size_bytes==0 || partition_size_bytes%512!=0) throw std::runtime_error("Invalid partition size.");
  if ((disk.lba_bumb*512)+partition_size_bytes>=disk.data.size()) throw std::runtime_error("Provided partition size would overflow the disk.");
  if (in_size>partition_size_bytes) throw std::runtime_error("Provided buffer is too big for provided partition size.");
  if (disk.part_count==128) throw std::runtime_error("Too much partitions.");
  partition_entry entry{};
  entry.attributes=attributes;
  size_t lba_count=partition_size_bytes/512;
  entry.start_lba=disk.header.first_usable_lba+disk.lba_bumb;
  entry.last_lba=entry.start_lba+lba_count-1;
  auto gpt_name=gpt_name_to_utf16le(name);
  memcpy(entry.name,gpt_name.data(),sizeof(gpt_name));
  auto type=vybuild::disk::uuid4_to_bytes(type_guid);
  auto unique=vybuild::disk::uuid4_to_bytes(unique_guid);
  memcpy(entry.type_guid,type.data(),16);
  memcpy(entry.unique_guid,unique.data(),16);
  disk.parts[disk.part_count]=entry;
  memcpy(disk.data.data()+disk.lba_bumb*512,in_buf,in_size);
  disk.part_count++;
  disk.lba_bumb+=lba_count;
  return;
}
void vybuild::disk::add_partition(string key,size_t partition_size_bytes,string type_guid,string unique_guid,uint64_t attributes,string name,uint8_t *in_buf,size_t in_size) {
  if (disk_map.find(key)==disk_map.end()) throw std::runtime_error(key+" isn't a valid disk key.");
  disk_container& disk=disk_map[key];
  add_partition(disk,partition_size_bytes,type_guid,unique_guid,attributes,name,in_buf,in_size);
  return;
}
static void write_sector(vector<uint8_t>& out,size_t lba,const void* data,size_t size) {
  size_t offset=lba*512;
  if (out.size()<offset+size) out.resize(offset+size,0);
  memcpy(out.data()+offset,data,size);
  return;
}
vector<uint8_t> export_disk(disk_container& cont) {
  vector<uint8_t> disk;
  size_t total_sectors=(cont.header.backup_lba+1);
  disk.resize(total_sectors*512,0);
  write_sector(disk,0,&cont.mbr,sizeof(cont.mbr));
  gpt_header header=cont.header;
  auto parts_crc=crc32(0,(uint8_t*)cont.parts.data(),sizeof(cont.parts));
  header.partitions_entries_crc32=parts_crc;
  header.header_crc32=0;
  header.header_crc32=crc32(0,(uint8_t*)&header,sizeof(header));
  write_sector(disk,1,&header,sizeof(header));
  write_sector(disk,cont.header.partition_entries_lba,cont.parts.data(),sizeof(cont.parts));
  if (!cont.data.empty()) {
    size_t data_lba=cont.header.first_usable_lba;
    write_sector(disk,data_lba,cont.data.data(),cont.data.size());
  }
  write_sector(disk,cont.header.backup_lba-(sizeof(cont.parts)/512),cont.parts.data(),sizeof(cont.parts));
  gpt_header backup=cont.header;
  backup.current_lba=cont.header.backup_lba;
  backup.backup_lba=1;
  backup.header_crc32=0;
  backup.header_crc32=crc32(0,(uint8_t*)&backup,sizeof(backup));
  write_sector(disk,cont.header.backup_lba,&backup,sizeof(backup));
  return disk;
}
vector<uint8_t> vybuild::disk::export_disk(string key) {
  if (disk_map.find(key)==disk_map.end()) throw std::runtime_error(key+" isn't a valid disk key.");
  disk_container& fat=disk_map[key];
  return export_disk(fat);
}