#include <stdio.h>
#include <stdlib.h>
#include <stdarg.h>
#include <string.h>
#include <unistd.h>
#include <signal.h>
#include <time.h>
#include <errno.h>
#include <sys/utsname.h>

#ifdef HAVE_LIBVIRT
# include <libvirt/libvirt.h>
# include <libxml/parser.h>
# include <libxml/tree.h>
# include <libxml/xpath.h>
#endif
#ifdef HAVE_XENSTORE
# include <xs.h>
#endif

#include "list.h"
#include "xs_tools.h"
#include "apps.h"

/* ------------------------------------------------------------- */

#define BUFSIZE     64
#define PRE_STATE   "%16s: "

struct dom {
    int                domid;
    char               name[BUFSIZE];
    char               tty[BUFSIZE];
    int                connected;
    int                destroyed;
    int                watched;
    struct list_head   next;
};
static LIST_HEAD(doms);
static int domcnt;

static char *screen_rc      = "/etc/xen/xenscreenrc";
static char *screen_unlink  = NULL;
static char *screen_session = "xencon";
static char *screen_title   = "mon";
static int  screen_detached;
static int  screen_logging;
static int  screen_shell;

static char builtin_screen_rc[] =
"multiuser on\n"
"\n"
"# status line\n"
"sorendition =s wb\n"
"hardstatus lastline \"%{=b bw} xen |%{-} %-w%{= yb} %50>%n* %t %{-}%+w%<\"\n"
"\n"
"# logging\n"
"logfile /var/log/xen/console.%t\n"
"logfile flush 1\n"
"logtstamp on\n"
"\n"
"# disable killing windows\n"
"bind k\n"
"bind K\n"
"bind ^K\n"
"\n"
"# misc other useful settings\n"
"shell /bin/bash\n"
"defscrollback 5000\n"
"compacthist on\n"
"defutf8 on\n"
"termcapinfo xterm hs@\n"
"\n"
;

/* ------------------------------------------------------------- */

static struct dom *find_dom(int domid)
{
    struct dom *dom;
    struct list_head *item;

    list_for_each(item, &doms) {
	dom = list_entry(item, struct dom, next);
	if (dom->domid == domid)
	    return dom;
    }
    return NULL;
}

static struct dom *get_dom(int domid)
{
    struct dom *dom;

    dom = find_dom(domid);
    if (!dom) {
	dom = malloc(sizeof(*dom));
	memset(dom,0,sizeof(*dom));
	dom->domid = domid;
	list_add_tail(&dom->next, &doms);
    }
    return dom;
}

/* ------------------------------------------------------------- */

static int termsig;

static void catchsig(int sig)
{
    termsig = sig;
}

/* ------------------------------------------------------------- */

static int screen_command(char *arg0, ...)
{
    va_list args;
    char *argv[64];
    int i = 0;

    argv[i++] = "screen";
    argv[i++] = "-X";
    argv[i++] = "-S";
    argv[i++] = screen_session;

    argv[i++] = arg0;
    va_start(args, arg0);
    while (i < array_size(argv)) {
	argv[i] = va_arg(args, char*);
	if (NULL == argv[i])
	    break;
	i++;
    }
    va_end(args);

    return run_application_va(1, "screen", argv);
}

static int screen_attach(char *window)
{
    return run_application(1, "screen", "screen",
			   "-S", screen_session, "-r", "-x",
			   "-p", window ? window : "=",
			   NULL);
}

