#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <gccore.h>
#include <wiiuse/wpad.h>
#include <fat.h>
#include "wdvd.h"


static void *xfb = NULL;
static GXRModeObj *rmode = NULL;
bool fatinited = false;

void init()
{
	IOS_ReloadIOS(5);
	
	// Initialise the video system
	VIDEO_Init();
	
	// This function initialises the attached controllers
	WPAD_Init();
	
	// Obtain the preferred video mode from the system
	// This will correspond to the settings in the Wii menu
	rmode = VIDEO_GetPreferredMode(NULL);

	// Allocate memory for the display in the uncached region
	xfb = MEM_K0_TO_K1(SYS_AllocateFramebuffer(rmode));
	
	// Initialise the console, required for printf
	console_init(xfb,20,20,rmode->fbWidth,rmode->xfbHeight,rmode->fbWidth*VI_DISPLAY_PIX_SZ);
	
	// Set up the video registers with the chosen mode
	VIDEO_Configure(rmode);
	
	// Tell the video hardware where our display memory is
	VIDEO_SetNextFramebuffer(xfb);
	
	// Make the display visible
	VIDEO_SetBlack(FALSE);

	// Flush the video register changes to the hardware
	VIDEO_Flush();

	// Wait for Video setup to complete
	VIDEO_WaitVSync();
	if(rmode->viTVMode&VI_NON_INTERLACE) VIDEO_WaitVSync();


	// The console understands VT terminal escape codes
	// This positions the cursor on row 2, column 0
	// we can use variables for this with format codes too
	// e.g. printf ("\x1b[%d;%dH", row, column );
	printf("\x1b[2;0H");

}


bool initdvd()
{
	u64 id;
	unsigned char idstring[9];

	printf("Initializing DVD...\n");

	if (WDVD_Init() == false)
	{
		printf("Error initializing DVD\n");

		return false;
	}

	if (WDVD_Reset() == false)
	{
		printf("Error resetting DVD\n");

		return false;
	}

	if (WDVD_LowReadDiskId(&id) != 1)
	{
		printf("Error reading disc ID\n");

		return false;
	}

	memset(idstring, 0, sizeof(idstring));
	memcpy(idstring, &id, 8);
	

	printf("Disc ID: %s\n", idstring);

	return true;
}

#define DUMP_DVD 0
#define BENCH_DVD 1
#define BENCH_SD 2

bool dumpdvd(unsigned int part, unsigned long partsize, unsigned int dvdtype, unsigned int benchmark)
{
	char filename[255];
  FILE *f;
  
  unsigned long DVD_TYPE[3] = {
    0x46090000, //Single Layer Wii DVD
    0x7ed38000, //Dual Layer Wii DVD
    0x15C16000, //Gamecube DVD
  };
  
  unsigned long PART_SIZE[3] = {
    1500 * 1024 * 1024,
    750  * 1024 * 1024,
    375  * 1024 * 1024
  };
  
  if(benchmark == BENCH_DVD)
  {
    printf("Benchmarking DVD Drive\n");
  }
  else if (benchmark == BENCH_SD)
  {
    printf("Benchmarking SD Card\n");
  }
  
  if(benchmark != BENCH_DVD)
  {
  	if (!fatinited)
  	{
  		printf("Initializing libfat...\n");
  		if(!fatInitDefault())
  		{
  			printf("Error cannot initialize libfat. Is SD card inserted?\n");

  			return false;
  		}

  		fatinited = true;
  	}
  }
  

	if (!initdvd())
	{
		return false;
	}

  if(benchmark == DUMP_DVD)
  {
  	sprintf(filename, "/%d.bin", part);
  	printf("Writing file %s to SD\n", filename);

  	f = fopen(filename, "wb");

  	if (f == 0)
  	{
  		printf("Error opening file\n");

  		return false;
  	}
  }
  else if(benchmark == BENCH_SD)
  {
    sprintf(filename, "/benchmark.bin");
    printf("Writing file %s to SD\n", filename);

  	f = fopen(filename, "wb");

  	if (f == 0)
  	{
  		printf("Error opening file\n");

  		return false;
  	}
  }
  

	unsigned long offset = 0;
	unsigned long size = 0;

	unsigned int maxsize = 50 * 1024;
	int countdown = 0;
	unsigned char* buffer;
  //Size of SL DVD = 4699979776 Bytes
  if(benchmark != BENCH_SD)
  {
    offset = (PART_SIZE[partsize] >> 2) * part;
    size = PART_SIZE[partsize];
    if((offset + (size >> 2))>DVD_TYPE[dvdtype])
    {
      size = DVD_TYPE[dvdtype] - offset;
      size <<= 2;
    }
  }
  else
  {
    offset = 0x2EE00000;
    size = 10 * 1024 * 1024;
  }


	unsigned long beginoffset = offset;
	unsigned long endoffset = offset + size/4;
	unsigned long totaloffset = size/4;
	int tries = 5;
	
	while(size != 0)
	{
		int currentsize = size>maxsize? maxsize : size;
		int result = 0;


		WPAD_ScanPads();

		u32 pressed = WPAD_ButtonsDown(0);

		if ( pressed & WPAD_BUTTON_HOME )
		{
      if(benchmark == DUMP_DVD)
			  printf("\nStopping dump... hit home again to exit or press A to continue\n");
      else
        printf("\nStopping Benchmark... hit home again to exit or press A to continue\n");
      
      WPAD_ScanPads();
      pressed = WPAD_ButtonsDown(0);
      while(!(pressed & WPAD_BUTTON_A) && !(pressed & WPAD_BUTTON_HOME))
      {
        VIDEO_WaitVSync();
        WPAD_ScanPads();
        pressed = WPAD_ButtonsDown(0);
      }
      if(pressed & WPAD_BUTTON_HOME)
      {
        exit(0);
  }
      
      if(benchmark != BENCH_DVD)
			  fclose(f);

			return false;
		}

tryagain:

		result = WDVD_LowUnencryptedRead(&buffer, currentsize, offset);

		if (result != 1)
		{
			if (tries)
			{
				tries--;
				goto tryagain;
			}


			printf("Error reading sector data\n");

      if(benchmark != BENCH_DVD)
			  fclose(f);
			return false;
		}

		offset += currentsize/4;
		size -= currentsize;
		tries = 5;

    if(benchmark != BENCH_DVD)
    {
  		if (fwrite(buffer, sizeof(unsigned char), currentsize, f) != currentsize)
  		{
  			printf("Error writing to SD\n");
  			fclose(f);
  			return false;
  		}
    }

		if (countdown == 0)
		{
			countdown = 10;

			unsigned long diffoffset = offset - beginoffset;
			unsigned long progress = diffoffset  / (totaloffset / 100);

			printf("\r     %lu%% / %08lu / %08lu / %08lu", progress, diffoffset, offset, endoffset);
		}
		else
		{
			countdown--;
		}
	}


	if (size == 0)
	{
		unsigned long diffoffset = offset - beginoffset;
		unsigned long progress = diffoffset  / (totaloffset / 100);

		printf("\r     %lu / %08lu / %08lu / %08lu", progress, diffoffset, offset, endoffset);
	}

  if(benchmark != BENCH_DVD)
	  fclose(f);

  if(benchmark == DUMP_DVD)
	  printf("\nSUCCESS - Written part %d.bin to SD\n", part);
  else
    printf("\nBENCHMARK COMPLETE\n");
  printf("\nPress A to Continue or Home to exit\n");
  
  WPAD_ScanPads();
  u32 pressed = WPAD_ButtonsDown(0);
  while(!(pressed & WPAD_BUTTON_A) && !(pressed & WPAD_BUTTON_HOME))
  {
    VIDEO_WaitVSync();
    WPAD_ScanPads();
    pressed = WPAD_ButtonsDown(0);
  }
  if(pressed & WPAD_BUTTON_HOME)
  {
    exit(0);
  }

	return true;
}


