#!/usr/bin/ruby
require 'yaml'
require 'pp'
require 'shellwords'
require 'getoptlong'
require 'ostruct'
require 'json'

MAPPINGS = [
	# Syntax:
	# [name_in_yml, docker_option_name, lambda_for_creating_yml_for_container=nil, options={}]
	# Options are:
	#   type: Sets type of this option
	#     switch: Only adds the docker_option_name to the command, ignoring the value.
	#     name: Value will be set to the name of the container (from yml filename or root element)
	#     hidden: Will not be added to the run command
	#   allow_empty: May be empty.
	#   escape: If set to false, will not be shellescaped

	['stdin_open', '--interactive', lambda{|c,i| c.Config.OpenStdin}, {:type=>:switch}],
	['tty', '--tty', lambda{|c,i| c.Config.Tty}, {:type=>:switch}],
	#['detach', '--detach', {:type=>:switch}],
	['remove', '--rm', lambda{|c,i| c.HostConfig.AutoRemove}, {:type=>:switch}],
	['container_name', '--name', nil, {:type=>:name}],
	['ports', '--publish', lambda{|c,i| c.HostConfig.PortBindings.to_h.collect{|k,v| "#{v.first.HostPort}:#{k.to_s.gsub("/tcp", "")}" }}],
	['restart', '--restart', lambda{|c,i| c.HostConfig.RestartPolicy.Name}],
	['environment', '--env', lambda{|c,i| e=c.Config.Env.delete_if{|v| i.Config.Env.include?(v)}; e}],
	['volumes', '--volume', lambda{|c,i| c.HostConfig.Binds}],
	['mem_limit', '--memory', lambda{|c,i| (mem=c.HostConfig.Memory)==0 ? nil : mem}],
	['links', '--link'],
	['stop_signal', '--stop-signal', lambda{|c,i| c.Config.StopSignal!="SIGTERM" ? c.Config.StopSignal : nil}],
	['stop_grace_period', '--stop-timeout', lambda{|c,i| c.Config.StopTimeout}],
	['devices', '--device'],
	['net', '--net'],
	['ipv6', '--ipv6'],
	['user', '--user'],
	['networks', '--network', lambda{|c,i| c.NetworkSettings.Networks.to_h.keys.map(&:to_s)}, {:type=>:hidden}],
	['entrypoint', '--entrypoint', lambda{|c,i| (ep=c.Config.Entrypoint) == i.Config.Entrypoint ? nil : ep}, {:allow_empty=>true}],
	['labels', '--label', lambda{|c,i| l=Hash[c.Config.Labels.to_h.delete_if{|k,v| i.Config.Labels[k]==v rescue false}.map{|k,v| [k.to_s, v]}] ; l}],
	['hostname', '--hostname'],
	['shm_size', '--shm-size'],
	['working_dir', '--workdir', lambda{|c,i| (wd=c.Config.WorkingDir) == i.Config.WorkingDir ? nil : wd}],
	['init', '--init', lambda{|c,i| c.Config.Init}, {:type=>:switch}],
	['image', nil, lambda{|c,i| c.Config.Image}],
	['command', nil, lambda{|c,i| (cmd=c.Config.Cmd.join(" ") rescue nil) == (i.Config.Cmd.join(" ") rescue nil)? nil : cmd}, {:escape=>false}],
	['test', nil, nil, {type: :hidden}],
	['auto_update', nil, nil, {type: :hidden}],
	['ofelia', nil, nil, {type: :hidden}]
]


def get_sample(name="container")
	return <<~HEREDOC
		image: "my/image:1.2.3"
		restart: always
		detach: true
		test: true # Having this set to true will prevent dup from creating a label de.fabianonline.dup
		auto_update: true # Enabling this adds a label to allow watchtower to automatically update this container
		  
		ports:
		  - "1234:1234"
		  - "8080"
		  - "1001:1001/udp"
		  
		environment:
		  VERSION: 1.7.0
		  TZ:
		  
		volumes:
		  - "/etc/localtime:/etc/localtime:ro"
		  
		labels:
		  nginx_virtual_host: container.home.schle.nz
		  nginx_port: 80
		  nginx_additional_ports: "443, 5050"
		  nginx_allow: fabian # "user" or "@group" or "user, user, @group" or "all"
		  nginx_no_robots: true
		  nginx_public_paths: "/public, /api"
		  nginx_type: # "http" (default), "https", "fastcgi", "skip" (doesn't create any entries)
		  nginx_client_max_body_size: "25M"
		  
		networks:
		  - "nginx"
		  - "mosquitto"
		  - "bridge"
		  
		user: "1000:1000"
		command: "/bin/bash"
		entrypoint: "/script.sh"
		build: "/data/dir"
		pull: false
		container_name: "Something"
		hostname: "foo"
		mem_limit: 125M
		stdin_open: false
		net: host
		tty: false
		remove: false
		stop_signal: SIGUSR1
		stop_grace_period: 5
		shm_size: 256M
		
		before_build:
		  - "echo 'Starting build'"
		after_build:
		  - "echo 'After build'"
		before_run:
		  - "echo 'Starting run'"
		after_run:
		  - "echo 'Run has finished.'"
		    - "docker restart container"
		  
		devices:
		  - "/dev/ttyUSB0:/dev/ttyUSB0"
		
		ofelia:
		  - name: update # No spaces!
		    schedule: "@every 4h" # go cron format! See https://pkg.go.dev/github.com/robfig/cron?utm_source=godoc
		    command: uname -a
		    user: foo
		    tty: false
		    no-overlap: true
		  - name: Test 2
		HEREDOC
