#include <stdio.h>
#include <string.h>
#include <stdlib.h>
#include <errno.h>
#include <syslog.h>
#include <misc.h>
#include <fcntl.h>
#include <unistd.h>
#include <linux/types.h>
#include <linux/nbd.h>
#include <netinet/in.h>
#include <tlmplib.h>
#include <tlmpnet.h>

#include "drsnap.h"

extern const char *version;

static DEBUG_KEY DRSNAP ("drsnap","dnsnap program");

#ifdef WORDS_BIGENDIAN
static unsigned long long ntohll(unsigned long long a) {
        return a;
}
#else
static unsigned long long ntohll(unsigned long long a) {
        unsigned lo = a & 0xffffffff;
        unsigned hi = a >> 32U;
        lo = ntohl(lo);
        hi = ntohl(hi);
        return ((unsigned long long) lo) << 32U | hi;
}
#endif
#define htonll ntohll


static SSTRING branch ("prod");

struct STATS{
	int reads,writes;
	unsigned long long bytes_read,bytes_written;
	STATS(){
		reads = writes = 0;
		bytes_read = bytes_written = 0;
	}
};

struct CLIENTINFO: public ARRAY_OBJ {
	union {
		nbd_request req;
		char breq[28];
	}u;
	int sofar;		// How much was received for writing
	int inheader;	// How many bytes of the header was received
	STATS stats;
	SNAPFILE *sn;
	bool dev_null;
	CLIENTINFO(const char *_fname){
		sn = snapfile_find (branch.c_str());
		if (sn != NULL){
			if (sn->is_inuse()){
				tlmp_error ("Data file %s already in use",sn->getpath());
				sn = NULL;
			}else{
				sn->set_inuse();
			}
		}
		sofar = 0;
		memset (u.breq,0,sizeof(u.breq));
		inheader = 0;
		dev_null = false;
	}
	~CLIENTINFO(){
		if (sn != NULL) sn->unset_inuse();
	}
	bool is_ok() const
	{
		return sn != NULL;
	}
	void sendok (int tofd){
		nbd_reply rep;
		rep.magic = ntohl(NBD_REPLY_MAGIC);
		rep.error = 0;
		memcpy (rep.handle,u.req.handle,sizeof(rep.handle));
		::write (tofd,&rep,sizeof(rep));
	}
	int read (int tofd){
		int ret = -1;
		if (sn->checkrange(u.req.from,u.req.len)!=-1){
			sendok (tofd);
			ret = sn->read (u.req.from,u.req.len,tofd);
		}
		return ret;
	}
	int write (const char *buf, int len){
		int ret = -1;
		if (sn->checkrange(u.req.from,len)!=-1){
			ret = sn->write (u.req.from,buf,len);
		}
		return ret;
	}
	unsigned long long getsize() const
	{
		return sn->getsize();
	}
	int snap()
	{
		int ret = -1;
		SNAPFILE *newsn = sn->snap();
		if (newsn != NULL){
				sn->unset_inuse();
				newsn->set_inuse();
				sn = newsn;
				ret = 0;
		}
		return ret;
	}
	int split(const char *newbranch)
	{
		int ret = -1;
		if (sn->snap('B',newbranch) != NULL) ret = snap();
		return ret;
	}
	const char *getpath() const
	{
		return sn->getpath();
	}
};

#define INIT_PASSWD "NBDMAGIC"

static int nbd_sendinit (_F_TCPSERVER_V1 *c, CLIENTINFO *cl)
{
	c->send (INIT_PASSWD);
	static unsigned long long magic = htonll(0x00420281861253LL);
	c->send (&magic,sizeof(magic));
	unsigned long long size = htonll(cl->getsize());
	// fprintf (stderr,"send size %Lu\n",cl->getsize());
	c->send (&size,sizeof(size));
	char stuff[128];
	memset (stuff,0,sizeof(stuff));
	c->send (stuff,sizeof(stuff));
	return 0;
}

