#include "libisofs.h"

#include "eltorito.h"
#include "volume.h"
#include "util.h"

#include <assert.h>
#include <malloc.h>
#include <time.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <unistd.h>
#include <string.h>
#include <stdio.h>
#include <fcntl.h>


struct el_torito_validation_entry {
	uint8_t header_id			BP(1, 1);
	uint8_t platform_id			BP(2, 2);
	uint8_t reserved			BP(3, 4);
	uint8_t id_string			BP(5, 28);
	uint8_t checksum			BP(29, 30);
	uint8_t key_byte1			BP(31, 31);
	uint8_t key_byte2			BP(32, 32);
};

struct el_torito_default_entry {
	uint8_t boot_indicator		BP(1, 1);
	uint8_t boot_media_type		BP(2, 2);
	uint8_t load_seg			BP(3, 4);
	uint8_t system_type			BP(5, 5);
	uint8_t unused1				BP(6, 6);
	uint8_t sec_count			BP(7, 8);
	uint8_t block				BP(9, 12);
	uint8_t unused2				BP(13, 32);
};

struct el_torito_section_header_entry {
	uint8_t header_indicator	BP(1, 1);
	uint8_t platform_id			BP(2, 2);
	uint8_t number				BP(3, 4);
	uint8_t character			BP(5, 32);
};

struct el_torito_section_entry {
	uint8_t boot_indicator		BP(1, 1);
	uint8_t boot_media_type		BP(2, 2);
	uint8_t load_seg			BP(3, 4);
	uint8_t system_type			BP(5, 5);
	uint8_t unused1				BP(6, 6);
	uint8_t sec_count			BP(7, 8);
	uint8_t block				BP(9, 12);
	uint8_t selec_criteria		BP(13, 13);
	uint8_t vendor_sc			BP(14, 32);
};

/**
 * This table should be written with accuracy values at offset
 * 8 of boot image, when used ISOLINUX boot loader
 */
struct boot_info_table {
	uint8_t bi_pvd				BP(1, 4);  /* LBA of primary volume descriptor */
	uint8_t bi_file				BP(5, 8);  /* LBA of boot file */
	uint8_t bi_length			BP(9, 12); /* Length of boot file */
	uint8_t bi_csum				BP(13, 16); /* Checksum of boot file */
	uint8_t bi_reserved			BP(17, 56); /* Reserved */	
};

struct partition_desc {
	uint8_t boot_ind;
	uint8_t begin_chs[3];
	uint8_t type;
	uint8_t end_chs[3];
	uint8_t start[4];
	uint8_t size[4];
};

struct hard_disc_mbr {
	uint8_t code_area[440];
	uint8_t opt_disk_sg[4];
	uint8_t pad[2];
	struct partition_desc partition[4];
	uint8_t sign1;
	uint8_t sign2;
};

static struct el_torito_boot_image *
create_image(struct iso_tree_node *image, 
             enum eltorito_boot_media_type type)
{
	struct el_torito_boot_image *boot;
	int boot_media_type = 0;
	int load_sectors = 0; /* number of sector to load */
	unsigned char partition_type = 0;
	
	switch (type) {
	case ELTORITO_FLOPPY_EMUL:
		switch (image->attrib.st_size) {
		case 1200 * 1024:
			boot_media_type = 1; /* 1.2 meg diskette */
			break;
		case 1440 * 1024:
			boot_media_type = 2; /* 1.44 meg diskette */
			break;
		case 2880 * 1024:
			boot_media_type = 3; /* 2.88 meg diskette */
			break;
		default:
			fprintf(stderr, "Invalid image size %d Kb. Must be one of 1.2, 1.44"
			               "or 2.88 Mb", (int) image->attrib.st_size / 1024);
			libisofs_errno = ELTORITO_WRONG_IMAGE_SIZE;
			return NULL;
			break;	
		}
		/* it seems that for floppy emulation we need to load 
		 * a single sector (512b) */
		load_sectors = 1;
		break;
	case ELTORITO_HARD_DISC_EMUL:
		{
		size_t i;
		int fd;
		struct hard_disc_mbr mbr;
		int used_partition;
		
		/* read the MBR on disc and get the type of the partition */
		fd = open(((struct iso_tree_node_file*)image)->loc.path, O_RDONLY);
		if ( fd == -1 ) {
			fprintf(stderr, "Can't open image file\n");
			return NULL;
		}
		if ( read(fd, &mbr, sizeof(mbr)) ) {
			fprintf(stderr, "Can't read MBR from image file\n");
			close(fd);
			return NULL;
		}
		close(fd);
		
		/* check valid MBR signature */
		if ( mbr.sign1 != 0x55 || mbr.sign2 != 0xAA ) {
			fprintf(stderr, "Invalid MBR. Wrong signature.\n");
			return NULL;
		}
		
		/* ensure single partition */
		used_partition = -1;
		for (i = 0; i < 4; ++i) {
			if (mbr.partition[i].type != 0) {
				/* it's an used partition */
				if (used_partition != -1) {
					fprintf(stderr, "Invalid MBR. At least 2 paritions: %d and " 
									"%d, are being used\n", used_partition, i);
					return NULL;
				} else
					used_partition = i;
			}
		}
		partition_type = mbr.partition[used_partition].type;
		}
		boot_media_type = 4;
		
		/* only load the MBR */
		load_sectors = 1;
		break;
	case ELTORITO_NO_EMUL:
		boot_media_type = 0;
		break;	
	}
	
