#include <exec/types.h>
#include <exec/memory.h>
#include <exec/ports.h>
#include <exec/io.h>
#include <devices/trackdisk.h>
#include <dos/rdargs.h>
#include <dos/doshunks.h>
#include <clib/exec_protos.h>
#include <clib/dos_protos.h>

#include <stdio.h>
#include <stdlib.h>

const STRPTR PROG_Name = "BBInstall";
const STRPTR PROG_Ver  = "1";
const STRPTR PROG_Rev  = "7";
const STRPTR PROG_Date = "1st September 1997";

/* Information for C:Version */
const STRPTR VerStr = "$VER: BBInstall 1.7 (1st September 1997)";


const STRPTR BA_Template = "DRIVE/A,LOAD/K,SAVE/K,DOSTYPE/N,CODE/S,MOVESYS/S";

struct BBIArgs {
  STRPTR DRIVE;
  STRPTR LOAD;
  STRPTR SAVE;
  ULONG *DOSTYPE;
  ULONG  CODE;
  ULONG  MOVESYS;
};

struct RDArgs *rdargs;
struct MsgPort *diskport;
struct IOExtTD *diskreq;
BYTE td_open;
ULONG *bb_buffer;
FILE *bbfile;


/* Initialize global data */
void clear_globals(void)
{
  rdargs = NULL;
  diskport = NULL;
  diskreq = NULL;
  td_open = NULL;
  bb_buffer = NULL;
  bbfile = NULL;
}

/* Exit cleanly */
void clean_exit(ULONG rc)
{
  /* Close BootBlock FILE */
  if (bbfile) fclose(bbfile);
  
  /* Free Allocated Memory */
  if (bb_buffer) FreeMem(bb_buffer,1024);
  
  /* Close TrackDisk Unit */
  if (td_open)
  {
    /* Flush any buffers */
    diskreq->iotd_Req.io_Command = CMD_UPDATE;
    DoIO((struct IORequest *)diskreq);
    
    /* Turn off the motor */
    diskreq->iotd_Req.io_Command = TD_MOTOR;
    diskreq->iotd_Req.io_Length = 0;
    DoIO((struct IORequest *)diskreq);
    
    /* Close trackdisk.device */
    CloseDevice((struct IORequest *)diskreq);
  }
  
  /* Delete IO Request */
  if (diskreq) DeleteIORequest(diskreq);
  
  /* Delete Message Port */
  if (diskport) DeleteMsgPort(diskport);
  
  /* Free arguments */
  if (rdargs) FreeArgs(rdargs);
  
  exit(rc);
}


void LoadCode(ULONG *buf, FILE *infile, STRPTR infilename, ULONG moveSysBase)
{
  ULONG read_value;
  ULONG code_size;
  int i;

  /* Find a HUNK_CODE within input file */
  do {
    fread(&read_value,4,1,infile);
  } while ((read_value != HUNK_CODE) && (!feof(infile)));
  
  if (read_value != HUNK_CODE)
  {
    fprintf(stderr,"%s has no code hunk.\n",infilename);
    clean_exit(10);
  }
  
  /* Get size of code */
  fread(&code_size,4,1,infile);
  if (moveSysBase)
  {
    fread(&read_value,4,1,infile);
    
    if (read_value == 0x2C780004)
      code_size--;
    else
      fseek(infile,-4,SEEK_CUR);
  }
  if (code_size > 253)
  {
    fprintf(stderr,"%s is too large to fit into a boot block.\n",infilename);
    clean_exit(10);
  }
  
  /* read code into bootblock array and clear leftover space */
  fread(buf+3,4,code_size,infile);
  for (i=code_size+3; i<256; i++)
    buf[i] = 0;
}

void LoadData(ULONG *buf, FILE *infile)
{
  int i;
  ULONG read_value;
  
  /* Clear buffer */
  for (i=0; i<256; i++)
    buf[i] = 0;
  
  fread(&read_value,4,1,infile);
  
  if ((read_value & 0xFFFFFF00) == 0x444F5300)
  {
    buf[0] = read_value;
    fread(buf+1,4,255,infile);
  }
  else
  {
    buf[3] = read_value;
    fread(buf+4,4,252,infile);
  }
}

void SaveCode(ULONG *buf, FILE *outfile, ULONG moveSysBase)
{
  ULONG exe_head[8] = {HUNK_HEADER,0,1,0,0,253,
                       HUNK_CODE,253};
  ULONG exe_foot[1] = {HUNK_END};
  ULONG exe_sysbase[1] = {0x2C780004};    /* MOVEA.L 4.W,A6 */
  
  if (moveSysBase)
  {
    exe_head[5]++;
    exe_head[7]++;
  }
    
  fwrite(exe_head,4,8,outfile);
  
  if (moveSysBase)
    fwrite(exe_sysbase,4,1,outfile);
 
  fwrite(buf+3,4,253,outfile);
  fwrite(exe_foot,4,1,outfile);
}

void SaveData(ULONG *buf, FILE *outfile)
{
  fwrite(buf,4,256,outfile);
}

ULONG calcsum(ULONG *buf)
{
  ULONG sum,prevsum;
  int i;
  
  sum = buf[0];
  prevsum = 0;
  if (sum < 0)
    sum++;
  
  for (i=2; i<256; i++)
  {
    prevsum = sum;
    sum += buf[i];
    if (prevsum > sum)
      sum++;
  }
  
  return -(sum+1);
}

void OutputBanner(void)
{
  printf("\033[1m%s v%s.%s © %s Karl J. Ots\033[0m\n\n",PROG_Name,
                                                        PROG_Ver,
                                                        PROG_Rev,
                                                        PROG_Date);
}

