#include <stdio.h>
#include <string.h>
#include <unistd.h>
#include <errno.h>
#include <stdlib.h>
#include <misc.h>
#include <popen.h>
#include <dialog.h>
#include <tlmplib.h>
#include <tlmpnet.h>
#include "fastscp.m"
#include <sys/time.h>
#include <sys/stat.h>
#include "fastscp.h"
#include <string>

using namespace std;

static long long getnow ()
{
	struct timeval tv;
	gettimeofday (&tv,NULL);
	return tv.tv_sec *(long long)1000000 + tv.tv_usec;
}


struct STATS{
	const char *fname;
	long long size;		// Size of the file if known
	long long sent;		// Bytes sent so far
	long long start;	// Transmission start time
	struct {
		long long now;	// Time of the last report
		long long sent;
	}last;
	long long statrate;
	void init (const char *_fname, long long _size){
		fname = _fname;
		size = _size;
		sent = 0;
		start = getnow();
		last.now = start;
		last.sent = 0;
		statrate = 1024*1024;
	}
	void print(long long _sent, long long start, long long now);
	void print_start();
	void print_end();
	void printif (int _sent);
};

static bool silent = false;

static void print_head (bool receiving)
{
	if (!silent){
		fprintf (stderr,"%-20s %9s %9s %9s %4s %10s %9s\n"
			,MSG_U(I_FILE,"File name")
			,MSG_U(I_SIZE,"Size(k)")
			,receiving ? MSG_U(I_RECV,"Recv(k)") : MSG_U(I_SENT,"Sent(k)")
			,MSG_U(I_RATE,"Rate(m/s)"),MSG_U(I_PERCENT,"%")
			,MSG_U(I_DURATION,"Duration")
			,MSG_U(I_ETA,"Eta"));
	}
}

void STATS::print (
	long long _sent,	// Amount sent since the last print
	long long _start,	// timestamp of the last print
	long long now)		// Current timestamp
{
	if (!silent){
		long long diff_time = now - _start;
		long long rate = 0;
		if (diff_time > 0){
			rate = (_sent * 1000000) / diff_time;
		}
		long long eta = 0;
		if (sent > 0){
			eta = (long long)((now - start) * ((double)size/sent) / 1000000);
		}
		long percent = 0;
		if (size != 0){
			percent = (sent*100)/size;
		}
		double frate = (double)rate/(1024*1024);
		fprintf (stderr,"%-20.20s %9Ld %9Ld %9.3lf %3d%% %10Ld %9Ld\r"
			,fname,size/1024,sent/1024,frate
			,percent
			,(now-start)/1000000
			,eta);
		fflush (stderr);
	}
}

void STATS::print_start ()
{
	print (0,start,start);
}

void STATS::printif (int _sent)
{
	sent += _sent;
	long long diff = sent - last.sent;
	if (diff >= statrate){
		long long now = getnow();
		print (diff,last.now,now);
		last.sent = sent;
		// Avoid printing stats too fast
		if (now - last.now < 100000){
			statrate *= 10;
			// printf ("new statrate %Ld\n",statrate);
		}
		last.now = now;
	}
}

/*
	Print last report (when the file has been completly sent)
*/
void STATS::print_end()
{
	if (!silent){
		print (sent,start,getnow());
		fprintf (stderr,"\n");
	}
}
/*
	Decompose an SCP user@server:/file syntax.
	Leave the ssh variable as user@server or just server.
	host will contain only server.
	file will contain only the file or nothing
*/
static void fastscp_parsessh (
	string &ssh,
	string &host,
	string &file)
{
	auto pt = ssh.find(':');
	if (pt != string::npos){
		file = ssh.substr(pt + 1);
		ssh = ssh.substr(0,pt);
	}
	pt = ssh.find('@');
	if (pt != string::npos){
		host = ssh.substr(pt+1);
	}else{
		host = ssh;
	}
}

#define CRBUF_SIZE 20000
static unsigned long long crbuf[CRBUF_SIZE];
static unsigned long long *ptcr = crbuf;
static unsigned long long *endcr = crbuf+CRBUF_SIZE;
static bool docrypt = false;

static void fastscp_crypt (char *buf, int n)
{
	if (docrypt){
		int modn = n%8;
		if (modn != 0) n += 8-modn;
		unsigned long long *endbuf = (unsigned long long*)(buf+n);
		for (unsigned long long *ptbuf = (unsigned long long*)buf; ptbuf < endbuf; ptbuf++){
			*ptbuf ^= *ptcr++;
			if (ptcr == endcr) ptcr = crbuf;
		}
	}
}