end

def action_sample
	puts get_sample
	exit 1
end

def action_help
	puts "-h, --help              Show this help."
	puts "-n, --dry-run           Don't execute any commands."
	puts "-s, --sample            Outputs a sample yml file."
	puts "-c, --create            Create a new container yml file."
	puts "-p, --pull              (Try to) Pull the image(s) before starting the container."
	puts "-l, --list              Lists all available containers."
	puts "    --host              Starts a container with net: host - which implies not using ports, networks and/or links."
	puts "    --regenerate        Generates YAML content from a running container."
	puts "    --all               Use all containers defined in yml files."
	puts "-v, --verbose           More verbose output."
	exit 1
end

def run_cmd(cmd, ignore_returnvalue=false, catch_interrupt=false, ignore_dry_run: false)
	verbose "+ #{cmd}"
	returnvalue=false
	begin
		returnvalue = ($dry_run && !ignore_dry_run) ? true : system("bash -c #{cmd.shellescape}")
	rescue Interrupt
		raise if not catch_interrupt
	end
	raise "Command returned a non-zero exit value." if returnvalue!=true && !ignore_returnvalue
	return returnvalue
end

def collect_commands(&block)
	$commands = []
	def run_cmd(cmd, ignore_returnvalue=false, catch_interrupt=false, ignore_dry_run: false); $commands<<cmd; return 0; end

	yield

	return $commands
end

def esc(obj, escape=true)
	return obj unless escape
	return obj.to_s.shellescape
end

def action_create(container)
	file = container.file
	raise "File #{file} already exists" if File.exists?(file)
	
	File.open(file, "w") {|f| f.write(get_sample(container))}
	
	exec("/usr/bin/editor #{file.shellescape}")
	exit 1 # will never be reached because exec replaces this process. But just in case... ;-)
end

def action_run(container)
	# $net_host
	container.each{|c| c.run(container.count>1)}
end

def list(base_dir)
	result = {:containers=>{}, :errors=>[]}
	Dir[File.join(base_dir, "*.yml")].sort.each do |file|
		basename = File.basename(file)[0..-5]
		yaml = File.open(file, "r") {|f| YAML.load(f.read)}
		keys = yaml.keys
		
		if keys.include? basename
			additional = keys.reject{|k| k==basename}
			result[:containers][basename] = additional
		else
			result[:errors] << "#{file}:\n    Missing a key called #{basename}\n    Has keys:\n#{keys.map{|k| "        #{k}"}.join("\n")}"
		end
	end
	return result
end

def action_list(base_dir)
	list = list(base_dir)
	list[:containers].each do |container, additional|
		puts container
		additional.each{|c| puts " + #{c}"}
	end
	
	list[:errors].each do |e|
		puts "ERROR: #{e}"
		puts
	end
	
	return list[:errors].empty? ? 0 : 1
end

def action_completion(base_dir, complete_str)
	#puts "in completion"
	return unless base_dir
	#puts "str is #{complete_str}"
	containers = list(base_dir)[:containers].keys.select{|c| c.start_with?(complete_str)}
	puts containers.join(" ")
end	

def action_regenerate(container)
	container.each do |c|
		data = c.regenerate
		p collect_commands{ c.run }
	end
end