	boot = calloc(1, sizeof(struct el_torito_boot_image));
	boot->bootable = 1;
	boot->image = (struct iso_tree_node_file *) image;
	boot->type = boot_media_type;
	boot->load_size = load_sectors;
	boot->partition_type = partition_type;
	return boot;
}

static struct iso_tree_node_boot_catalog*
create_boot_catalog_node(struct iso_tree_node_dir *parent, 
				         const char *name)
{
	struct iso_tree_node_boot_catalog *boot;
	
	assert( parent && name );
	
	boot = calloc(1, sizeof(struct iso_tree_node_boot_catalog));
	boot->node.attrib.st_mode = S_IFREG | 0444;
	boot->node.attrib.st_atime = 
		boot->node.attrib.st_mtime = 
		boot->node.attrib.st_ctime = time(NULL);
	boot->node.type = LIBISO_NODE_BOOTCATALOG;
	boot->node.name = strdup(name);
	iso_tree_add_child(parent, (struct iso_tree_node*) boot);
	return boot;
}

struct el_torito_boot_image *
iso_volume_create_boot_catalog(struct iso_volume *volume, 
                               struct iso_tree_node *image,
                               enum eltorito_boot_media_type type,
                               struct iso_tree_node_dir *dir, 
                               char *name)
{
	struct el_torito_boot_image *boot_image;
	struct iso_tree_node_boot_catalog *boot_node;
	struct el_torito_boot_catalog *catalog;
	
	assert( volume && !volume->bootcat);
	assert( image && ISO_ISREG(image) && dir && name);
	
	boot_image = create_image(image, type);
	if ( !boot_image ) 
		return NULL;
	
	/* creates the catalog with the given default image */
	catalog = malloc(sizeof(struct el_torito_boot_catalog));
	catalog->nentries = 1;
	catalog->entries = malloc(sizeof(struct el_torito_boot_image *));
	catalog->entries[0] = boot_image;
	catalog->file = malloc(sizeof(struct iso_file));
	catalog->file->size = 2048;
	
	/* add catalog file */
	boot_node = create_boot_catalog_node(dir, name);
	boot_node->catalog = catalog;
	
	volume->bootcat = catalog;
	
	return boot_image;
}

void
el_torito_set_load_seg(struct el_torito_boot_image *bootimg, int segment)
{
	if (bootimg->type != ELTORITO_NO_EMUL)
		return;
	bootimg->load_seg = segment;
}

void
el_torito_set_load_size(struct el_torito_boot_image *bootimg, int sectors)
{
	if (bootimg->type != ELTORITO_NO_EMUL)
		return;
	bootimg->load_size = sectors;
}

void
el_torito_set_no_bootable(struct el_torito_boot_image *bootimg)
{
	bootimg->bootable = 0;
}

void
el_torito_set_write_boot_info(struct el_torito_boot_image *bootimg)
{
	bootimg->patch_isolinux = 1;
}

void el_torito_boot_catalog_free(struct el_torito_boot_catalog *cat)
{
	size_t i;
	
	assert(cat);
	
	for(i = 0; i < cat->nentries; ++i) {
		free(cat->entries[i]);
	}
	free(cat->entries);
	free(cat->file);
	free(cat);
}

void el_torito_get_image_files(struct ecma119_write_target *t)
{
	size_t i;
	struct el_torito_boot_catalog *cat = t->catalog;
	assert(cat);
	
	for(i = 0; i < cat->nentries; ++i) {
		
		struct iso_tree_node_file *image = cat->entries[i]->image;
		struct iso_file *file = iso_file_table_lookup(t->file_table, image);
		if ( file == NULL ) {
			file = iso_file_new(image);
			iso_file_table_add_file(t->file_table, file);
		}
		cat->entries[i]->file = file;
	}
}

/**
 * Write the Boot Record Volume Descriptor
 */