#if 0
static int nbd_seek (CLIENTINFO *c)
{
	int ret = 0;
	unsigned long long end = c->u.req.from + c->u.req.len;
	debug_printf (DRSNAP,"Trying to seek %Lu + %lu = %Lu > %Lu\n"
		,c->u.req.from,c->u.req.len,end,c->size);
	if (end > c->size){
		tlmp_error ("Trying to seek pass end of file: %Lu + %lu > %Lu"
			,c->u.req.from,c->u.req.len,c->size);
		ret = -1;
	}else if (lseek64 (c->fd,c->u.req.from,SEEK_SET)==-1){
		tlmp_error ("Can't seek to position %lu (%s)"
			,c->u.req.from,strerror(errno));
		ret = -1;
	}
	return ret;
}
#endif

static void nbd_ntoh (nbd_request &req)
{
	req.magic = ntohl(req.magic);
	req.type = ntohl(req.type);
	req.from = ntohll(req.from);
	req.len = ntohl(req.len);
}

#if 0
static void nbd_sendok (_F_TCPSERVER_V1 *c, CLIENTINFO *cl)
{
	nbd_reply rep;
	rep.magic = ntohl(NBD_REPLY_MAGIC);
	rep.error = 0;
	memcpy (rep.handle,cl->u.req.handle,sizeof(rep.handle));
	c->send (&rep,sizeof(rep));
}


static int nbd_doread (_F_TCPSERVER_V1 *c, CLIENTINFO *cl)
{
	int ret = -1;
	if (nbd_seek(cl)!= -1){
		nbd_sendok (c,cl);
		int len = cl->u.req.len;
		char buf[len];
		if (read(cl->fd,buf,len)!=-1){
			c->send (buf,len);
			ret = 0;
		}else{
			tlmp_error("Can't read %s (%s)",cl->fname.c_str()
				,strerror(errno));
		}
	}
	return ret;
}
#endif

