#!{{ parameter_dict.get('python-path') }}
# BEWARE: This file is operated by slapos node
# BEWARE: It will be overwritten automatically

import hashlib
import os
import socket
import subprocess
import gzip
import shutil
from random import shuffle
import glob
import re
import json
import operator

# XXX: give all of this through parameter, don't use this as template, but as module
qemu_img_path = {{ repr(parameter_dict["qemu-img-path"]) }}
qemu_path = {{ repr(parameter_dict["qemu-path"]) }}
disk_size = {{ repr(parameter_dict["disk-size"]) }}
disk_type = {{ repr(parameter_dict["disk-type"]) }}

network_adapter = {{ repr(parameter_dict["network-adapter"]) }}

socket_path = '{{ parameter_dict.get("socket-path") }}'
nbd_list = (('{{ parameter_dict.get("nbd-host") }}',
              {{ parameter_dict.get("nbd-port") }}),
            ('{{ parameter_dict.get("nbd2-host") }}',
              {{ parameter_dict.get("nbd2-port") }}))
default_cdrom_iso = '{{ parameter_dict.get("default-cdrom-iso") }}'

nat_rules = '{{ parameter_dict.get("nat-rules") }}'.strip()
use_tap = '{{ parameter_dict.get("use-tap") }}'.lower()
use_nat = '{{ parameter_dict.get("use-nat") }}'.lower()
set_nat_restrict = '{{ parameter_dict.get("nat-restrict") }}'.lower()
enable_vhost = '{{ parameter_dict.get("enable-vhost") }}'.lower()
tap_interface = '{{ parameter_dict.get("tap-interface") }}'
listen_ip = '{{ parameter_dict.get("ipv4") }}'
mac_address = '{{ parameter_dict.get("mac-address") }}'
tap_mac_address = '{{ parameter_dict.get("tap-mac-address") }}'
tap_ipv6_addr = '{{ parameter_dict.get("tap-ipv6-addr") }}'
numa_list = '{{ parameter_dict.get("numa", "") }}'.split()
ram_max_size = '{{ parameter_dict.get("ram-max-size") }}'
init_ram_size = {{ parameter_dict.get("init-ram-size") }}
init_smp_count = {{ parameter_dict.get("init-smp-count") }}
pid_file_path = '{{ parameter_dict.get("pid-file-path") }}'
external_disk_number = {{ parameter_dict.get("external-disk-number") }}
external_disk_size = {{ parameter_dict.get("external-disk-size") }}
external_disk_format = {{ repr(parameter_dict["external-disk-format"]) }}
external_disk = {{ parameter_dict['external-disk'] }}
if int(external_disk_number) > 0 and len(external_disk) > 0:
  raise ValueError("external-disk-number and external-disk are mutually exclusive.")
instance_root = '{{ parameter_dict['instance-root'] }}'
disk_storage_dict = {}
disk_storage_list = """{{ parameter_dict.get("disk-storage-list") }}""".split('\n')
map_storage_list = []
etc_directory = '{{ parameter_dict.get("etc-directory") }}'.strip()
httpd_port = {{ parameter_dict.get("httpd-port") }}
netcat_bin = '{{ parameter_dict.get("netcat-binary") }}'.strip()
cluster_doc_host = '{{ parameter_dict.get("cluster-doc-host") }}'
cluster_doc_port = {{ parameter_dict.get("cluster-doc-port") }}
auto_ballooning = '{{ parameter_dict.get("auto-ballooning") }}' in ('true', 'True', '1')
vm_name = '{{ parameter_dict.get("name") }}'

# If a device (ie.: /dev/sdb) is provided, use it instead
# the disk_path with disk_format
disk_info_list = []
for disk_device_path in '{{ parameter_dict.get("disk-device-path", "") }}'.split():
  if disk_device_path.startswith("/dev/"):
    disk_info_list.append({
      'path': disk_device_path,
      'format': "raw",
      'aio': "{{ parameter_dict['disk-aio'] or 'native' }}",
      'cache': "{{ parameter_dict['disk-cache'] or 'none' }}",
    })

if not disk_info_list:
  disk_info_list.append({
{%- for k in 'path', 'format', 'aio', 'cache' %}
{%-   set v = parameter_dict['disk-' + k] %}
{%-   if v %}
    {{ repr(k) }}: {{ repr(v) }},
{%-   endif %}
{%- endfor %}
  })

