#include "pata.h"
#include "hd.h"
#include "vfs.h"
#include "memory.h"
#include "string.h"

struct hd_dev **hds;
unsigned char hd_count = 0;

struct trim_struct {
	struct block_buffer **bbs;
	int count;
	int top;
	int filled;
};

int init_hd(struct ide_dev *ide) {
	int i;
	int j;
	int k;
	struct hd_dev *tmp_hd;
	char *buffer;
	struct partition_table *part;
	char name[8];
	struct vfs_device_t *hd_vdev;
	int fstry;
	buffer = (char *)malloc(512);

	for (i=0;i<4;i++) {
		if (ide->ide_devices[i].reserved == 1) {
			if (ide->ide_devices[i].type == IDE_ATA) {
				ide_read_sectors(ide, i, 1, 0, 0x10, (unsigned int)buffer);
				for (j=0;j<4;j++) {
					part = (struct partition_table *)(buffer + (0x1be + (j * 16)));	
					if (part->type == 0x1 || part->type == 0x4 || part->type == 0x6 || part->type ==  0xb || part->type == 0xc || part->type == 0xe || part->type == 0xf || part->type == 0x83) {
						tmp_hd = (struct hd_dev *)malloc(sizeof(struct hd_dev));

						tmp_hd->ide_dev = ide;
						tmp_hd->part = j;
						tmp_hd->part_offset = part->relative_sector;
						tmp_hd->part_len = part->total_sectors;
						tmp_hd->drive = i;
						tmp_hd->block_buffers = hashmap_new();
						if (hd_count == 0) {
							hds = (struct hd_dev **)malloc(sizeof(struct hd_dev *));
						} else {
							hds = (struct hd_dev **)realloc(hds, sizeof(struct hd_dev *) * (hd_count + 1));
						}

						hds[hd_count] = tmp_hd;
						strcpy(name, "disk");
						k = hd_count;
						if ((hd_count / 100) > 0) {
							name[4] = (k / 100) + '0';
							k = k % 100;
							name[5] = (k / 10) + '0';
							k = k % 10;
							name[6] = k + '0';
							name[7] = '\0';
						} else if ((hd_count / 10) > 0) {
							name[4] = (k / 10) + '0';
							k = k % 10;
							name[5] = k + '0';
							name[6] = '\0';
						} else {
							name[4] = k + '0';
							name[5] = '\0';
						}

						hd_count++;
						
						hd_vdev = (void *)0;
						fstry = 1;
						
						while (hd_vdev == (void *)0 && fstry < 4) {
							hd_vdev = vfs_register_device(0x100 + hd_count - 1, name, fstry);
							fstry++;
						}
						
					}
				}
			}
		}
	}

	free(buffer);
}


char *hd_convert_key(int number, char *buffer) {
	char *ptr;
	int i;
	
	buffer[15] = '\0';
	
	ptr = &buffer[15];
	
	if (number == 0) {
		ptr--;
		*ptr = '0';
		return ptr;
	}
	
	while (number > 0 && ptr != buffer) {
		ptr--;
		i = number % 16;
		*ptr = "fedcba9876543210123456789abcdef"[15 + i];
		number /= 16;
	}
	
	return ptr;
}

int hd_sync_iter(any_t item, any_t data) {
	struct block_buffer *bb = (struct block_buffer *)data;
	int hd = (int)item;
	
	if (bb->dirty == 1) {
		ide_write_sectors(hds[hd]->ide_dev, hds[hd]->drive, 1, bb->lba, 0x10, (unsigned int)bb->buffer);
		bb->dirty = 0;
	}
	
	return MAP_OK;
}

void hd_sync(int hd) {
	if (hd < 0 || hd >= hd_count) {
		return;
	}
	
	hashmap_iterate(hds[hd]->block_buffers, hd_sync_iter, (any_t)hd);
}

int hd_trim_iter(any_t item, any_t data) {
	struct trim_struct *trim = (struct trim_struct *)item;
	struct block_buffer *bb = (struct block_buffer *)data;
	int j;
	
	
	int bbtop;
	
	if (trim->filled < trim->count) {
		if (trim->top == 0 || bb->accessed < trim->top) {
			trim->top = bb->accessed;
		}
		trim->bbs[trim->filled] = bb;
		trim->filled++;
	} else if (bb->accessed < trim->top) {
		trim->top = bb->accessed;
		bbtop = 0;
		for (j=1;j<trim->count;j++) {
			if (trim->bbs[bbtop]->accessed < trim->bbs[j]->accessed) {
				bbtop = j;
			}
		}
			
		trim->bbs[bbtop] = bb;
	}
	
	return MAP_OK;
}

