#!/usr/bin/env ucode
'use strict';
import { basename, readlink, open, stdout } from "fs";
let udebug = require("udebug");
let uloop = require("uloop");
let libubus = require("ubus");
uloop.init();
let ubus = libubus.connect();

let opts = {
	select: []
};
let service_config;
let trace_config;

const usage_message = `
Usage: ${basename(sourcepath())} [<options>] <command> [<args>]

  Options:
    -f				Ignore errors on opening rings
    -e <service list>		List of services to enable before running the command
    -d <duration>:		Only fetch data up to <duration> seconds old
    -o <file>|-			Set output file for snapshot/stream (or '-' for stdout)
    -i <process>[:<name>]	Select debug buffer for snapshot/stream
    -s <path>			Use udebug socket <path>
    -q				Suppress warnings/error messages
    -l				Select only debug buffers containing log messages
    -t				Include timestamps in logstream output

  Commands:
    reset:			Reset enabled debug buffers to configured settings
    set	<service list>:		Set enabled debug buffers
    list:			List available debug buffers
    snapshot:			Create a pcapng snapshot of debug buffers
    set_flag [<name>=0|1 ...]	Set ring buffer flags
    get_flags			Get ring buffer flags
    stream:			Stream packet data as pcap
    logstream:			Stream syslog data as text
    logdump:			Dump syslog data snapshot as text

  Service list: space separated list of services matching the config
    - <name>			Enable service <name>
    - <name>.<opt>=<val>	Set service <name> option <opt> to <val>

`;

function _warn(str) {
	if (opts.quiet)
		return;

	warn(str);
}

function usage() {
	warn(usage_message);
	exit(1);
}

function parse_service_value(svc, option, val)
{
	if (option) {
		svc[option] = val;
		val = '1';
	}

	svc.enabled = val;
}

function parse_service_entry(name, option, val)
{
	if (wildcard(name, "kernel:*")) {
		name = substr(name, 7);
		trace_config ??= [];
		push(trace_config, name);
		return;
	}

	if (index(name, "*") >= 0) {
		for (let svcname, svc in service_config) {
			if (!wildcard(svcname, name))
				continue;

			parse_service_value(svc, option, val);
		}
		return;
	}

	let svc = service_config[name];
	if (svc) {
		parse_service_value(svc, option, val);
		return;
	}

	if (!option && !+val)
		return;

	svc = service_config[name] = {
		enabled: "1"
	};

	if (option)
		parse_service_value(svc, option, val);
}

function parse_service_list(val) {
	if (!service_config) {
		service_config = ubus.call('udebug', 'get_config', { override: false })?.service;
		if (!service_config) {
			_warn('Failed to get current service config');
			exit(1);
		}
		for (let name, svc in service_config)
			svc.enabled = "0";
	}

	let val_list = split(val, /[ \t]+/);
	for (let cur in val_list) {
		cur = split(cur, "=");
		let val = cur[1] ?? "1";

		cur = split(cur[0], ".", 2);
		parse_service_entry(cur[0], cur[1], val);
	}
}

while (substr(ARGV[0], 0, 1) == "-") {
	let opt = substr(shift(ARGV), 1);
	switch(opt) {
	case 'd':
		opts.duration = +shift(ARGV);
		break;
	case 's':
		opts.socket = shift(ARGV);
		break;
	case 'i':
		push(opts.select, shift(ARGV));
		break;
	case 'o':
		opts.output_file = shift(ARGV);
		break;
	case 'e':
		parse_service_list(shift(ARGV));
		break;
	case 'q':
		opts.quiet = true;
		break;
	case 'f':
		opts.force = true;
		break;
	case 'l':
		opts.log_only = true;
		break;
	case 't':
		opts.timestamp = true;
		break;
	default:
		usage();
	}
}

let procs = {};
let selected = [];
let rings = {};
let subscriber;
let pcap, log_out;

function ring_selected(ring) {
	if (!length(opts.select))
		return true;

	for (let sel in opts.select) {
		let match = split(sel, ":", 2);
		if (wildcard(ring.proc_name, match[0]) &&
		    (!match[1] || wildcard(ring.ring_name, match[1])))
			return true;
	}

	return false;
}

function poll_data() {
	let data = [];
	for (let ring_id in rings) {
		let ring = rings[ring_id];
		let s = ring[1].fetch();
		if (s)
			push(data, s);
	}
	if (length(data) > 0) {
		if (log_out) {
			udebug.foreach_packet(data, (entry, data, timestamp) => {
				if (!length(data))
					return;
				if ((opts.timestamp && !log_out.write(sprintf("[%.6f] ", timestamp / 1000000.0))) ||
					!log_out.write(data) || !log_out.write("\n"))
					uloop.end();
				log_out.flush();
			});
		}
		if (pcap && pcap.write(data) == null)
			uloop.end();
	}
}

function open_ring(ring, poll) {
	let ring_name =` ${ring.proc_name}:${ring.ring_name}`;
	let ref = udebug.get_ring(ring);

	if (!ref)
		return null;
	if (opts.log_only && ref.get_info().format != udebug.FORMAT_STRING)
		return false;
	if (opts.duration)
		ref.set_fetch_duration(opts.duration);
	if (poll)
		ref.set_poll_cb(() => { poll_data() });

	let ring_id = ring.id + "";
	ring = [ ring_name, ref ];
	rings[ring_id] = ring;

	return ring;
}