smp_max_count = {{ parameter_dict.get("smp-max-count") }}
machine_options = '{{ parameter_dict.get("machine-options", "") }}'.strip()
cpu_model = '{{ parameter_dict.get("cpu-model") }}'.strip()

enable_device_hotplug = '{{ parameter_dict.get("enable-device-hotplug") }}'.lower()

logfile = '{{ parameter_dict.get("log-file") }}'

boot_image_url_list_json_config = '{{ parameter_dict.get("boot-image-url-list-json-config") }}'
boot_image_url_select_json_config = '{{ parameter_dict.get("boot-image-url-select-json-config") }}'
virtual_hard_drive_url_json_config = '{{ parameter_dict.get("virtual-hard-drive-url-json-config") }}'
virtual_hard_drive_gzipped = '{{ parameter_dict.get("virtual-hard-drive-gzipped") }}'.strip().lower()

def md5Checksum(file_path):
    with open(file_path, 'rb') as fh:
        m = hashlib.md5()
        while True:
            data = fh.read(8192)
            if not data:
                break
            m.update(data)
        return m.hexdigest()

def getSocketStatus(host, port):
  s = None
  for af, socktype, proto, canonname, sa in socket.getaddrinfo(
      host, port, socket.AF_UNSPEC, socket.SOCK_STREAM):
    try:
      s = socket.socket(af, socktype, proto)
      s.connect(sa)
      return s
    except socket.error:
      if s:
        s.close()
        s = None

def getMapStorageList(disk_storage_dict, external_disk_number):
  map_disk_file = os.path.join(etc_directory, '.data-disk-ids')
  last_disk_num_f = os.path.join(etc_directory, '.data-disk-amount')
  id_list = []
  last_amount = 0
  map_f_exist = os.path.exists(map_disk_file)
  if os.path.exists(last_disk_num_f):
    with open(last_disk_num_f, 'r') as lf:
      last_amount = int(lf.readline())
  if map_f_exist:
    with open(map_disk_file, 'r') as mf:
      # ID are writen in one line: data1 data3 data2 ...
      content = mf.readline()
      for id in content.split(' '):
        if id in disk_storage_dict:
          id_list.append(id)
        else:
          # Mean that this disk path has been removed (disk unmounted)
          last_amount -= 1
  for key in disk_storage_dict:
    if not key in id_list:
      id_list.append(key)

  if id_list:
    if not map_f_exist:
      # shuffle the list to not write disk in data1, data2, ... everytime
      shuffle(id_list)
    if external_disk_number < last_amount:
      # Drop created disk is not allowed
      external_disk_number = last_amount
    with open(map_disk_file, 'w') as mf:
      mf.write(' '.join(id_list))
    with open(last_disk_num_f, 'w') as lf:
      lf.write('%s' % external_disk_number)
  return id_list, external_disk_number

# Use downloaded virtual-hard-drive-url
if len(disk_info_list) == 1 and not os.path.exists(disk_info_list[0]['path']) and virtual_hard_drive_url_json_config != '':
  print('Using virtual hard drive...')
  with open(virtual_hard_drive_url_json_config) as fh:
    image_config = json.load(fh)
  if image_config['error-amount'] == 0:
    image = image_config['image-list'][0]
    downloaded_image = os.path.join(image_config['destination-directory'], image['destination'])
    # previous version was using disk in place, but here it would result with
    # redownload, so copy it
    if virtual_hard_drive_gzipped == 'true':
      try:
        with open(disk_info_list[0]['path'], 'wb') as d_fh:
          with gzip.open(downloaded_image, 'rb') as s_fh:
            shutil.copyfileobj(s_fh, d_fh)
      except Exception:
        if os.path.exists(disk_info_list[0]['path']):
          os.unlink(disk_info_list[0]['path'])
        raise
    else:
      try:
        shutil.copyfile(downloaded_image, disk_info_list[0]['path'])
      except Exception:
        if os.path.exists(disk_info_list[0]['path']):
          os.unlink(disk_info_list[0]['path'])
        raise
  else:
    raise ValueError('virtual-hard-drive-url not ready yet')

# Create disk if doesn't exist
# XXX: move to Buildout profile
if len(disk_info_list) == 1 and not disk_info_list[0]['path'].startswith("/dev/") and not os.path.exists(disk_info_list[0]['path']):
  print('Creating virtual hard drive...')
  subprocess.check_call([qemu_img_path, 'create' ,'-f', disk_info_list[0]['format'],
      disk_info_list[0]['path'], '%sG' % disk_size])
  print('Done.')