static void
write_boot_vol_desc(struct ecma119_write_target *t, uint8_t *buf)
{
	struct el_torito_boot_catalog *cat = t->catalog;
	struct ecma119_boot_rec_vol_desc *vol = 
			(struct ecma119_boot_rec_vol_desc*)buf;
	
	assert(cat);
	
	vol->vol_desc_type[0] = 0;
	memcpy(vol->std_identifier, "CD001", 5);
	vol->vol_desc_version[0] = 1;
	memcpy(vol->boot_sys_id, "EL TORITO SPECIFICATION", 23);
	iso_lsb(vol->boot_catalog, cat->file->block, 4);
}

static void 
write_validation_entry(struct ecma119_write_target *t, uint8_t *buf)
{
	size_t i;
	int checksum;
	
	struct el_torito_validation_entry *ve = 
		(struct el_torito_validation_entry*)buf;
	ve->header_id[0] = 1;
	ve->platform_id[0] = 0; /* 0: 80x86, 1: PowerPC, 2: Mac */
	ve->key_byte1[0] = 0x55;
	ve->key_byte2[0] = 0xAA;
	
	/* calculate the checksum, to ensure sum of all words is 0 */
	checksum = 0;
	for (i = 0; i < sizeof(struct el_torito_validation_entry); i += 2) {
		checksum -= buf[i];
		checksum -= (buf[i] << 8);
	}
	iso_lsb(ve->checksum, checksum, 2);
}

static void
patch_boot_file(struct el_torito_boot_image *img)
{
	struct boot_info_table info;
	int fd;
	uint32_t checksum;
	ssize_t len;
	uint8_t buf[4];
	
	memset(&info, 0, sizeof(info));
	
	/* open image */
	fd = open(img->image->loc.path, O_RDWR);
	if ( fd == -1 ) {
		//TODO what do do? exit or just continue? 
		fprintf(stderr, "Can't patch boot image %s\n", img->image->loc.path);
		close(fd);
		return;
	}
	
	/* compute checksum, as the the sum of all 32 bit words in boot image
	 * from offset 64 */
	checksum = 0;
	lseek(fd, (off_t) 64, SEEK_SET);
	
	//TODO this can (must) be optimizied by reading to a longer buffer
	while ( (len = read(fd, buf, 4) ) == 4 ) {
		checksum += iso_read_lsb(buf, 4);
	}
	if ( len != 0 ) {
		/* error reading file, or file length not multiple of 4 */
		//TODO what do do? exit or just continue? 
		fprintf(stderr, "Can't patch boot image %s\n", img->image->loc.path);
		close(fd);
		return;
	}
	
	/* fill boot info table */
	iso_lsb(info.bi_pvd, 16, 4); //FIXME this should be changed when we implement ms
	iso_lsb(info.bi_file, img->file->block, 4);
	iso_lsb(info.bi_length, img->image->node.attrib.st_size, 4);
	iso_lsb(info.bi_csum, checksum, 4);
	
	/* patch file */
	lseek(fd, (off_t) 8, SEEK_SET);
	write(fd, &info, sizeof(info));
	
	close(fd);
} 

void 
el_torito_patch_image_files(struct ecma119_write_target *t)
{
	size_t i;
	struct el_torito_boot_catalog *cat = t->catalog;
	assert(cat);
	
	for (i = 0; i < cat->nentries; ++i) {
		struct el_torito_boot_image *img = cat->entries[i];
		if ( img->patch_isolinux )
			patch_boot_file(img);
	}
}

/**
 * Write one section entry.
 * Currently this is used for both default and other entries since we
 * put selection criteria no 0 (no sel. criteria)
 */
static void
write_section_entry(uint8_t *buf, struct el_torito_boot_image *img)
{
	struct el_torito_section_entry *se = 
		(struct el_torito_section_entry*)buf;

	se->boot_indicator[0] = img->bootable ? 0x88 : 0x00;
	se->boot_media_type[0] = img->type;
	iso_lsb(se->load_seg, img->load_seg, 2);
	se->system_type[0] = 0; //TODO need to get the partition type
	iso_lsb(se->sec_count, img->load_size, 2);
	iso_lsb(se->block, img->file->block, 4);
}

/**
 * Write El-Torito Boot Catalog
 */
static void
write_catalog(struct ecma119_write_target *t, uint8_t *buf)
{
	struct el_torito_boot_catalog *cat = t->catalog;
	assert(cat);
	assert(cat->nentries >= 1 && cat->nentries < 63);
	
	write_validation_entry(t, buf);
	
	/* write default entry */
	write_section_entry(buf + 32, cat->entries[0]);

	//TODO write all images	
}

void
el_torito_wr_boot_vol_desc(struct ecma119_write_target *t, uint8_t *buf)
{
	assert(t->catalog);
	ecma119_start_chunking(t,
			       write_boot_vol_desc,
			       2048,
			       buf);
}

void
el_torito_wr_catalog(struct ecma119_write_target *t, uint8_t *buf)
{
	assert(t->catalog);
	ecma119_start_chunking(t,
			       write_catalog,
			       2048,
			       buf);
}