static void fastscp_cryptinit (FILE *fout)
{
	if (docrypt){
		FILE *fin = fopen ("/dev/urandom","r");
		if (fin == NULL){
			tlmp_error (MSG_U(E_URANDOM,"Can't open /dev/urandom (%s), can't crypt\n")
				,strerror(errno));
			exit (-1);
		}else{
			if (fread(crbuf,1,sizeof(crbuf),fin)!=sizeof(crbuf)){
				tlmp_error (MSG_U(E_READURANDOM,"Can't read /dev/urandom (%s), can't crypt\n")
					,strerror(errno));
				exit (-1);
			}else{
				// Now we must transmit the crypt buffer to the backend
				// We always send the full crypt buffer (full 100000)
				// even if we are using less. This way, an attacker spying
				// the ssh packets can't figure out the size of the
				// buffer. So it can't easily figure the period of the
				// crypt loop.
				for (unsigned i=0; i<CRBUF_SIZE; i+= 20){
					fprintf (fout,"crbuf ");
					for (unsigned j=i; j<i+20; j++){
						fprintf (fout,"%016Lx",crbuf[j]);
					}
					fprintf (fout,"\n");
				}
				// We compute the length of the buffer
				// based on the first long long word
				endcr = crbuf + (CRBUF_SIZE/2 + crbuf[0]%(CRBUF_SIZE/2));
				// fprintf (stderr,"Crypt buf size %u\n",endcr-crbuf);
			}
			fclose (fin);
			fflush (fout);
		}
	}
}

/*
	Receive one file
*/
static void fastscp_receive (
	const char *fname,
	FILE *fout,
	int netfd,
	long long &received)
{
	STATS stsin;
	stsin.init (fname,0);
	stsin.print_start();

	received = 0;
	char buf[1024*1024];
	int n;
	while ((n=read(netfd,buf,sizeof(buf)))>0){
		received += n;
		fastscp_crypt (buf,n);
		fwrite (buf,1,n,fout);
		stsin.printif (n);
	}
	if (fout != stdout) fclose (fout);
	stsin.print_end();
}
	
/*
	Send one file over the fast socket
*/
static void fastscp_send (const string &pathname, FILE *fout, int netfd)
{
	const char *fname = pathname.c_str();
	if (file_type(fname)!=0){
		tlmp_error (MSG_U(E_NOTAFILE,"%s is not a file\n"),fname);
		return;
	}
	FILE *fin = strcmp(fname,"-")==0 ? stdin : fopen64 (fname,"r");
	if (fin == NULL){
		tlmp_error (MSG_U(E_CANTOPEN,"Can't open file %s (%s)\n")
			,fname,strerror(errno));
	}else{
		const char *pt = strrchr(fname,'/');
		if (pt == NULL){
			pt = fname;
		}else{
			pt++;
		}
		fprintf (fout,"file %s\n",pt);
		fflush (fout);
		long long size = 0;
		if (fin != stdin){
			struct stat64 st;
			if (fstat64(fileno(fin),&st)!=-1){
				size = st.st_size;
			}
		}
		char buf[1024*1024];
		int n;
		STATS sts;
		sts.init (fname,size);
		sts.print_start();
		while ((n=fread(buf,1,sizeof(buf),fin))>0){
			fastscp_crypt (buf,n);
			write (netfd,buf,n);
			sts.printif (n);
		}
		sts.print_end();
		if (fin != stdin) fclose (fin);
		fprintf (fout,"sent %Ld\n",sts.sent);
		fflush (fout);
	}
}
	