class Container
	def initialize(name)
		@name = name
	end
	
	def exists?
		run_cmd("docker inspect --format='1' #{@name.shellescape} >/dev/null 2>&1", true, ignore_dry_run: true)
	end
	
	def load
		raise "File #{self.filename} not found." unless File.exists?(self.filename)

		@data = File.open(self.filename, "r") {|f| YAML.load(f.read)}

		if @data.has_key?(@name)
			if @data.keys.count > 1
				raise "Expected #{self.filename} to contain all data on top level or underneath a single root element named #{@name}."
			end
			@data = @data[@name]
		end
	end
	
	def filename; "%s/%s.yml" % [$base_dir, @name]; end

	def add_label(key, value)
		if @data["labels"].is_a? Array
			@data["labels"] << "#{key}=#{value}"
		elsif @data["labels"].is_a? Hash
			@data["labels"][key] = value
		else
			raise "data['labels'] is of an unexpected type: #{@data["labels"].class}"
		end
	end
	
	def build_run_command
		$net_host = true if @data["net"]=="host"
		cmd = ["docker", "create"]
		cmd << "--net" << @data["networks"][0] if @data["networks"] && @data["networks"].count>0 && !$net_host
		cmd << "--net" << "host" if $net_host

		@data["labels"] ||= []
		
		if !@data["test"]
			add_label("de.fabianonline.dup", "true")
		end

		if @data["auto_update"]
			add_label("com.centurylinklabs.watchtower.enable", "true")
		end

		if @data["ofelia"] && @data["ofelia"].count > 0
			add_label("ofelia.enabled", "true")

			@data["ofelia"].each_with_index do |job, index|
				name = job["name"] || "job-#{index}"
				name = name.downcase.gsub(/[^a-z0-9\-]/, "-").gsub(/-+/, "-")
				schedule = job["schedule"]
				next unless schedule
				command = job["command"]
				next unless command
				add_label("ofelia.job-exec.#{name}.schedule", schedule)
				add_label("ofelia.job-exec.#{name}.command", command)
				add_label("ofelia.job-exec.#{name}.user", job["user"]) if job["user"]
				add_label("ofelia.job-exec.#{name}.tty", job["tty"].to_s) if job["tty"]!=nil
				ol = job["no-overlap"]
				add_label("ofelia.job-exec.#{name}.no-overlap", ol!=nil ? ol.to_s : "true")
			end
		end

		MAPPINGS.each do |mapping|
			yml_name, cmd_name, reverse_lambda, options = *mapping
			next if $net_host && %w(links net ports).include?(yml_name)
			next if options && options[:type]==:hidden
			options ||= {}
			if !@data[yml_name]
				if options[:type]==:name
					cmd << cmd_name << esc(@name)
				end
				next
			end
			
			if options[:allow_empty] && @data[yml_name]==""
				cmd << cmd_name << '""'
				next
			end
			 
			
			if options[:type]==:switch
				cmd << cmd_name
				next
			end
			
			if @data[yml_name].is_a?(Array)
				@data[yml_name].each {|val| cmd << cmd_name << esc(val, options[:escape])}
			elsif @data[yml_name].is_a?(Hash)
				@data[yml_name].each {|key, val| cmd << cmd_name << "#{esc(key, options[:escape])}=\"#{esc(val, options[:escape])}\""}
			else
				cmd << cmd_name << esc(@data[yml_name], options[:escape])
			end
		end
		
		return cmd.compact.join(" ")
	end
	
	def build
		(@data["before_build"] || []).each{|cmd| run_cmd(cmd)}
		cmd = ["docker", "build", "-t", @data["image"].shellescape, @data["build"].shellescape]
		run_cmd(cmd.join(" "))
		(@data["after_build"] || []).each{|cmd| run_cmd(cmd)}
	end
	
	def pull
		run_cmd("docker pull #{@data["image"].shellescape}", true)
	end
	
	def stop_and_remove
		verbose "Stopping and removing old container..."
		run_cmd("docker stop #{@name.shellescape} >/dev/null 2>&1", true)
		run_cmd("docker rm -f #{@name.shellescape} >/dev/null 2>&1", true)
	end
	
	def create_networks
		if !$net_host && (@data["networks"]||=[]).count>0
			networks = `docker network ls --format '{{.Name}}'`.split
			@data["networks"].each do |network|
				if networks.include?(network)
					verbose "Network #{network} exists."
				else
					verbose "Creating network #{network}..."
					run_cmd("docker network create #{network.shellescape}", true)
				end
			end
		end
	end
	
	def connect_to_networks
		if !$net_host && @data["networks"].count>0
			@data["networks"].each do |network|
				verbose "Connecting container to network #{network}..."
				run_cmd("docker network connect #{network.shellescape} #{@name.shellescape} >/dev/null", true)
			end
		end
	end
	
	def attach
		verbose "Attaching container..."
		run_cmd("docker attach #{@name.shellescape}")
	end
	
	def log
		puts "Opening log... (press Ctrl-C to exit)"
		puts
		run_cmd("docker logs -f #{@name.shellescape}", true, true)
	end
	
	def run(force_background = false)
		raise "No data loaded. Maybe call .load first?" if @data.nil?
		
		if @data["build"]
			self.build
		end
		
		(@data["before_run"] || []).each{|cmd| run_cmd(cmd)}

		if @data["pull"] || $pull
			self.pull
		end
		
		if self.exists?
			puts "Recreating #{@name}..."
		else
			puts "Creating   #{@name}..."
		end
		
		self.stop_and_remove
		
		self.create_networks
		
		verbose "Creating container..."
		cmd = build_run_command
		run_cmd(cmd + " >/dev/null")
		
		self.connect_to_networks
		
		verbose "Starting container..."
		run_cmd("docker start #{@name.shellescape} >/dev/null")
		
		(@data["after_run"] || []).each{|cmd| run_cmd(cmd)}
		
		if !force_background
			if ! @data["detach"]
				self.attach
			else
				self.log
			end
		end
	end
	
	def regenerate
		c_data = `docker inspect #{@name}`
		c = JSON.parse(c_data, object_class: OpenStruct).first
		i_data = `docker inspect #{c.Image}`
		i = JSON.parse(i_data, object_class: OpenStruct).first
		
		data = {}
		
		MAPPINGS.each do |m|
			name, cmd, lmbd, opts = *m
			
			next unless lmbd
			result = lmbd.call(c, i)
			next if opts && opts[:type]==:switch && result==false
			next if result==nil || result==[]
			
			data[name] = result
		end
		
		@data = data
		return data
	end
