#include <stdio.h>      /* standard C i/o facilities */
#include <stdlib.h>     /* needed for atoi() */
#include <unistd.h>  	/* defines STDIN_FILENO, system calls,etc */
#include <sys/types.h>  /* system data type definitions */
#include <sys/socket.h> /* socket specific definitions */
#include <netinet/in.h> /* INET constants and stuff */
#include <arpa/inet.h>  /* IP address conversion stuff */
#include <netdb.h>	/* gethostbyname */
#include <sys/time.h>


/* code for building messages */

enum opcodes {RRQ=1,WRQ=2,DATA=3,ACK=4,ERR=5};

/* build rrq or wrq messages */
/* assumes that msg can hold at least 516 bytes! */
/* returns size of message or -1 */

int build_request(char *msg, unsigned short opcode, char *filename, char *mode) {

  if (strlen(filename)>100) {
    printf("File name too long\n");
    return(-1);
  }

  if (strlen(mode)>100) {
    printf("Transfer mode too long\n");
    return(-1);
  }

  if ((opcode != RRQ) && (opcode != WRQ)) {
    printf("build_request: invalid opcode (%d)\n",opcode);
    return(-1);
  }

  opcode=htons(opcode);
  memcpy(msg,&opcode,2);
  strcpy(msg+2,filename);
  strcpy(msg+2+strlen(filename)+1,mode);
  return(2+strlen(filename)+1+strlen(mode)+1);
}


/* build ack message */

int build_ack(char *msg, unsigned short block) {
  unsigned short opcode;

  opcode=htons(ACK);
  memcpy(msg,&opcode,2);
  block = htons(block);
  memcpy(msg+2,&block,2);
  return(4);
}

unsigned short get_shortval( char *p ) {
  unsigned short tmp;

  memcpy(&tmp,p,2);
  tmp = ntohs(tmp);
  return(tmp);
}

char *opcode_strings[] = {  "RRQ", "WRQ", "DATA", "ACK", "ERR" };

int get_opcode(char *msg) {
  unsigned short opcode;
  opcode = get_shortval(msg);
  return(opcode);
}


int get_blocknum(char *msg) {
  unsigned short blocknum;
  blocknum = get_shortval(msg+2);
  return(blocknum);
}

/* print out a TFTP message */
void print_msg( char *msg, int len ) {

  unsigned short opcode,blocknum,errcode;

  opcode = get_shortval(msg);

  if (opcode> ERR) {
    printf("Error - invalid opcode: %d\n",opcode);
    return;
  }

  printf("%s: ",opcode_strings[opcode-1]);

  switch(opcode) {
  case RRQ:
  case WRQ:
    printf("<%s> <%s>\n",msg+2,msg+2+strlen(msg+2)+1);
    break;
  case DATA:
    blocknum = get_shortval(msg+2);
    printf("BLOCK: %d (%d bytes)\n",blocknum,len-4); 
    break;
  case ACK:
    blocknum = get_shortval(msg+2);
    printf("BLOCK: %d\n",blocknum); 
    break;
  case ERR: 
    errcode = get_shortval(msg+2);
    printf("CODE: %d <%s>\n",errcode,msg+4);
    break;
  }

}




int wait_for_response(int sk,struct sockaddr_in *s,char *buff, int max) {
  int len=max;
  fd_set rfds;
  struct timeval t = { 2,0 }; /* 2 seconds */ 
  int received;

  FD_ZERO(&rfds);
  FD_SET(sk,&rfds);
  if (select(sk+1,&rfds,NULL,NULL,&t)==0) {
    return(0);
  } else if (FD_ISSET(sk,&rfds)) {
    received = recvfrom(sk,buff,1000,0,(struct sockaddr*)s,&len);
    if (received<0) {
      perror("Error calling recvfrom\n");
      return(0);
    }
    return(received);
  }
  return(0);
}