<mod>
int main (int argc, char *argv[])
{
	glocal int ret = -1;
	glocal bool remdebug = false;
	glocal bool remdirect = false;
	glocal bool remfadvise = false;
	glocal bool recurse = false;
	glocal.ret = <call tlmpprogram>(argc,argv,"fastscp");
	<f init>
		setproginfo ("fastscp",VERSION
			,MSG_U(I_FASTSCP
				,"Very fast copy partially over ssh\n"
				 "\n"
				 "fastscp file [file ...] user@host:/dir[/filename]\n"
				 "fastscp user@host:/dir/filename file\n"
				 "command | fastscp -- - user@host:/dir/filename\n"
				 "fastscp user@host:/dir/filename - | command\n"
			));
		setarg ('c',"crypt",MSG_U(I_CRYPT,"Crypt the fast channel"),docrypt,false);
		setarg ('q',"quiet",MSG_U(I_QUIET,"Silent mode"),silent,false);
		setarg ('r',"recurse",MSG_U(I_RECURSE,"Copy directory tree"),glocal.recurse,false);
		setarg (' ',"remotedebug",MSG_U(I_REMDEBUG,"Produce debugging on the remote side"),glocal.remdebug,false);
		//setarg ('o',"remotedirect",MSG_U(I_REMDIRECT,"Utilise O_DIRECT pour ecrire les fichiers"),glocal.remdirect,false);
		setarg ('a',"remotefadvise",MSG_U(I_REMADVISE,"Utilise posix_fadvise pour ecrire les fichiers"),glocal.remfadvise,false);
	</f>
	<f main>
		glocal int ret = -1;
		if (argc < 2){
			usage();
		}else{
			bool ok = true;
			glocal char **argv = argv;
			glocal string host;
			glocal int argc = argc;
			glocal string file;
			glocal bool receiving = false;
			glocal FILE *fileout = NULL;
			string ssh = argv[0];
			auto pt = ssh.find(':');
			if (pt != string::npos){
				// We are reading from the ssh server
				if (argc != 2){
					ok = false;
					usage();
				}else{
					fastscp_parsessh (ssh,glocal.host,glocal.file);
					glocal.receiving = true;
	
					const char *arg = argv[argc-1];
					glocal.fileout = strcmp(arg,"-")==0
						? stdout : fopen64 (arg,"w");
					if (glocal.fileout == NULL){
						tlmp_error (MSG_R(E_CANTOPEN),arg,strerror(errno));
						ok = false;
					}
				}
			}else{
				argc--;
				glocal.argc = argc;
				ssh = argv[argc];

				fastscp_parsessh (ssh,glocal.host,glocal.file);
			}
			if (ok){
				glocal int netfd = -1;	// Socket to send/receive data
				glocal int argno = 0;	// File to transmit
				glocal long long received = 0;
				glocal string walk_path;
				glocal string sent_target_dir;
				glocal FILE *fout = NULL;
				<obj COROUTINE walk>();
				<f run>
					glocal.sent_target_dir = glocal.file;
					fprintf (glocal.fout,"dir %d %s\n"
						,glocal.argc > glocal.argno + 1,glocal.file.c_str());
					while (glocal.argno < glocal.argc){
						const char *arg = glocal.argv[glocal.argno++];
						if (file_type(arg)==1){
							if (!glocal.recurse){
								tlmp_error (MSG_U(E_ISDIR,"Argument %s is a directory, but option -r not active, ending\n"),arg);
								exit (-1);
							}
							<call walkfs>(arg);
							<f onefile>
								if (file_type(path)==0){
									glocal.walk_path = path;
									string target_dir;
									const char *pt = strrchr(relpath,'/');
									if (pt!=NULL){
										target_dir = glocal.file + "/" + string(relpath,pt-relpath);
									}else{
										target_dir = glocal.file;
									}
									if (target_dir != glocal.sent_target_dir){
										glocal.sent_target_dir = target_dir;
										//printf ("target_dir = %s\n",glocal.sent_target_dir.c_str());
										fprintf (glocal.fout,"dir 1 %s\n",glocal.sent_target_dir.c_str());
									}
									if(glocal.COROUTINE.yield()==-1) end();
								}
							</f>
							<f recurse>
								return true;
							</f>
							</call>
						}else{
							glocal.walk_path = arg;
							if(yield()==-1) break;
						}
					}
				</f>
				</obj>
				glocal COROUTINE *walk = &walk;
				string cmd = string_f("ssh %s /usr/lib/tlmp/fastend",ssh.c_str());
				<call walkpopen>(cmd,1000);
				<f oneerr>
					tlmp_error (MSG_U(E_ERRBACKEND,"Back end error: %s\n"),line);
					return 0;
				</f>
				<f oneline>
					int ret = -1;
					//fprintf (stderr,"host :%s: port :%s:\n",glocal.host.c_str(),line);
					glocal.fout = fout;
					if (glocal.netfd == -1){
						char verline[100];
						int len = snprintf (verline,sizeof(verline),"VERSION%d "
							,FASTSCP_PROTO_VERSION);
						if (strncmp(line,verline,len)!=0){
							tlmp_error (MSG_U(E_PROTOVER,"Invalid backend version. Expect %s, got %s\n"),verline,line);
						}else{
							const char *port = line + len;
							glocal.netfd = cmdsock_connect (glocal.host.c_str(),port,100,5);
							if (glocal.netfd == -1){
								tlmp_error (MSG_U(E_CANTCONNECT
									,"Can't connect to fast data port %s\n")
									,port);
							}else{
								if (glocal.remdebug) fprintf (glocal.fout,"debug\n");
								if (glocal.remdirect) fprintf (glocal.fout,"odirect\n");
								if (glocal.remfadvise) fprintf (glocal.fout,"fadvise\n");
								fastscp_cryptinit (fout);
								print_head(glocal.receiving);
								if (glocal.receiving){
									fprintf (glocal.fout,"get %s\n",glocal.file.c_str());
									fflush  (glocal.fout);
									fastscp_receive (glocal.file.c_str()
										,glocal.fileout,glocal.netfd,glocal.received);
									ret = 0;
								}else{
									if (glocal.walk->next()){
										fastscp_send (glocal.walk_path,glocal.fout,glocal.netfd);
									}
									ret = 0;
								}
							}
						}
					}else if (strcmp(line,"recok")==0){
						if (glocal.walk->next()){
							fastscp_send (glocal.walk_path,glocal.fout,glocal.netfd);
							ret = 0;
						}else{
							glocal.ret = 0;
						}
					}else if (strncmp(line,"sent ",5)==0){
						long long sent;
						sscanf (line+5,"%Ld",&sent);
						if (sent == glocal.received) glocal.ret = 0;
					}else{
						tlmp_error (MSG_U(E_BKMSG,"Unknown backend message: %s\n")
							,line);
					}
					return ret;
				</f>
				</call>
				// fprintf (stderr,"end walkpopen\n");
			}
		}
		return glocal.ret;
	</f>
	</call>
	return glocal.ret;
}
</mod>