end

def verbose(str); puts(str) if $verbose; end

action = :run
container = nil
$dry_run = false
$verbose = false
$pull = false
needs_container = true
completion_str = ""
silent_basedir = false
$net_host = false

if ENV['DUP_DIR']
	$base_dir = ENV['DUP_DIR']
else
	$base_dir = File.join(Dir.home, ".dup")
	puts "Environment variable DUP_DIR is not set. Looking for .yml files in #{$base_dir}" unless silent_basedir
end

opts = GetoptLong.new(
	[ '--sample', '-s', GetoptLong::NO_ARGUMENT ],
	[ '--create', '-c', GetoptLong::NO_ARGUMENT ],
	[ '--help', '-h', GetoptLong::NO_ARGUMENT ],
	[ '--dry-run', '-n', GetoptLong::NO_ARGUMENT ],
	[ '--list', '-l', GetoptLong::NO_ARGUMENT ],
	[ '--pull', '-p', GetoptLong::NO_ARGUMENT ],
	[ '--update', '-u', GetoptLong::NO_ARGUMENT ],
	[ '--_completion', GetoptLong::OPTIONAL_ARGUMENT ],
	[ '--host', GetoptLong::NO_ARGUMENT ],
	[ '--regenerate', GetoptLong::NO_ARGUMENT ],
	[ '--verbose', '-v', GetoptLong::NO_ARGUMENT ],
	[ '--all', GetoptLong::NO_ARGUMENT ],
)

opts.each do |opt, arg|
	case opt
		when '--sample'
			action_sample
		when '--create'
			action = :create
		when '--help'
			action_help
		when '--dry-run'
			puts "Dry-run. Not going to execute any command."
			$dry_run = true
		when '--list'
			action = :list
			needs_container = false
		when '--pull'
			$pull = true
		when '--update'
			action = :update
			needs_container = false
		when '--_completion'
			action = :_completion
			completion_str = arg
			needs_container = false
			silent_basedir = true
		when '--host'
			$net_host = true
		when '--regenerate'
			action = :regenerate
			needs_container = false
		when '--verbose'
			$verbose = true
		when '--all'
			container = Dir[$base_dir + "/*.yml"].sort.map{|f| Container.new(File.basename(f, ".yml"))}
	end
end

container ||= ARGV.map{|c| Container.new(c)}

if needs_container
	if container.empty?
		raise "No container given."
	end
	
	container.each(&:load)
end

if action == :create
	action_create(container)
elsif action == :run
	action_run(container)
elsif action == :list
	action_list(base_dir)
elsif action == :update
	action_update(base_dir)
elsif action == :regenerate
	action_regenerate(container)
	container.each {|c| puts c.build_run_command}
elsif action == :_completion
	action_completion($base_dir, completion_str)
end