#!/usr/bin/env python3
import sys
import argparse
import os
import re
import subprocess
import xmltodict
import ctypes

uid = os.getuid()
MOUNTPATH = f"/dev/shm/{uid}"
DISKMOUNTPATH = f"/var/tmp/{uid}"


off_t = ctypes.c_long  # YMMV
__NR_statfs = 137  # YMMV


class linux_statfs(ctypes.Structure):
    _fields_ = [
        ('f_type', ctypes.c_ulong),
        ('f_bsize', ctypes.c_ulong),
        ('f_blocks', ctypes.c_ulong),
        ('f_bfree', ctypes.c_ulong),
        ('f_bavail', ctypes.c_ulong),
    ]


def statfs(s):
    buf = ctypes.ARRAY(ctypes.c_char, 1024)()
    bstr = ctypes.create_string_buffer(s.encode())
    _statfs = ctypes.CDLL(None).syscall
    _statfs.restype = ctypes.c_int
    _statfs.argtypes = ctypes.c_long, ctypes.POINTER(ctypes.c_char), ctypes.POINTER(ctypes.c_char)

    res = _statfs(__NR_statfs, bstr, buf)
    if res == 0:
        return linux_statfs.from_buffer(buf, 0)
    return None


# ###############################################################
#                         XML Functions
# ###############################################################

def get_token(s, pos):
    """
    Function returns next token in the string
    """
    delimiters = "[/=]"
    rval = ""
    i = pos
    for c in s[pos:]:
        if c in delimiters:
            if i == pos:
                rval += c
                return (rval,i + 1)
            else:
                return (rval,i)
        rval += c
        i += 1
    return (rval,i)


DEBUG = False


def debug(msg):
    if DEBUG:
        print(f"[DEBUG] {msg}", file=sys.stderr)


def inside_sqbrackets(xml, s, pos, level=1):
    node = xml
    empty = ""
    pfx = f"{empty:>{level*4}}"
    debug(f"{pfx}inside_sqbrackets: processing {node=}")
    # get nex token
    tok, pos = get_token(s, pos)
    debug(f"{pfx}{tok=}")
    if tok == "":
        debug(f"{pfx}Empty token, Returning {node=}")
        return node, pos
    if tok == "]":
        debug(f"{pfx}Closed bracket. Returning {node=}")
        return node, pos
    if tok == "/":
        debug(f"{pfx}Skipping delimiter")
        _node, pos = inside_sqbrackets(node, s, pos, level + 1)
        debug(f"{pfx}{tok=} Returning {_node=}")
        return _node, pos
    elif tok == "=":
        # get next token and filter by value
        tok, pos = get_token(s, pos)
        if tok == "":
            # TODO: ^ add 'or any special symbol'
            raise Exception(f"Syntax error, Expected value at {pos=}")
        else:
            debug(f"{pfx}Compare value {node=} {tok=}")
            if node == tok:
                return node, pos
            else:
                return None, pos
    else:
        #  key
        if isinstance(node, dict):
            if tok in node:
                _node, pos = inside_sqbrackets(node[tok], s, pos, level + 1)
                if _node is None:
                    debug(f"{pfx}{tok=} found in {node=} but got None back")
                    return None, pos
                debug(f"{pfx}{tok=} found in {node=}. Returning {node=}")
                return node, pos
            debug(f"{pfx}{tok=} not found in {node=}")
            return None, pos

        elif isinstance(node, list):
            _node = list()
            debug(f"{pfx}Processing list {node=}")
            for x in node:
                if tok in x:
                    y, _pos = inside_sqbrackets(x[tok], s, pos)
                    if y is not None:
                        _node.append(x)
            debug(f"{pfx} List processing done. Returning {_node=}")
            return _node, _pos

        else:
            debug(f"{pfx} node is neither dict or list. Returning None")
            return None, pos