# Check and create external disk
additional_disk_list = []
for storage in disk_storage_list:
  if storage:
    key, val = storage.split(' ')
    disk_storage_dict[key.strip()] = val.strip()

map_storage_list, external_disk_number = getMapStorageList(disk_storage_dict,
                                                      int(external_disk_number))

assert len(map_storage_list) == len(disk_storage_dict)
if disk_storage_dict:
  if external_disk_number > 0:
    index = 0
    while (index < len(disk_storage_dict)) and (index < external_disk_number):
      path = disk_storage_dict[map_storage_list[index]]
      if os.path.exists(path):
        disk_filepath = os.path.join(path,
                                  'kvm_virtual_disk.%s' % external_disk_format)
        disk_list = glob.glob('%s.*' % os.path.join(path, 'kvm_virtual_disk'))
        if disk_list == []:
          print('Creating one additional virtual hard drive...')
          process = subprocess.check_call([qemu_img_path, 'create' ,'-f', external_disk_format,
              disk_filepath, '%sG' % external_disk_size])
        else:
          # Cannot change or recreate if disk is exists
          disk_filepath = disk_list[0]
        additional_disk_list.append(disk_filepath)
      else:
        print('Data folder %s was not used to create external disk %r' % (index +1))
      index += 1

# Generate network parameters
# XXX: use_tap should be a boolean
tap_network_parameter = []
nat_network_parameter = []
numa_parameter = []
number = -1
if use_nat == 'true':
  number += 1
  rules = 'user,id=lan%s' % number
  for rule in nat_rules.split():
    proto = 'tcp'
    rule = rule.split(':')
    if len(rule) == 1:
      port = int(rule[0])
    elif len(rule) == 2:
      proto = rule[0]
      port = int(rule[1])

    rules += ',hostfwd={proto}:{hostaddr}:{hostport}-:{guestport}'.format(
      proto=proto,
      hostaddr=listen_ip,
      hostport=port + 10000,
      guestport=port
    )

  if httpd_port > 0:
    rules += ',guestfwd=tcp:10.0.2.100:80-cmd:%s %s %s' % (netcat_bin,
                                                        listen_ip, httpd_port)
  if cluster_doc_host and cluster_doc_port > 0:
    rules += ',guestfwd=tcp:10.0.2.101:443-cmd:%s %s %s' % (netcat_bin,
                                           cluster_doc_host, cluster_doc_port)
  if set_nat_restrict == 'true':
    rules += ',restrict=on'
  if use_tap == 'true' and tap_ipv6_addr != '':
    rules += ',ipv6=off'
  nat_network_parameter = ['-netdev', rules,
          '-device', '%s,netdev=lan%s,mac=%s' % (network_adapter, number, mac_address)]
if use_tap == 'true':
  number += 1
  vhost = ''
  if enable_vhost == 'true':
    vhost = ',vhost=on'
  tap_network_parameter = ['-netdev',
          'tap,id=lan%s,ifname=%s,script=no,downscript=no%s' % (number,
            tap_interface, vhost),
          '-device', '%s,netdev=lan%s,mac=%s' % (network_adapter, number, tap_mac_address)]

smp = '%s,maxcpus=%s' % (init_smp_count, smp_max_count)
ram = '%sM,slots=128,maxmem=%sM' % (init_ram_size, ram_max_size)

kvm_argument_list = [qemu_path,
  '-enable-kvm', '-smp', smp, '-name', vm_name, '-m', ram, '-vga', 'std',
  '-vnc', '%s:1,password=on,websocket=on' % listen_ip,
  '-boot', 'order=cd,menu=on',
  '-qmp', 'unix:%s,server,nowait' % socket_path,
  '-pidfile', pid_file_path, '-msg', 'timestamp=on',
  '-D', logfile,
  '-nodefaults',
  # switch to tablet mode for the mouse to have it synced with a client, see https://wiki.gentoo.org/wiki/QEMU/Options#USB
  '-usbdevice', 'tablet',
]
for disk_info in disk_info_list:
  kvm_argument_list += (
    '-drive',
    'file=%s,if=%s,discard=on%s' % (
      disk_info['path'], disk_type,
      ''.join(',%s=%s' % x for x in disk_info.items() if x[0] != 'path'))
  )