void mainloop()
{
  unsigned char gamedisctype[3][13] = {
    "Wii SL DVD",
    "Wii DL DVD",
    "Gamecube DVD"
  };
    
  unsigned int sizes[3] = {
    1500, 750, 375
  };
    
  unsigned long part_number=0;
  unsigned long max_part[3][3] = {
    { 2 , 5 , 0 },  //1500 MiB, Wii SL, Wii DL, GC
    { 5 , 11, 1 },  //750 MiB
    { 11, 23, 3 },  //375 MiB
  };
  unsigned long dvd_type=0;
  unsigned long part_size=0;
  int i;
  
	printf("Wii disc dumper v1.2 - nitrotux, CaitSith2\n");
		
	printf("\n");
  printf("Hit 1 to Set part Number\n");
  printf("Hit 2 to Set part Size\n");
  printf("Hit B to Set Game Disc type\n");
  
  printf("Hold + (and hit A/1/2) to benchmark DVD drive\n");
  printf("Hold - (and hit A/1/2) to benchmark SD Card\n");

	while(1) 
	{
    WPAD_ScanPads();

		u32 pressed = WPAD_ButtonsDown(0);
    u32 held = WPAD_ButtonsHeld(0);
    held &= (WPAD_BUTTON_PLUS | WPAD_BUTTON_MINUS);
    
    printf ("\x1b[%d;%dH", 3, 0 );
    printf ("                                                                               \r");
    if(held == 0)
      printf("Hit A to dump %s part %lu, size %d MiB to SD\n", &gamedisctype[dvd_type][0], part_number, sizes[part_size]);
    else if(held == WPAD_BUTTON_PLUS)
      printf("Hit A to benchmark DVD Drive, %s part %lu, size %d MiB\n", &gamedisctype[dvd_type][0], part_number, sizes[part_size]);
    else if(held == WPAD_BUTTON_MINUS)
      printf("Hit A to benchmark SD card\n");
    for(i=0;i<10;i++)
    {
      printf("\x1b[%d;%dH",9+i,0);
      printf ("                                                                               \r");
    }
    printf ("\x1b[%d;%dH", 9, 0 );
		

		if ( pressed & WPAD_BUTTON_HOME )
			exit(0);

		if ( pressed & WPAD_BUTTON_A )
    {
      if ( held & WPAD_BUTTON_PLUS )
        dumpdvd(part_number,part_size,dvd_type,BENCH_DVD);
      else if (held & WPAD_BUTTON_MINUS )
        dumpdvd(part_number,part_size,dvd_type,BENCH_SD);
      else
			  dumpdvd(part_number,part_size,dvd_type,DUMP_DVD);
    }
    
    if (pressed & WPAD_BUTTON_B )
    {
      dvd_type++;
      if(dvd_type==3)
        dvd_type=0;
      part_number=0;
    }

		if ( pressed & WPAD_BUTTON_1 )
		{
      part_number++;
      if(part_number>max_part[part_size][dvd_type])
        part_number=0;
    }

		if ( pressed & WPAD_BUTTON_2 )
		{
      part_size++;
      if(part_size==3)
        part_size=0;
      part_number=0;
    }


		VIDEO_WaitVSync();
	}


	while(1) 
	{
		VIDEO_WaitVSync();
	}
}

//---------------------------------------------------------------------------------
int main(int argc, char **argv) 
{
	init();

	mainloop();

	return 0;
}