void hd_buffer_trim() {
	int hd;
	struct trim_struct trim;
	int i;

	for (hd=0;hd<hd_count;hd++) {
		if (hashmap_length(hds[hd]->block_buffers) > 20000) {
			trim.count = hashmap_length(hds[hd]->block_buffers) - 20000;
			trim.top = 0;
			trim.bbs = (struct block_buffer **)malloc(sizeof(struct block_buffer *) * trim.count);
			trim.filled = 0;
			for(i=0;i<trim.count;i++) {
				trim.bbs[i] = (void *)0;
			}
			
			hashmap_iterate(hds[hd]->block_buffers, hd_trim_iter, (any_t)&trim);
			
			for(i=0;i<trim.count;i++) {
				if (trim.bbs[i] != (void *)0) {
					hashmap_remove(hds[hd]->block_buffers, trim.bbs[i]->key);
					if (trim.bbs[i]->dirty == 1) {
						ide_write_sectors(hds[hd]->ide_dev, hds[hd]->drive, 1, trim.bbs[i]->lba, 0x10, (unsigned int)trim.bbs[i]->buffer);
					}
					free(trim.bbs[i]);
				} else {
					kprintf("TRIM = NULL o.O\n");
				}
			}		
			
			free(trim.bbs);
		}
	}
}

void hd_write_block(int hd, unsigned int blockno, char *source, unsigned int blocksize) {
	char key[16];	
	struct block_buffer *bb;
	int i;
	
	if (hd < 0 || hd >= hd_count) {
		return;
	}

	int count = blocksize / 512;
	int lba = (blockno * count) + hds[hd]->part_offset;

	if ((blockno * count) + count > hds[hd]->part_len) {
		return;
	}
	
	for (i=0;i<count;i++) {
		if (hashmap_get(hds[hd]->block_buffers, hd_convert_key((lba + (i * 512)), key), &bb) == MAP_OK) {
			bb->dirty = 1;
			bb->accessed++;
			memcpy(bb->buffer, &source[i * 512], 512);
		} else {
			bb = (struct block_buffer *)malloc(sizeof(struct block_buffer));
			bb->lba = lba + (i * 512);
			bb->dirty = 1;
			bb->accessed = 1;
			strcpy(bb->key,  hd_convert_key((lba + (i * 512)), key));
			memcpy(bb->buffer, &source[i * 512], 512);
			hashmap_put(hds[hd]->block_buffers,bb->key, bb);
		}
	}
}


char *hd_read_block(int hd, unsigned int blockno, char *dest, unsigned int blocksize) {
	struct block_buffer *bb;
	char key[16];
	int i;
	
	if (hd < 0 || hd >= hd_count) {
		return (void *)0;
	}

	int count = blocksize / 512;
	int lba = (blockno * count) + hds[hd]->part_offset;

	if ((blockno * count) + count > hds[hd]->part_len) {

		return (void *)0;
	}

	for (i=0;i<count;i++) {
		if (hashmap_get(hds[hd]->block_buffers, hd_convert_key((lba + (i * 512)), key), &bb) == MAP_OK) {
			bb->accessed++;
			memcpy(&dest[i * 512], bb->buffer, 512);
		} else {
			bb = (struct block_buffer *)malloc(sizeof(struct block_buffer));
			bb->lba = lba + (i * 512);
			bb->accessed = 1;
			strcpy(bb->key,  hd_convert_key((lba + (i * 512)), key));
			ide_read_sectors(hds[hd]->ide_dev, hds[hd]->drive, 1, bb->lba, 0x10, bb->buffer);
			memcpy(&dest[i * 512], bb->buffer, 512);
			
			hashmap_put(hds[hd]->block_buffers, bb->key, bb);
		}
	}




	return dest;
}

void hd_write_blocks(int hd, unsigned int blockstart, unsigned int blockcount, char *source, unsigned int blocksize) {
	int i;
	
	if (hd < 0 || hd >= hd_count) {
		return;
	}
	int blocks;

	int count = (blocksize * blockcount) / 512;

	blocks = count;
	while (blockstart * (blocksize / 512) + blocks > hds[hd]->part_len) {
		blocks--;
		if (blocks <= 0) {
			return;
		}
	}

	for (i=0;i<blocks;i++) {
		hd_write_block(hd, blockstart * (blocksize / 512) + i, &source[i * 512], 512);
	}

	return;
}

int hd_read_blocks(int hd, unsigned int blockstart, unsigned int blockcount, char *dest, unsigned int blocksize) {
	int i;
	
	if (hd < 0 || hd >= hd_count) {
		return 0;
	}
	int blocks;

	int count = (blocksize * blockcount) / 512;

	blocks = count;
	while (blockstart * (blocksize / 512) + blocks > hds[hd]->part_len) {
		blocks--;
		if (blocks <= 0) {
			return 0;
		}
	}
	for (i=0;i<blocks;i++) {
		hd_read_block(hd, blockstart * (blocksize / 512) + i, &dest[i * 512], 512);
	}

	return blocks;
}