static void try_attach_domain(struct dom *dom, int boot)
{
    int rc;
    
    if (dom->connected)
	return;
    if (!strlen(dom->name))
	return;
    if (!strlen(dom->tty))
	return;

    fprintf(stderr, PRE_STATE "%s (%d) @ %s\n",
	    "connecting", dom->name, dom->domid, dom->tty);

    if (0 != access(dom->tty, R_OK)) {
	fprintf(stderr, "        error: no access to tty %s\n", dom->tty);
	return;
    }

#if 0
    /* known-racy, but better than nothing ... */
    rc = run_application(1, "fuser", "fuser", "-s", dom->tty, NULL);
    if (0 == rc) {
	fprintf(stderr,"        error: tty %s already in use\n", dom->tty);
	return;
    }
#endif

    if (screen_logging)
	rc = screen_command("screen", "-L", "-t", dom->name, dom->tty,
			    NULL);
    else
	rc = screen_command("screen", "-t", dom->name, dom->tty,
			    NULL);

#if 0
    /*
     * Hmm, not exactly the most elegant way to do this, has
     * some ugly glitches too.
     *
     * Switches back from the new window to the previous one.
     * Better would be to not switch in the first place,
     * seems screen can't do that though :-(
     */
    if (!boot)
	rc = screen_command("other", NULL);
#endif

    dom->connected = 1;
    domcnt++;
}

static void try_release_domain(struct dom *dom)
{
    if (!dom->destroyed)
	return;
    fprintf(stderr, PRE_STATE "%s (%d)%s\n", "disappeared",
	    dom->name, dom->domid, dom->connected ? " [conn]" : "");
    if (dom->connected)
	domcnt--;
    list_del(&dom->next);
    free(dom);
}

static void builtin_screen_setup(void)
{
    char *rc = strdup("/tmp/xenscreen-XXXXXX");
    int fd;

    fd = mkstemp(rc);
    if (-1 == fd) {
	fprintf(stderr,"mkstmp(%s): %s\n", rc, strerror(errno));
	return;
    }
    write(fd, builtin_screen_rc, sizeof(builtin_screen_rc));
    close(fd);

    fprintf(stderr, "Config file \"%s\" doesn't exist, using builtin config (%s).\n",
	    screen_rc, rc);
    screen_rc     = rc;
    screen_unlink = rc;
    return;
}

/* ------------------------------------------------------------- */
/* libvirt bits                                                  */

#ifdef HAVE_LIBVIRT

static int
libvirt_xml_xpath_str(const char *doc, char *xpath, char *dest, int len)
{
    xmlDocPtr xml = NULL;
    xmlXPathObjectPtr obj = NULL;
    xmlXPathContextPtr ctxt = NULL;
    int ret = -1;

    xml = xmlReadDoc((const xmlChar *) doc, "domain.xml", NULL,
                     XML_PARSE_NOENT | XML_PARSE_NONET |
                     XML_PARSE_NOWARNING);
    if (!xml)
        goto cleanup;
    ctxt = xmlXPathNewContext(xml);
    if (!ctxt)
        goto cleanup;

    obj = xmlXPathEval(BAD_CAST xpath, ctxt);
    if ((obj == NULL) || (obj->type != XPATH_STRING) ||
        (obj->stringval == NULL) || (obj->stringval[0] == 0)) {
        goto cleanup;
    }
    snprintf(dest, len, "%s", obj->stringval);
    ret = 0;

 cleanup:
    if (obj)
        xmlXPathFreeObject(obj);
    if (ctxt)
        xmlXPathFreeContext(ctxt);
    if (xml)
        xmlFreeDoc(xml);
    return ret;
}

static void libvirt_scan(virConnectPtr conn, int boot)
{
    int i, count, *ids;
    virDomainPtr vdom;
    struct dom *dom;
    const char *name, *xml;
    struct list_head *item;

    list_for_each(item, &doms) {
	dom = list_entry(item, struct dom, next);
	dom->destroyed = 1;
    }

    count = virConnectNumOfDomains(conn) + 4;
    ids = malloc(count * sizeof(int));
    count = virConnectListDomains(conn, ids, count);
    for (i = 0; i < count; i++) {
	dom = find_dom(ids[i]);
	if (dom) {
	    /* have it */
	    dom->destroyed = 0;
	    if (dom->connected)
		continue;
	    /* try again a few times in case we have no tty,
	     * it may show up a little later ... */
	    if (dom->watched > 8)
		continue;
	    dom->watched++;
	}
	/* new one */
	vdom = virDomainLookupByID(conn, ids[i]);
	name = virDomainGetName(vdom);
	xml  = virDomainGetXMLDesc(vdom, 0);
	if (!name || !xml)
	    continue;
	// fprintf(stderr, "\n-- xmldesc --\n%s\n--\n", xml);
	dom = get_dom(ids[i]);
	snprintf(dom->name, sizeof(dom->name), "%s", name);
	libvirt_xml_xpath_str(xml, "string(/domain/devices/console/@tty)",
			      dom->tty, sizeof(dom->tty));
	fprintf(stderr, "[ libvirt poll debug: %s (%d): tty=\"%s\" ]\n",
		dom->name, dom->watched, dom->tty);
	try_attach_domain(dom, boot);
    }
    free(ids);

    list_for_each(item, &doms) {
	dom = list_entry(item, struct dom, next);
	if (dom->destroyed)
	    try_release_domain(dom);
    }
}
#endif