def parse(xml, s, pos):
    node = xml
    tok, pos = get_token(s, pos)
    if tok == "":
        return node, pos
    elif tok == "\x5b":
        node, pos = inside_sqbrackets(node, s, pos)
        return parse(node, s, pos)
    elif tok == "/":
        return parse(node, s, pos)
    elif tok == "]":
        return parse(node, s, pos)
    else:
        if isinstance(node, dict):
            if tok in node:
                return parse(node[tok], s, pos)
            return None, pos
        elif isinstance(node, list):
            _node = list()
            for x in node:
                if tok in x:
                    y, _pos = parse(x[tok], s, pos)
                    if y is not None:
                        _node.append(y)
            return _node, _pos
    debug(f"parse: {tok=} {pos=}")
    return node, pos


def apply_single_update_xml(xml, path, key, value):
    node = xml
    node, _ = parse(node, path, 0)
    debug(f"Got the {node=} by {path=}")
    if isinstance(node, dict):
        # If key doesn't exists. Create it
        # if key in node:
        node[key] = value
    elif isinstance(node, list):
        for x in node:
            # if key in x:
            x[key] = value
    debug(f"{node=} after changing {key=} {value=}")
    return xml


def update_xml(xml, xml_updates=None):
    if xml_updates is None:
        return
    if len(xml_updates) == 0:
        return
    for update in xml_updates:
        (path, value) = update.split("::")
        print(path, value)
        apply_single_update_xml(xml, os.path.dirname(path), os.path.basename(path), value)


_SLASH_RE = re.compile(r'/+')


def rewrite_file_paths(xml, rewrites):
    """Rewrite <source file='...'/> prefixes on all disks.

    Each rewrite is 'OLD::NEW'. ONe emits both '/var/lib/one/datastores/'
    and '/var/lib/one//datastores/' — collapse runs of '/' before matching
    so a single rewrite handles either form.
    """
    if not rewrites:
        return
    pairs = []
    for r in rewrites:
        old, new = r.split("::", 1)
        pairs.append((_SLASH_RE.sub('/', old), new))

    devices = xml.get('domain', {}).get('devices', {})
    disks = devices.get('disk', [])
    if isinstance(disks, dict):
        disks = [disks]
    for disk in disks:
        src = disk.get('source')
        if not isinstance(src, dict):
            continue
        f = src.get('@file')
        if not f:
            continue
        norm = _SLASH_RE.sub('/', f)
        for old, new in pairs:
            if norm.startswith(old):
                src['@file'] = new + norm[len(old):]
                break


# ##########################################################
#                   Args handler
# ##########################################################
def parse_args():
    parser = argparse.ArgumentParser(prog='oneagent',
                                     description='ONe agent on the VMC side which helps to do complicated things')
    parser.add_argument('-u', '--update-xml', type=str, nargs="*", help="Update XML during operation")
    parser.add_argument('-r', '--prefix-rewrite', action="append", default=[],
                        help="Rewrite <source file='...'/> prefix on restore: OLD::NEW (repeatable)")
    parser.add_argument('-p', '--vnc-password', type=str, nargs="?", help="Update VNC password after restoring")
    parser.add_argument('-n', '--dry-run', action="store_true", help="Don't run actual commands or edit data")
    parser.add_argument('action', type=str, choices=["save", "restore", "help"], help="save, restore")
    parser.add_argument('action_object', type=str, help="Action params")
    return parser.parse_args()


args = parse_args()


# ##########################################################
#                      Logging functions
# ##########################################################
def error(msg: str):
    print(f"[ERROR] {msg}", file=sys.stderr)


def panic(msg: str):
    print(f"[FATAL] {msg}", file=sys.stderr)
    sys.exit(-1)


# ##########################################################
#                      Run functions
# ##########################################################
def run(cmd: list, on_failure_fn, msg: str):
    print(f"[INFO] Running command: {cmd}", file=sys.stderr)
    if not args.dry_run:
        rc = subprocess.run(cmd, capture_output=True, text=True)
        if rc.returncode != 0:
            print(rc.stderr, file=sys.stderr)
            on_failure_fn(msg)
            return None
        return rc.stdout
    # in case of dry run return empty string to continue the script
    return ""


def run_or_log(cmd: list, msg: str):
    return run(cmd, on_failure_fn=error, msg=msg)


def run_or_panic(cmd: list, msg: str):
    return run(cmd, on_failure_fn=panic, msg=msg)