<mod>
int main (int argc, char *argv[])
{
	glocal int ret;
	glocal bool daemon = false;
	glocal const char *port = NULL;
	glocal const char *ctrlport = NULL;
	glocal const char *file_prefix = NULL;
	glocal const char *log = NULL;
	glocal const char *pidfile = NULL;
	//glocal const char *unixsock = NULL;
	glocal.ret = <call tlmpprogram>(argc,argv);
	<f init>
		setproginfo ("drsnapd",version
			,"DRBD userland server with snapshot\n");

		setarg ('c',"ctrlport","Control port/Administration"
			,glocal.ctrlport,true);
		setarg ('p',"port","Listen on TCP port",glocal.port,true);
		setarg ('D',"directory","Directory holding the volumes"
			,glocal.file_prefix,true);
		//setarg ('u',"unixctl","Control socket"
		//	,glocal.unixsock,false);

		setgrouparg ("Misc. options");
		setarg ('d',"daemon","Detach and run in background",glocal.daemon,false);
		setarg ('l',"log","Collect statistics in a file",glocal.log,false);
		setarg ('P',"pidfile","Pid file for daemon mode",glocal.pidfile,false);
	</f>
	<f showerror>
		if (glocal.daemon){
			syslog (LOG_ERR,"%s",msg);
		}
		fprintf (stderr,"%s\n",msg);
	</f>
	<f main_noarg>
		int ret = -1;
		if (snapfile_load(glocal.file_prefix)!=-1){
			glocal FILE *flog = NULL;
			if (glocal.log != NULL){
				glocal.flog = fopen (glocal.log,"a");
				if (glocal.flog == NULL){
					tlmp_error ("Can't open log file %s (%s)",glocal.log
						,strerror(errno));
				}
			}
			<obj TCPSERVER o>(glocal.port,10);
			<f newclient>
				CLIENTINFO *c = new CLIENTINFO(glocal.file_prefix);
				if (!c->is_ok()){
					tlmp_error ("No volume found");
					endclient = true;
				}else{
					info.data = c;
					settcpnodelay(true);
					if (nbd_sendinit (this,c)==-1){
						endclient = true;
					}
				}
			</f>
			<f endclient>
				// CLIENTINFO *c = (CLIENTINFO*)info.data;
			</f>
			<f receive>
				CLIENTINFO *c = (CLIENTINFO*)info.data;
				// We may receive several request back to back
				unsigned linelen = info.linelen;
				while (linelen > 0){
					unsigned needed = sizeof(c->u.breq) - c->inheader;
					if (needed > 0){
						if (linelen < needed){
							// Header still incomplete
							memcpy (c->u.breq + c->inheader,line,linelen);
							c->inheader += linelen;
							break;
						}else{
							// We have a full header and maybe some data
							memcpy (c->u.breq + c->inheader,line,needed);
							nbd_ntoh (c->u.req);
							line += needed;
							linelen -= needed;
							c->inheader += needed;
							if (c->u.req.magic != NBD_REQUEST_MAGIC){
								tlmp_error ("Bad magic for request %08lx != %08lx"
									,c->u.req.magic,NBD_REQUEST_MAGIC);
								endclient = true;
								break;
							}
							if (c->u.req.type == NBD_CMD_READ){
								if (glocal.flog != NULL){
									fprintf (glocal.flog,"read %s %Lu %d\n"
										,c->getpath()
										,c->u.req.from,c->u.req.len);
								}
								c->stats.reads++;
								c->stats.bytes_read += c->u.req.len;
								if (c->dev_null){
									tlmp_error ("Reading from %s while in /dev/null mode"
										,c->getpath());
								}
								if (c->read(no)==-1){
									endclient = true;
									break;
								}
								c->inheader = 0;
							}else if (c->u.req.type == NBD_CMD_WRITE){
								if (glocal.flog != NULL){
									fprintf (glocal.flog,"write %s %Lu %d\n"
										,c->getpath()
										,c->u.req.from,c->u.req.len);
								}
								c->stats.writes++;
								c->stats.bytes_written += c->u.req.len;
								// While we write a chunk, we change
								// the from,len in the nbd_request.
							}else if (c->u.req.type == NBD_CMD_DISC){
								// We do nothing with this command
								c->sendok(no);
								endclient = true;
							}else{
								tlmp_error ("Unsupported command %d: from %Lu len %lu"
									,c->u.req.type,c->u.req.from,c->u.req.len);
								endclient = true;
								break;
							}
						}
					}else if (c->u.req.type == NBD_CMD_WRITE){
						int len = linelen;
						int remain = c->u.req.len;
						if (len > remain) len = remain;
						if (!c->dev_null && c->write (line,len)==-1){
							endclient = true;
							break;
						}else{
							line += len;
							c->u.req.from += len;
							c->u.req.len -= len;
							linelen -= len;
							if (c->u.req.len == 0){
								c->inheader = 0;
								c->sendok (no);
							}
						}
					}else{
						tlmp_error ("Receiving data, lost");
						endclient = true;
						break;
					}
				}
			</f>
			</obj>
			glocal TCPSERVER_V1 *pto = &o;
			<obj TCPSERVER ctrl>(glocal.ctrlport,10);
			<f newclient>
				sendf ("drsnap %s\n",version);
				send ("help: snap branch quit end\n");
			</f>
			<f endclient>
			</f>
			<f receive>
				SSTRING word;
				line = str_copyword(word,line);
				line = str_skip(line);
				if (word == "snap"){
					// Move the first (only) connection to a newly created snapshot
					void *data;
					if (glocal.pto->iter_init(data)!=-1){
						CLIENTINFO *c = (CLIENTINFO*)data;
						c->snap();
						sendf ("Ok: Snapshot done and active: %s\n",c->getpath());
					}else{
						send ("*** No connection\n");
					}
				}else if (word == "/dev/null"){
					// Change the operation mode of the first (only) connection
					// This is for testing purpose. All writes are lost
					// when in /dev/null mode.
					void *data;
					if (glocal.pto->iter_init(data)!=-1){
						CLIENTINFO *c = (CLIENTINFO*)data;
						SSTRING keyw;
						str_copyword (keyw,line);
						if (keyw == "on"){
							c->dev_null = true;
							sendf ("Ok: %s in /dev/null mode\n",c->getpath());
						}else if (keyw == "off"){
							c->dev_null = false;
							sendf ("Ok: %s in normal mode (non /dev/null)\n",c->getpath());
						}else{
						}
					}else{
						send ("*** No connection\n");
					}
				}else if (word == "split"){
					// Move the first (only) connection to a newly created snapshot
					// Create another snapshot so we can access the parent
					void *data;
					if (glocal.pto->iter_init(data)!=-1){
						CLIENTINFO *c = (CLIENTINFO*)data;
						c->split(line);
						sendf ("Ok: Snapshot done and active: %s\n",c->getpath());
					}else{
						send ("*** No connection\n");
					}
				}else if (word == "splitrev"){
					// Create a snapshot of a specific revision
					// The snapshot will be in a new branch.
					SSTRING rev,newbranch;
					line = str_copyword (rev,line);
					line = str_skip(line);
					line = str_copyword (newbranch,line);
					line = str_skip(line);
					if (rev.is_empty() || newbranch.is_empty()
						|| line[0] != '\0'){
						send ("*** Invalid syntax: splitrev revision newbranch\n");
					}else{
						SNAPFILE *sn = snapfile_find_from_rev(rev.getval());
						if (sn != NULL){
							SNAPFILE *newsn = sn->snap('B',newbranch.c_str());
							if (newsn != NULL){
								sendf ("Ok: Snapshot done: %s\n",newsn->getpath());
							}else{
								send ("*** Can't create snapshot\n");
							}
						}
					}
				}else if (word == "branch"){
					SNAPFILE *sn = snapfile_find (line);
					if (sn == NULL){
						sendf ("*** Unknown branch: %s\n",line);
					}else{
						branch = line;
						sendf ("Ok: Default branch is now %s\n",branch.c_str());
					}
				}else if (word == "diff"){
					// We want to extract the difference between two
					// snapshots. The source and destination are identified
					// by a utime and events pair (from raid1 superblock).
					// This is done this way because this is what the client
					// can easily see. To assemble his RAID1, the client
					// has access to a connected NBD device and a local
					// device. He can extract the superblock information of
					// both.
					SSTRING from_utime,from_events,to_utime,to_events;
					line = str_copyword(from_utime,line);
					line = str_copyword(from_events,line);
					line = str_copyword(to_utime,line);
					str_copyword(to_events,line);
					if (to_events.is_empty()){
						sendf ("diff from-utime from-events to-utime to-events\n");
					}else{
						unsigned long long from_utimev = atoll(from_utime.c_str());
						unsigned long long from_eventsv = atoll(from_events.c_str());
						unsigned long long to_utimev = atoll(to_utime.c_str());
						unsigned long long to_eventsv = atoll(to_events.c_str());
						if (snapfile_diff (from_utimev,from_eventsv
							,to_utimev,to_eventsv,no) == -1){
							send ("*** Can't find the snapshots\n");
						}
						send ("Done\n");
					}
				}else if (word == "stats"){
					void *data;
					if (glocal.pto->iter_init(data)!=-1){
						while (1){
							CLIENTINFO *c = (CLIENTINFO*)data;
							sendf ("# %s read %d/%Lu writes %d/%Lu /dev/null=%s\n"
								,c->getpath()
								,c->stats.reads,c->stats.bytes_read
								,c->stats.writes,c->stats.bytes_written
								,c->dev_null ? "on" : "off");
							if (glocal.pto->iter_next(data)==-1) break;
						}
					}
					send ("Ok: \n");
				}else if (word == "quit"){
					endclient = true;
				}else if (word == "end"){
					send ("Ending server\n");
					exit (0);
				}else{
					sendf ("???\n");
				}
			</f>
			</obj>
			if (o.is_ok() && ctrl.is_ok()){
				o.setrawmode (true);
				<obj NETEVENT_MANAGER ne>();
				<f idle>
					if (glocal.flog != NULL) fflush (glocal.flog);
				</f>
				</obj>
				ne.add (o);
				ne.add (ctrl);
				if (glocal.daemon){
					int fd = open ("/dev/null",O_RDONLY);
					if (fd == -1){
						tlmp_error ("Can't open /dev/null, can't detach");
					}else{
						dup2 (fd,0);
						dup2 (fd,1);
						dup2 (fd,2);
					}
					setsid();
					pid_t pid = fork();
					if (pid == 0){
						if (glocal.pidfile != NULL){
							<call savefile>(glocal.pidfile,false);
							<f dowrite>
								fprintf (fout,"%d\n",getpid());
								return 0;
							</f>
							</call>
						}
					}else if (pid == (pid_t)-1){
						tlmp_error ("Can't fork (%s)",strerror(errno));
					}else{
						exit (0);
					}
				}
				ne.loop(3);
				ret = 0;
				if (glocal.flog != NULL) fclose (glocal.flog);
			}
		}
		return ret;
	</f>
	</call>
	return glocal.ret;
}
</mod>