rgx = re.compile('^[\w*\,][\=\d+\-\,\w]*$')
for numa in numa_list:
  if rgx.match(numa):
    numa_parameter.extend(['-numa', numa])
kvm_argument_list += numa_parameter

if tap_network_parameter == [] and nat_network_parameter == []:
  print('Warning : No network interface defined.')
else:
  kvm_argument_list += nat_network_parameter + tap_network_parameter

for disk in additional_disk_list:
  kvm_argument_list.extend([
            '-drive', 'file=%s,if=%s' % (disk, disk_type)])

# support external-disk parameter
# allow empty index if only one disk is provided
if len(external_disk) > 1:
  for key, value in external_disk.items():
    if 'index' not in value:
      raise ValueError('index is missing and more than one disk is present in external-disk configuration')
for disk_info in sorted(external_disk.values(), key=operator.itemgetter('index')):
  if disk_info['path'].startswith('rbd:') or disk_info['path'].startswith('/'):
    path = disk_info['path']
  else:
    path = os.path.join(instance_root, disk_info['path'])

  drive_argument_list = [
    "file=%s" %( path,),
    "if=%s" % (disk_type,),
    "cache=%s" % (disk_info.get('cache', 'writeback'),)
  ]
  if disk_info.get('format', 'autodetect') != 'autodetect':
    drive_argument_list.append(
      "format=%s" % (disk_info['format'],)
    )
  kvm_argument_list += (
    '-drive',
    ",".join(drive_argument_list)
  )

if auto_ballooning:
  kvm_argument_list.extend(['-device', 'virtio-balloon-pci,id=balloon0'])

machine_option_list = machine_options.split(',')
if machine_options and len(machine_option_list) > 0:
  name = 'type'
  if '=' in machine_option_list[0]:
    name, val = machine_option_list[0].split('=')
  else:
    val = machine_option_list[0]
  machine_option_list[0] = 'type=%s' % val
  if name == 'type':
    machine = ''
    for option in machine_option_list:
      key, val = option.split('=')
      machine += ',%s=%s' % (key, val)

    kvm_argument_list.extend(['-machine', machine])

if cpu_model:
  rgx = re.compile('^[\w*\,-_][\=\d+\-\,\w]*$')
  if rgx.match(cpu_model):
    kvm_argument_list.extend(['-cpu', cpu_model])

def handle_image(config, name):
  with open(config) as fh:
    image_config = json.load(fh)
  if image_config['error-amount'] == 0:
    for image in sorted(image_config['image-list'], key=lambda k: k['image-number']):
      destination = os.path.join(image_config['destination-directory'], image['destination'])
      if os.path.exists(destination):
        kvm_argument_list.extend([
          '-drive',
          'file=%s,media=cdrom' % (destination,)
        ])
      else:
       raise ValueError('%s not ready yet' % (name,))
  else:
    raise ValueError('%s not ready yet' % (name,))

# Try to connect to NBD server (and second nbd if defined).
# If not available, don't even specify it in qemu command line parameters.
# Reason: if qemu starts with unavailable NBD drive, it will just crash.
for nbd_ip, nbd_port in nbd_list:
  if nbd_ip and nbd_port:
    s = getSocketStatus(nbd_ip, nbd_port)
    if s is None:
      # NBD is not available : launch kvm without it
      print('Warning : Nbd is not available.')
    else:
      # NBD is available
      # We close the NBD socket else qemu won't be able to use it apparently
      s.close()
      kvm_argument_list.extend([
          '-drive',
          'file=nbd:[%s]:%s,media=cdrom' % (nbd_ip, nbd_port)])
else:
  # Note: Do not get tempted to use virtio-scsi-pci, as it does not work with
  #       Debian installation CDs, rendering it uninstallable
  if boot_image_url_select_json_config:
    # Support boot-image-url-select
    handle_image(boot_image_url_select_json_config, 'boot-image-url-select')
  if boot_image_url_list_json_config:
    # Support boot-image-url-list
    handle_image(boot_image_url_list_json_config, 'boot-image-url-list')
  # Always add by default the default image
  kvm_argument_list.extend([
      '-drive', 'file=%s,media=cdrom' % default_cdrom_iso
  ])


print('Starting KVM: \n %s' % ' '.join(kvm_argument_list))
os.execv(qemu_path, kvm_argument_list)