/* ------------------------------------------------------------- */
/* xenstore bits                                                 */

#ifdef HAVE_XENSTORE
static void xenstore_scan(struct xs_handle *xenstore)
{
    xs_transaction_t xst;
    char **vec = NULL;
    int domid;
    unsigned int count, i;
    char path[BUFSIZE];
    struct dom *dom;

    /* look for running domains */
    if (!(xst = xs_transaction_start(xenstore))) {
	fprintf(stderr,"Oops, can't start xenstore transaction\n");
	exit(1);
    }
    vec = xs_directory(xenstore, xst, "/local/domain", &count);
    xs_transaction_end(xenstore, xst, 0);
    
    for (i = 0; i < count; i++) {
	domid = atoi(vec[i]);
	dom = get_dom(domid);
	snprintf(path, sizeof(path), "/local/domain/%d/name", domid);
	xenstore_read(xenstore, path, dom->name, sizeof(dom->name));
	snprintf(path, sizeof(path), "/local/domain/%d/console/tty", domid);
	xenstore_read(xenstore, path, dom->tty, sizeof(dom->tty));
	try_attach_domain(dom, 1);
    }
    if (vec)
	free(vec);
}

static void xenstore_update(struct xs_handle *xenstore)
{
    char **vec = NULL;
    int domid;
    unsigned int count, rc;
    char path[BUFSIZE], value[BUFSIZE];
    struct dom *dom;

    vec = xs_read_watch(xenstore, &count);
    if (NULL == vec) {
	fprintf(stderr,"xs_read_watch() failed\n");
	exit(1);
    }
    if (2 != sscanf(vec[XS_WATCH_PATH], "/local/domain/%d/%64s", &domid, path)) {
	if (1 != sscanf(vec[XS_WATCH_PATH], "/local/domain/%d", &domid))
	    goto cleanup;
	strcpy(path, "");
    }
    dom = get_dom(domid);
    
    if (0 == strcmp(path,"")) {
	rc = xenstore_read(xenstore, vec[XS_WATCH_PATH], value, sizeof(value));
	if (0 != rc)
	    dom->destroyed = 1;
	
    } else if (0 == strcmp(path, "console/tty")) {
	rc = xenstore_read(xenstore, vec[XS_WATCH_PATH], value, sizeof(value));
	if (0 != rc)
	    goto cleanup;
	strcpy(dom->tty, value);
	
    } else if (0 == strcmp(path, "name")) {
	rc = xenstore_read(xenstore, vec[XS_WATCH_PATH], value, sizeof(value));
	if (0 != rc)
	    goto cleanup;
	strcpy(dom->name, value);
	fprintf(stderr, PRE_STATE "%s (%d)\n", "new domain", dom->name, dom->domid);
	
    } else {
	goto cleanup;

    }

    try_attach_domain(dom, 0);
    try_release_domain(dom);

cleanup:
    if (vec)
	free(vec);
}
#endif

/* ------------------------------------------------------------- */