# ##########################################################
#                      Main program
# ##########################################################
def get_free_mem_mb():
    with open('/proc/meminfo') as file:
        for line in file:
            if line.startswith("MemFree"):
                return int(line.split()[1]) / 1024
    return None


def mountpoint(size_mb):
    if size_mb + 4096 >= get_free_mem_mb():
        if args.dry_run:
            return DISKMOUNTPATH
        # Use disk instead
        os.makedirs(DISKMOUNTPATH, exist_ok=True, mode=0o777)
        fsinfo = statfs(DISKMOUNTPATH)
        if fsinfo is None:
            panic("Can't find mountpoint")
        else:
            if size_mb + 8192 >= fsinfo.f_bsize * fsinfo.f_bfree / (1024 * 1024):
                panic("Not enough free space on disk to store dump")
            else:
                return DISKMOUNTPATH
    else:
        if args.dry_run:
            return MOUNTPATH
        # use memory
        os.makedirs(MOUNTPATH, exist_ok=True, mode=0o777)
        return MOUNTPATH


def save(vmname, mountpath):
    cmd = ["virsh", "save", f"{vmname}", f"{mountpath}/{vmname}.state"]
    if args.dry_run:
        return f"{mountpath}/{vmname}.state"

    rc = subprocess.run(cmd, capture_output=True, text=True)
    if rc.returncode != 0:
        errmsg = rc.stderr
        cmd = ["rm", "-f", f"{mountpath}/{vmname}.state"]
        subprocess.run(cmd, capture_output=True, text=True)
        panic(errmsg)
    return f"{mountpath}/{vmname}.state"


def dommem(vmname):
    out = run_or_panic(["virsh", "dommemstat", f"{vmname}"],
                       "Couldn't get domain memory stat")
    for line in out.split("\n"):
        if line.startswith("actual"):
            return int(line.split()[1])
    return None


def dumpxml(vmname, mountpath):
    xml = run_or_panic(["virsh", "save-image-dumpxml", f"{mountpath}/{vmname}.state", "--security-info"],
                       "Couldn't get domain XML from state file")
    return xmltodict.parse(xml)


def restore(statefile, xml_updates=None, path_rewrites=None):
    mountpath = os.path.dirname(statefile)
    # cut '.state' suffix
    vmname = os.path.basename(statefile)[:-6]

    cmd = ["virsh", "restore", f"{mountpath}/{vmname}.state"]

    if args.vnc_password is not None:
        if xml_updates is not None:
            xml_updates.append(f"/domain/devices/graphics[@type=vnc]/@passwd::{args.vnc_password}")
        else:
            xml_updates = [f"/domain/devices/graphics[@type=vnc]/@passwd::{args.vnc_password}"]

    if xml_updates is not None or path_rewrites:
        xml = dumpxml(vmname, mountpath)
        rewrite_file_paths(xml, path_rewrites)
        update_xml(xml, xml_updates)
        with open(f"{mountpath}/{vmname}.xml", "w") as f:
            f.write(xmltodict.unparse(xml))
        cmd = ["virsh", "restore", f"{mountpath}/{vmname}.state", "--xml", f"{mountpath}/{vmname}.xml"]

    run_or_panic(cmd, "Couldn't restore VM from state file")
    run_or_log(["virsh", "domtime", vmname, "--sync"],
               "Couldn't sync the time")
    run_or_log(["rm", "-f", f"{mountpath}/{vmname}.state", f"{mountpath}/{vmname}.xml"],
               f"Couldn't remove {mountpath}/{vmname}.state file")


def action_help(action_object):
    if action_object == "save":
        print(f"Saving virtual machine state into {MOUNTPATH}")
    if action_object == "restore":
        print(f"Retore virtual machine state from {MOUNTPATH}")


def migrate_save(vm, updates):
    path = mountpoint(size_mb=int(dommem(vm) / 1024))
    print(save(vm, path))
    if updates is not None:
        # xml = dumpxml(vm, path)
        # update_xml(xml, updates)
        pass


def main():

    if args.action == "help":
        action_help(args.action_object)
    if args.action == "save":
        migrate_save(args.action_object, args.update_xml)
    if args.action == "restore":
        restore(args.action_object, args.update_xml, args.prefix_rewrite)


if __name__ == '__main__':
    main()
