require 'optparse'
require 'highline'
require 'hirb'

module Vzlimit
  class CLI
    ACTIONS = [:limit, :list]
    VZ_CONF_DIR = '/etc/vz/conf'
    
    def initialize(stdout, arguments)
      @stdout = stdout
      @arguments = arguments
      @h = HighLine.new
      
      begin
        # Check action argument
				if !@arguments[0] || (!ACTIONS.include?(@arguments[0].to_sym) &&! @arguments[0] =~ /^-.*/)
          usage
        end
        
        # list
        if @arguments[0].to_sym == :list
          list
        end
        
        # limit
        if @arguments[0].to_sym == :set
          set
        end

				# help
				if @arguments[0] == '-h' || @arguments[0] == '--help'
					help
				end

				# version
        if @arguments[0] == '-v' || @arguments[0] == '--version'
          version
        end
      rescue
        @h.say $!
      end
    end
    
    def list
      # Scan directory /etc/vz/conf/
      conf_files = Dir.entries(VZ_CONF_DIR)
      
      vps_list = []
      
      conf_files.each do |file|
        if !(file =~ /^\d\d+\.conf$/)
          next;
        end
        
        vpsid = /^(\d\d+)/.match(file)[0]
        conf = File.open(VZ_CONF_DIR + '/' + file, "rb").read
        
        # Get memory limits
        null, memory_soft, memory_hard = /\nPRIVVMPAGES="(\d+):(\d+)"\n/.match(conf).to_a
        memory_soft = (memory_soft.to_f / 1024 * 4).to_i.to_s
        memory_hard = (memory_hard.to_f / 1024 * 4).to_i.to_s
        
        # Get diskspace limits
        null, diskspace_soft, diskspace_hard = /\nDISKSPACE="(\d+):(\d+)"\n/.match(conf).to_a
        diskspace_soft = (diskspace_soft.to_f / 1024).to_i.to_s
        diskspace_hard = (diskspace_hard.to_f / 1024).to_i.to_s
        
        # Get ip
        ip = /\nIP_ADDRESS="(\d+\.\d+\.\d+\.\d+)"\n/.match(conf)[1]
        
        vps_list << {
          :vpsid            => vpsid,
          :ip               => ip,
          :memory_soft      => @h.color(memory_soft + ' MB', :yellow),
          :memory_hard      => @h.color(memory_hard + ' MB', :red),
          :diskspace_soft   => @h.color(diskspace_soft + ' MB', :yellow),
          :diskspace_hard   => @h.color(diskspace_hard + ' MB', :red) 
        }
        
        #vps_list << [
        #  vpsid + "e[32mDONEe[0m",
        #  ip,
        #  memory_soft,
        #  memory_hard,
        #  diskspace_soft,
        #  diskspace_hard
        #]
      end
      
      #@h.say Hirb::Helpers::Table.render(vps_list, 
      #  :fields => [:vpsid, :ip, :memory_soft, :memory_hard, :diskspace_soft, :diskspace_hard],
      #  :headers => {:vpsid => 'CTID', :ip => 'IP', :memory_soft => 'MEMORY SOFT LIMIT', :memory_hard => 'MEMORY HARD LIMIT', :diskspace_soft => 'DISKSPACE SOFT LIMIT', :diskspace_hard => 'DISKSPACE HARD LIMIT'}
      #)
      HighLine.use_color=false
    
      @h.say sprintf("%-8s\t%-15s\t%-20s\t%-20s\t%-20s\t%-20s", 'CTID', 'IP', 'MEMORY SOFT LIMIT', 'MEMORY HARD LIMIT', 'DISKSPACE SOFT LIMIT', 'DISKSPACE HARD LIMIT') + "\n"
      @h.say sprintf("%-8s\t%-15s\t%-17s\t%-17s\t%-20s\t%-20s", '='*8, '='*15, '='*17, '='*17, '='*20, '='*20) + "\n"
    
      vps_list.each do |vps|
        @h.say sprintf("%-8s\t%-15s\t%-17s\t\t%-17s\t\t%-20s\t\t%-20s", vps[:vpsid], vps[:ip], vps[:memory_soft], vps[:memory_hard], vps[:diskspace_soft], vps[:diskspace_hard]) + "\n"
      end
      
      #@h.say sprintf("%10sworld", "hello");
      #@h.say @h.list vps_list, :columns_down
      
    end

    def set
      @arguments.shift
      usage = 'Usage: vzlimit set vpsid [options]'
      raise usage unless @arguments[0]
      
      vpsid = @arguments.shift
      memory = nil
      disk = nil
      
      while (arg = @arguments.shift) do
        if ['-r', '--ram'].include? arg
          memory = @arguments.shift
          
          if !(memory =~ /^[-+]?[0-9]*\.?[0-9]+:[-+]?[0-9]*\.?[0-9]+$/)
            raise "Usage: vzlimit set vpsid #{arg} soft:hard"
          end
        elsif ['-d', '--diskspace'].include? arg
          disk = @arguments.shift
          
          if !(disk =~ /^\d+:\d+$/)
            raise "Usage: vzlimit set vpsid #{arg} soft:hard"
          end
        else
          raise "Argument #{arg} not valid for action 'set'."
        end
      end
      
      if !memory &&! disk
        raise "Nothing to set."
      end
      
      # Limit memory
      limit_memory(vpsid, memory) if memory
      
      # Limit disk space
      limit_disk_space(vpsid, disk) if disk
    end
    
    def limit_memory(vpsid, memory)
      # Get separate soft and hard limit
      soft_mb, hard_mb = memory.split(':')
      
      # Desaster validation
      if soft_mb.to_f > hard_mb.to_f || soft_mb.to_f < 0 || hard_mb.to_f < 0
        raise "Memory limits not valid."
      end
      
      # Convert MB limits to mempages
      soft = (soft_mb.to_f * 1024 / 4).to_i
      hard = (hard_mb.to_f * 1024 / 4).to_i
      
      # Execute command 
      execute "vzctl set #{vpsid} --privvmpages #{soft}:#{hard} --save"
      
      # Success message
      @h.say @h.color "Memory limits have been set:", :green
      @h.say "  VPS-ID:             #{vpsid}"
      @h.say "  Soft memory limit:  #{soft_mb} MB"
      @h.say "  Hard memory limit:  #{hard_mb} MB"
    end
    
    def limit_disk_space(vpsid, disk)
      # Get separate soft and hard limit
      soft_mb, hard_mb = disk.split(':')
      
      # Desaster validation
      if soft_mb.to_f > hard_mb.to_f || soft_mb.to_f < 0 || hard_mb.to_f < 0
        raise "Disk limits not valid."
      end
      
      # Convert MB limits to 1K-block
      soft = (soft_mb.to_f * 1024).to_i
      hard = (hard_mb.to_f * 1024).to_i
      
      # Execute command 
      execute "vzctl set #{vpsid} --diskspace #{soft}:#{hard} --save"
      
      # Success message
      @h.say @h.color "Diskspace limits have been set:", :green
      @h.say "  VPS-ID:             #{vpsid}"
      @h.say "  Soft memory limit:  #{soft_mb} MB"
      @h.say "  Hard memory limit:  #{hard_mb} MB"
    end
    
    def execute(command)
      @h.say "Will execute the following command:"
      @h.say "=================================================================================="
      @h.say command
      @h.say "=================================================================================="
      
      unless @h.agree @h.color("Execute this command?", :red) + " [yes/no]  "
        raise 'User aborted.'
      end
      
      retval = system command
      raise @h.color "Error: Command execution failed, exit status #{$?}.", :red unless retval
    end
    
    def usage
      @h.say "Usage: vzlimit action [params...]"
      exit
    end

		def version
			@h.say "vzlimit version #{VERSION}"
		end

    def help
			@h.say "vzlimit list | set <veid>"
			@h.say "vzlimit set <veid>"
   		@h.say "\t[--ram, -r <soft-limit-in-mb>:<hard-limit-in-mb>]"
   		@h.say "\t[--diskspace, -d <soft-limit-in-mb>:<hard-limit-in-mb>]"
			@h.say "vzlimit --help, -h"
			@h.say "vzlimit --version, -v"
		end
  end
end