static void usage(FILE *fp)
{
    fprintf(fp,
	    "I'm managing xen consoles using screen.\n"
	    "\n"
	    "usage: xenscreen [options]\n"
	    "options:\n"
	    "   -h            print this text\n"
	    "   -b            print default screen config file\n"
	    "   -z            open a screen window with a shell\n"
#ifdef HAVE_LIBVIRT
	    "   -v uri        libvirt connect uri\n"
#endif
	    "\n"
	    "   -L            enable console output logging\n"
	    "   -c screenrc   screen config file              [%s]\n"
	    "   -S session    screen session name             [%s]\n"
	    "   -p window     preselect screen window\n"
	    "\n"
	    "-- \n"
	    "(c) 2006,07 Gerd Hoffmann <kraxel@redhat.com>\n",
	    screen_rc, screen_session);
}

int main(int argc, char *argv[])
{
    struct sigaction act,old;
    struct utsname uts;
    int maxfd;
    fd_set set;
    char *window = NULL;
    int nac, c;
    unsigned int rc, i;
    time_t last_ctrl_c = 0;
    char **nav;
    struct timeval tv;

#ifdef HAVE_LIBVIRT
    char *vir_url = getenv("VIRSH_DEFAULT_CONNECT_URI");
    virConnectPtr vir_conn = NULL;
#else
    void *vir_conn = NULL;
#endif

#ifdef HAVE_XENSTORE
    struct xs_handle *xenstore = NULL;
#else
    void *xenstore = NULL;
#endif

    for (;;) {
        if (-1 == (c = getopt(argc, argv, "hdbLzc:S:u:p:v:")))
            break;
        switch (c) {

	/* screen-like behaviour */
        case 'c':
	    screen_rc = optarg;
            break;
        case 'S':
	    screen_session = optarg;
            break;
	case 'p':
	    window = optarg;
	    break;
	case 'L':
	    screen_logging = 1;
	    break;
	case 'd':
	    screen_detached = 1;
	    break;

	/* other options */
	case 'u':
	    screen_unlink = optarg;
	    break;
	case 'z':
	    screen_shell = 1;
	    break;
	case 'b':
	    printf("%s", builtin_screen_rc);
	    exit(0);

#ifdef HAVE_LIBVIRT
	case 'v':
	    vir_url = optarg;
	    break;
#endif

        case 'h':
            usage(stdout);
            exit(0);
        default:
            usage(stderr);
            exit(1);
        }
    }

    if (!have_application("screen")) {
	fprintf(stderr, "screen not found in $PATH (not installed?), exiting.\n");
	exit(1);
    }

    memset(&uts, 0, sizeof(uts));
    if (0 == uname(&uts)) {
	char *h;
	if (NULL != (h = strstr(uts.nodename, ".")))
	    *h = 0;
	screen_title = malloc(strlen(uts.nodename) +4);
	sprintf(screen_title, "[%s]", uts.nodename);
    }

    if (NULL == getenv("STY") || NULL == strstr(getenv("STY"),screen_session)) {
	/* not running inside screen */
	if (!screen_detached) {
	    /* try to attach */
	    rc = screen_attach(window);
	    if (0 == rc)
		exit(0);
	} else {
	    /* This is a nop: just check if screen is running */
	    rc = screen_command("select", ".", NULL);
	    if (0 == rc) {
		fprintf(stderr,"Screen session \"%s\" already active, exiting.\n",
			screen_session);
		exit(0);
	    }
	}

	/* failing that, start a new screen session ... */
	fprintf(stderr,"Starting new screen session \"%s\".\n", screen_session);
	if (0 != access(screen_rc, R_OK))
	    builtin_screen_setup();
	nav = malloc(sizeof(char*) * (argc + 16));
	nac = 0;
	nav[nac++] = "screen";
	nav[nac++] = "-d";
	nav[nac++] = "-m";
	nav[nac++] = "-S";
	nav[nac++] = screen_session;
	nav[nac++] = "-c";
	nav[nac++] = screen_rc;
	nav[nac++] = "-t";
	nav[nac++] = screen_title;
	for (i = 0; argv[i] != NULL;)
	    nav[nac++] = argv[i++];
	if (screen_unlink) {
	    nav[nac++] = "-u";
	    nav[nac++] = screen_unlink;
	}
	nav[nac++] = NULL;
	rc = run_application_va(1, "screen", nav);
	/* ... and attach if asked for */
	if (0 == rc && !screen_detached) {
	    sleep(1); /* quick & dirty race work around */
	    rc = screen_attach(window);
	}
	exit(rc);
    }

    /* setup signal handler */
    memset(&act,0,sizeof(act));
    sigemptyset(&act.sa_mask);
    act.sa_handler = catchsig;
    sigaction(SIGTERM,&act,&old);
    sigaction(SIGINT,&act,&old);

    fprintf(stderr,
	    "###\n"
	    "### Managing Xen consoles using screen.\n"
	    "### This is the monitor process, at %s.\n"
	    "###\n"
	    "\n",
	    uts.nodename);
    if (screen_shell)
	screen_command("screen", "-t", "[shell]", "/bin/bash", NULL);

#ifdef HAVE_LIBVIRT
    if (vir_url) {
	fprintf(stderr, "trying libvirt (%s) ...\n", vir_url);
	vir_conn = virConnectOpenReadOnly(vir_url);
	if (vir_conn)
	    fprintf(stderr, "using libvirt\n");
	else
	    fprintf(stderr, "libvirt: can't connect\n");
    }
#endif

#ifdef HAVE_XENSTORE
    if (!vir_conn) {
	/* connect to xenstore */
	fprintf(stderr, "trying xenstore ...\n");
	xenstore = xenstore_open(1,1,1,1);
	if (xenstore) {
	    fprintf(stderr, "using xenstore\n");
	    xs_watch(xenstore, "/local/domain", "token");
	} else
	    fprintf(stderr, "xenstore: can't connect\n");
    }
#endif

    if (!vir_conn && !xenstore) {
	fprintf(stderr, "Failed to establish VM management connection.\n");
	fprintf(stderr, "Exiting in 10 seconds ...\n");
	sleep(10); /* give the user the chance to see the error */
	exit(1);
    }

    fprintf(stderr,"looking for existing domains\n");
#ifdef HAVE_LIBVIRT
    if (vir_conn)
	libvirt_scan(vir_conn, 1);
#endif
#ifdef HAVE_XENSTORE
    if (xenstore)
	xenstore_scan(xenstore);
#endif

    /* main loop */
    fprintf(stderr,"ok, watching out for changes now\n");
    for (;;) {
	if (termsig) {
	    if (!domcnt)
		break;
	    if (time(NULL) - last_ctrl_c < 3)
		break;
	    fprintf(stderr,
		    "\n"
		    "Got ^C - still %d domain(s) active - not quitting.\n"
		    "\n"
		    "You should better use detach instead (^A d).\n"
		    "Or kill all windows (^A \\) if you don't want\n"
		    "keep screen hanging around.\n"
		    "\n"
		    "Hit ^C within 3 secs again to quit nevertheless.\n"
		    "\n",
		    domcnt);
	    last_ctrl_c = time(NULL);
	    termsig = 0;
	}

	FD_ZERO(&set);
	maxfd = 0;

#ifdef HAVE_XENSTORE
	if (xenstore) {
	    int fd = xs_fileno(xenstore);
	    FD_SET(fd, &set);
	    if (maxfd < fd)
		maxfd = fd;
	    tv.tv_sec = 0;
	    tv.tv_usec = 0;
	}
#endif
#ifdef HAVE_LIBVIRT
	if (vir_conn) {
	    /* FIXME: polling once per second */
	    tv.tv_sec = 1;
	    tv.tv_usec = 0;
	}
#endif

	switch (select(maxfd+1, &set, NULL, NULL, tv.tv_sec ? &tv : NULL)) {
	case -1:
	    if (EINTR == errno)
		continue; /* termsig check */
	    perror("select");
	    break;
	case 0:
#ifdef HAVE_LIBVIRT
	    if (vir_conn)
		libvirt_scan(vir_conn, 0);
#endif
	    break;
	default:
	    break;
	}

#ifdef HAVE_XENSTORE
	if (xenstore && FD_ISSET(xs_fileno(xenstore), &set))
	    xenstore_update(xenstore);
#endif
    }

    if (screen_unlink)
	unlink(screen_unlink);
    return 0;
}