void OutputUsage(void)
{
  fprintf(stderr,"Usage: BBInstall %s\n",BA_Template);
  clean_exit(5);
}

int main(int argc, char *argv[])
{
  struct BBIArgs BBIArgs = {NULL, NULL, NULL, NULL, FALSE, FALSE};
  int unit;
  ULONG old_dostype;
  
  clear_globals();
  OutputBanner();
  
  /* Collect and analyse arguments */
  rdargs = ReadArgs(BA_Template,(ULONG *)&BBIArgs,NULL);
  if (!rdargs)
    OutputUsage();
  
  if ((BBIArgs.LOAD && BBIArgs.SAVE) ||
      (!BBIArgs.LOAD && !BBIArgs.SAVE))
    OutputUsage();


  /* Determine if DRIVE is valid */
  if (((BBIArgs.DRIVE[0] != 'D') && (BBIArgs.DRIVE[0] != 'd')) ||
      ((BBIArgs.DRIVE[1] != 'F') && (BBIArgs.DRIVE[1] != 'f')) ||
       (BBIArgs.DRIVE[2] >= '4') ||
       (BBIArgs.DRIVE[3] != ':'))
  {
    fprintf(stderr,"Invalid Drive %s\n",BBIArgs.DRIVE);
    clean_exit(5);
  }
  unit = BBIArgs.DRIVE[2] - '0';
  

  /* Open TrackDisk Unit */
  diskport = CreateMsgPort();
  if (!diskport)
  {
    fputs("Error: Couldn't create a Message Port (Out of Memory?)\n",stderr);
    clean_exit(20);
  }
  
  diskreq = (struct IOExtTD *)
            CreateIORequest(diskport,sizeof(struct IOExtTD));
  if (!diskreq)
  {
    fputs("Error: Couldn't create an IO request (Out of Memory?)\n",stderr);
    clean_exit(20);
  }
  
  td_open = !OpenDevice(TD_NAME,unit,(struct IORequest *)diskreq,NULL);
  if (!td_open)
  {
    fprintf(stderr,"Can't open trackdisk.device unit %d (DF%d:)!\n",unit,
                                                                    unit);
    clean_exit(20);
  }
  
  
  /* Allocate some memory */
  bb_buffer = (ULONG *)AllocMem(1024,MEMF_CLEAR|MEMF_CHIP);
  if (!bb_buffer)
  {
    fputs("Error: Out of memory.\n",stderr);
    clean_exit(10);
  }
  
  /* Read TrackDisk Unit */
  diskreq->iotd_Req.io_Length = 1024;
  diskreq->iotd_Req.io_Data = bb_buffer;
  diskreq->iotd_Req.io_Command = CMD_READ;
  diskreq->iotd_Req.io_Offset = 0;
  if (DoIO((struct IORequest *)diskreq))
  {
    fputs("An error occured\n",stderr);
    clean_exit(20);
  }
  
  old_dostype = bb_buffer[0];
    
  if (BBIArgs.LOAD)
  {
    bbfile = fopen(BBIArgs.LOAD,"rb");
    if (!bbfile)
    {
      fprintf(stderr,"Couldn't open \"%s\" for reading\n",BBIArgs.LOAD);
      clean_exit(10);
    }
    printf("Loading \"%s\" as %s to DF%d: ... \n",BBIArgs.LOAD,
                                            (BBIArgs.CODE ? "code" : "data"),
                                                  unit);
    
    /* Setup bootblock header (without checksum) */
    /* bb_buffer[0] = bb_buffer[0] */ /* Obviously... */
    bb_buffer[1] = 0;
    bb_buffer[2] = 0x00000370;
    
    if (BBIArgs.CODE)
      LoadCode(bb_buffer,bbfile,BBIArgs.LOAD,BBIArgs.MOVESYS);
    else
      LoadData(bb_buffer,bbfile);
    
    if (BBIArgs.DOSTYPE)
      bb_buffer[0] = (bb_buffer[0] & 0xFFFFFF00) |
                     (BBIArgs.DOSTYPE[0] & 0xFF);
    printf("  Dostype = 0x%02X ... ",(bb_buffer[0] & 0xFF));
    fflush(stdout);
    
    /* Calculate checksum and insert it into bootblock */
    bb_buffer[1] = calcsum(bb_buffer);
    printf("Checksum = 0x%08X ... ",bb_buffer[1]);
    fflush(stdout);
    
    /* bb_buffer in now ready to be inserted into bootblock */
    /* Write TD Unit */
    diskreq->iotd_Req.io_Length = 1024;
    diskreq->iotd_Req.io_Data = bb_buffer;
    diskreq->iotd_Req.io_Command = CMD_WRITE;
    diskreq->iotd_Req.io_Offset = 0;
    if (DoIO((struct IORequest *)diskreq))
    {
      fputs("Another error occured\n",stderr);
      clean_exit(20);
    }
  
    puts("Done.\n");
  }
  else  /* if (BBIArgs.SAVE) */ /* <-- We know this is set if LOAD isn't */
  {
    bbfile = fopen(BBIArgs.SAVE,"wb");
    if (!bbfile)
    {
      fprintf(stderr,"Couldn't open \"%s\" for writing.\n",BBIArgs.SAVE);
      clean_exit(10);
    }
    printf("Saving old BootBlock from DF%d: to \"%s\" as %s ... ",unit,
                                                             BBIArgs.SAVE,
                                           (BBIArgs.CODE ? "code" : "data"));
    fflush(stdout);
        
    if (BBIArgs.CODE)
      SaveCode(bb_buffer,bbfile,BBIArgs.MOVESYS);
    else
      SaveData(bb_buffer,bbfile);
    
    puts("Done.\n");
  }
  
  clean_exit(0);
  
  return 0;	/* Keep compiler happy */
}