function open_trace() {
	let ring = udebug.trace_ring("udebug-" + readlink("/proc/self"));

	ring.set_file("trace_clock", "tai");
	ring.set_file("set_event", join("\n", trace_config));
	ring.set_poll_cb(poll_data);
	rings.kernel = [ "kernel", ring ];
}

function open_log_out() {
	let out = opts.output_file;
	if (!opts.output_file || out == "-")
		log_out = stdout;
	else
		log_out = open(opts.output_file, "w");

	if (!log_out) {
		_warn(`Could not open output file\n`);
		exit(1);
	}
}

function open_pcap_out() {
	if (!opts.output_file) {
		_warn(`No output file\n`);
		exit(1);
	}
	let out = opts.output_file;
	if (out == "-")
		out = null;

	pcap = udebug.pcap_file(out);
	if (!pcap) {
		_warn(`Failed to open output\n`);
		exit(1);
	}
}

function stream_data(log) {
	if (log)
		open_log_out();
	else
		open_pcap_out();

	subscriber = ubus.subscriber((req) => {
		let type = req.type;
		let ring = req.data;
		let ring_id = ring.id + "";
		if (type == "remove") {
			ring = rings[ring_id];
			if (!ring)
				return;

			ring[1].close();
			delete rings[ring_id];
		} else if (type == "add") {
			open_ring(ring, true);
			poll_data();
		}
	}, null, [ "udebug" ]);
	for (let ring in selected) {
		if (open_ring(ring, true) == null) {
			_warn(`Failed to open ring ${ring.ring_name}\n`);
			if (opts.force)
				continue;

			exit(1);
		}
	}

	if (trace_config)
		open_trace();

	let done = () => { uloop.end(); };
	signal('SIGINT', done);
	signal('SIGTERM', done);

	poll_data();
	delete opts.duration;
	uloop.run();
}

let cmds = {
	list: function() {
		for (let proc in procs) {
		    print(`Process ${proc}:\n`);
		    for (let ring in procs[proc])
			print(`  - ${ring.ring_name}\n`);
		}
	},
	snapshot: function() {
		open_pcap_out();

		if (!length(selected)) {
			_warn(`No available debug buffers\n`);
			exit(1);
		}

		for (let ring in selected) {
			if (open_ring(ring) == null) {
				_warn(`Failed to open ring ${ring.proc_name}:${ring.ring_name}\n`);
				if (opts.force)
					continue;

				exit(1);
			}
		}

		poll_data();
		pcap.close();
	},
	set_flag: function() {
		for (let ring in selected) {
			if (!length(ring.flags))
				continue;

			let mask = 0, set = 0;
			for (let flag in ring.flags) {
				for (let change in ARGV) {
					change = split(change, "=", 2);
					let name = change[0];
					let val = !!int(change[1]);
					if (flag[0] == name)
						if (val)
							set |= flag[1];
						else
							mask |= flag[1];
				}
			}

			if (!(mask | set))
				continue;

			let r = open_ring(ring);
			if (!r)
				continue;

			r[1].change_flags(mask, set);
		}
	},
	get_flags: function() {
		for (let ring in selected) {
			if (!length(ring.flags))
				continue;

			let r = open_ring(ring);
			if (!r)
				continue;

			print(`${r[0]}\n`);
			let flags = r[1].get_flags();
			for (let flag in ring.flags)
				print(`\t${flag[0]}=${((flags & flag[1]) == flag[1]) ? 1 : 0 }\n`);
		}
	},
	stream: function() {
		stream_data(false);
	},
	logstream: function() {
		stream_data(true);
	},
	logdump: function() {
		open_log_out();

		if (!length(selected)) {
			_warn(`No available debug buffers\n`);
			exit(1);
		}

		for (let ring in selected) {
			if (open_ring(ring) == null) {
				_warn(`Failed to open ring ${ring.proc_name}:${ring.ring_name}\n`);
				if (opts.force)
					continue;

				exit(1);
			}
		}

		poll_data();
	},
	reset: function() {
		exit(0);
	},
	set: function() {
		exit(0);
	},
};

let cmd = shift(ARGV);
if (!cmds[cmd])
	usage();

if (cmd == "logstream" || cmd == "logdump")
	opts.log_only = true;

if (cmd == 'reset') {
	ubus.call('udebug', 'set_config', { override: true });
	exit(0);
}

if (cmd == 'set') {
	if (!length(ARGV))
		usage();
	for (let val in ARGV)
		parse_service_list(val);
}

if (service_config) {
	ubus.call('udebug', 'set_config', {
		override: true,
		service: service_config,
	});
}

if (service_config && cmd != 'set')
	sleep(1000);

let ring_list = ubus.call("udebug", "list");
if (!ring_list || !ring_list.results) {
	warn("Failed to get ring buffer list from udebugd\n");
	exit(1);
}

ring_list = ring_list.results;
for (let ring in ring_list) {
	if (!ring_selected(ring))
		continue;

	let proc = procs[ring.proc_name];
	if (!proc) {
	    proc = [];
	    procs[ring.proc_name] = proc;
	}
	push(proc, ring);
	push(selected, ring);
}

if (cmd != "list" && !udebug.init(opts.socket)) {
	_warn(`Failed to connect to udebug socket\n`);
	exit(1);
}

cmds[cmd]();