#!/bin/env python3
from axone import One
# , Vm, Datastore, Disk, Image
import sys
import os
import argparse
import base64

# import requests
import xmltodict
# import html
#
# import os
import subprocess

from dataclasses import dataclass

import datetime

STEP_PRECOPY = 0x1
STEP_FINAL   = 0x2


def parse_args():
    parser = argparse.ArgumentParser(prog='premigrate',
                                     description='Migrate ONe disks to another storage')
    # parser.add_argument('-p', '--persistent', action="store_true", help="Migrate only persistent disks")
    # parser.add_argument('-0', '--prep-run', action="store_true", help="Do only 0-step of the migration: pre-copy the disks content")
    # parser.add_argument('-1', '--final-run', action="store_true", help="Do only last step of the migration: final copy the disks content, save/restore, update db")
    # parser.add_argument('-a', '--all', action="store_true", help="Migrate all disks")
    parser.add_argument('-n', '--dry-run', action="store_true", help="Don't run actual commands or edit data")
    parser.add_argument('src_host', type=str, help="Source host (Example: 'slc-fd1-vmc1')")
    parser.add_argument('dst_host', type=str, help="Destination host (Example: 'slc-fd1-vmc2')")
    parser.add_argument('dst_path', type=str, help="Path to vm datastore (Example: '/var/lib/datastore/123/45')")
    parser.add_argument('vmid', type=str, help="Vm oid (Example: '45')")
    parser.add_argument('dsid', type=str, help="Target datastore oid (Example: '123')")
    parser.add_argument('caller', type=str, nargs='?', help="Caller script")
    args = parser.parse_args()
    # Inside the ONe driver both steps should be completed
    args.steps = STEP_PRECOPY | STEP_FINAL
    # The only thing we need to know from passed VM xml is the latest history
    # sequence number
    vmxml = base64.b64decode(sys.stdin.read()).decode('utf-8')
    args.history_seq = int(xmltodict.parse(vmxml)['VM']['HISTORY_RECORDS']['HISTORY']['SEQ'])
    return args


ZMIGRATOR = "/usr/local/bin/zmigrator"
ONEAGENT = "/usr/local/bin/oneagent"
ARGS = parse_args()
one = One()


@dataclass
class Dataset:
    host: str
    path: str


def pre_migrate_cmd(src: Dataset, dst: Dataset) -> str:
    return ["ssh", "-A", src.host, ZMIGRATOR, "-v", "-R", "-M", "-k", "-s",
            "--prep-run", "-p", src.path, f"{dst.host}:{dst.path}"]


def migrate_cmd(src: Dataset, dst: Dataset) -> str:
    return ["ssh", "-A", src.host, ZMIGRATOR, "-v", "-R", "-M", "-k", "-s",
            "-p", src.path, f"{dst.host}:{dst.path}"]


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


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


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)


def pre_migrate_disks(vm, tgtds, PERSISTENT, TRANSIENT):
    PRECOPY_DONE = False
    # Pre-migrate persistent images
    if PERSISTENT:
        for disk in [d for d in vm.disks() if d.is_persistent()]:
            src = disk.datastore()
            dst = tgtds.image
            if src.id == dst.id:
                continue
            src = Dataset(src.storage, f"{src.zpool}/one/{disk.image.id}")
            dst = Dataset(dst.storage, f"{dst.zpool}/one")

            cmd = pre_migrate_cmd(src, dst)
            PRECOPY_DONE = True

            run_or_log(cmd, f"Failed to copy the disk {disk.device}")

    # Pre-migrate non-persistent images in one shot
    if TRANSIENT:
        src = one.get_datastore(vm.dsid)
        dst = tgtds.system
        src = Dataset(src.storage, f"{src.zpool}/one/{vm.axc_vmtype}/system/{vm.id}")
        dst = Dataset(dst.storage, f"{dst.zpool}/one/{vm.axc_vmtype}/system")
        PRECOPY_DONE = True

        cmd = pre_migrate_cmd(src, dst)

        run_or_log(cmd, f"Failed to copy the dataset {src.path}")
    return PRECOPY_DONE


def copy_datastore(src_host, dst_host, dst_path, src_ds_id, dst_ds_id):
    dst_base = os.path.basename(dst_path)
    dst_dir = os.path.dirname(dst_path)
    src_dir = os.path.dirname(dst_dir) + f"/{src_ds_id}"
    run_or_panic(["ssh", dst_host, "mkdir", "-p", f"{dst_dir}"],
                 f"Couldn't create {dst_dir} on target {dst_host}")
    # Remove target dir only if it's on different hosts or datastores differ
    if src_ds_id != dst_ds_id or src_host != dst_host:
        run_or_panic(["ssh", dst_host, "rm", "-rf", f"{dst_path}"],
                     f"Couldn't clean up {dst_path} on target {dst_host}")

    if src_host != dst_host:
        cmd = f"ssh -A {src_host} 'tar -C {src_dir} --sparse -cf - {dst_base}" + \
            f"| ssh {dst_host} tar -C {dst_dir} --sparse -xf -'"
        print(f"[INFO] Running command: {cmd}", file=sys.stderr)
        if not ARGS.dry_run:
            rc = subprocess.run(cmd, shell=True, capture_output=True, text=True)
            if rc.returncode != 0:
                panic(f"Couldn't copy datastore {src_host}:{src_dir}/{dst_base} -> {dst_host}:{dst_dir}/")
        # Also seed the NEW dsid path on the source host. The oneagent restore
        # step below rewrites file-backed disk paths to the NEW dsid via -r,
        # and virsh restore opens CDROM backing files locally on src_host
        # before ONe's kvm/migrate hands the domain to dst_host. postmigrate
        # then removes it via 'rm -rf $DST_PATH on $SRC_HOST'.
        if src_ds_id != dst_ds_id:
            run_or_panic(["ssh", src_host, "mkdir", "-p", f"{dst_dir}"],
                         f"Couldn't create {dst_dir} on source {src_host}")
            run_or_panic(["ssh", src_host, "rm", "-rf", f"{dst_path}"],
                         f"Couldn't clean up {dst_path} on source {src_host}")
            run_or_panic(["ssh", src_host, "cp", "-a", f"{src_dir}/{dst_base}", dst_dir],
                         f"Couldn't seed {dst_path} on source {src_host}")
    else:
        # if we're on the same host - copy datastore locally
        run_or_panic(["ssh", dst_host, "cp", "-a", f"{src_dir}/{dst_base}", dst_dir],
                     f"Couldn't copy {src_dir}/{dst_base} to {dst_dir} on target {dst_host}")


# ##########################################################
#                      Main program
# ##########################################################
def main():
    if ARGS.caller is not None:
        print(f"Called from tm/{ARGS.caller}, but I have nothing to do", file=sys.stderr)
        sys.exit(0)

    # Setup initial variables
    PRECOPY_DONE = False
    PRECOPY = bool(ARGS.steps & STEP_PRECOPY)
    FINAL = bool(ARGS.steps & STEP_FINAL)
    ts = datetime.datetime.now().strftime("%Y%m%d%H%M%S")
    state_file = None
    vm = one.get_vm(ARGS.vmid, ARGS.history_seq - 1)
    target_ds = one.get_datastore(ARGS.dsid)
    ds = one.get_storage_pool(target_ds.zpool)
    # Enable migrating non-persistent disks only in case target datastore differs
    TRANSIENT = ds.system.id != vm.dsid
    # if there is no corresponding image datastore skip migration of persistent
    # disks. Also skip in case we don't migrate transient
    PERSISTENT = ds.image is not None and TRANSIENT
    updates = list()
    path_rewrites = list()

    if not vm.is_migratable_to(ds, PERSISTENT, TRANSIENT):
        panic("Cannot migrate the disks. Clusters incompatibility")

    # Copy datastore directory from source host to destination
    copy_datastore(ARGS.src_host, ARGS.dst_host, ARGS.dst_path, vm.dsid, ds.system.id)

    if PRECOPY:
        PRECOPY_DONE = pre_migrate_disks(vm, ds, PERSISTENT, TRANSIENT)
    else:
        PRECOPY_DONE = True

    # Save VM state on the VMC
    if FINAL and PRECOPY_DONE:
        cmd = ["ssh", vm.host.NAME, ONEAGENT, "save", vm.title]
        state_file = run_or_panic(cmd, msg=f"Failed to save VM state {vm.title} ({vm.NAME})")
        state_file = state_file.rstrip()

    # Migrate persistent images
    if FINAL and PERSISTENT:
        # for disk in vm.disks():
        #     if disk.is_persistent():
        for disk in [d for d in vm.disks() if d.is_persistent()]:
            src = disk.datastore()
            srcds = src
            dst = ds.image
            if src.id == dst.id:
                continue
            src = Dataset(src.storage, f"{src.zpool}/one/{disk.image.id}")
            dst = Dataset(dst.storage, f"{dst.zpool}/one")

            cmd = migrate_cmd(src, dst)
            errmsg = f"Failed to copy the disk {disk.device}"

            if run_or_log(cmd, msg=errmsg) is not None:
                updates.append(f"/domain/devices/disk[target/@dev={disk.device}]/source/host/@name::{ds.image.ip}"),

                one.update("vm", vm.id, [(f"{disk.xpath}/ISCSI_HOST", ds.image.ip),
                                         (f"{disk.xpath}/DATASTORE_ID", ds.image.id),
                                         (f"{disk.xpath}/DATASTORE", ds.image.zpool),
                                         (f"{disk.xpath}/POOL_NAME", ds.image.zpool)], ARGS.dry_run)
                one.update("image", disk.image.id, [("/IMAGE/DATASTORE_ID", ds.image.id),
                                                    ("/IMAGE/DATASTORE", ds.image.zpool)], ARGS.dry_run)
                # ctldel on source
                cmd = ["ssh", "-A", src.host, "sudo", "ctldel", f"{disk.source_iqn}", "1"]
                run_or_log(cmd, f"Failed to delete iscsi target {disk.source_iqn} on {src.host}")
                # ctladd on dest
                cmd = ["ssh", "-A", dst.host, "sudo", "ctladd", f"{disk.source_iqn}", "1", f"/dev/zvol/{ds.image.zpool}/one/{disk.image.id}"]
                run_or_panic(cmd, f"Failed to add iscsi target {disk.source_iqn} on {dst.host}")
                # move to to-del on source
                cmd = ["ssh", "-A", src.host, "sudo", "zfs", "rename", "-p",
                       src.path, f"{srcds.zpool}/one/to-del/{ts}-{disk.image.id}"]
                run_or_log(cmd, f"Failed to move dataset {src.path} into to-del")
                # reload ctld
                cmd = ["ssh", src.host, "sudo", "service", "ctld", "reload"]
                run_or_log(cmd, f"Failed to reload ctld on {src.host}")
                cmd = ["ssh", dst.host, "sudo", "service", "ctld", "reload"]
                run_or_panic(cmd, f"Failed to reload ctld on {dst.host}")

    # Migrate non-persistent images in one shot
    if FINAL and TRANSIENT:
        src = one.get_datastore(vm.dsid)
        srcds = src
        dst = ds.system
        src = Dataset(src.storage, f"{src.zpool}/one/{vm.axc_vmtype}/system/{vm.id}")
        dst = Dataset(dst.storage, f"{dst.zpool}/one/{vm.axc_vmtype}/system")

        cmd = migrate_cmd(src, dst)
        errmsg = f"Failed to copy the dataset {src.path}"

        if run_or_log(cmd, msg=errmsg) is not None:
            # Context CDROM (and any other file-backed disks in the system
            # datastore) got moved on-disk by copy_datastore() above, but the
            # saved libvirt XML still references the old dsid. Rewrite the
            # path prefix on restore so the guest resumes with the correct
            # <source file>.
            #
            # Safe for both same-host and cross-host: copy_datastore seeds the
            # NEW dsid path on src_host too when hosts differ, so the file
            # exists locally when virsh restore opens the backing file. For
            # cross-host, ONe's kvm/migrate uses plain 'virsh migrate --live'
            # (@shared=YES for efs_zvol_iscsi) and propagates the source XML
            # verbatim to dst — so whatever path we set on src is what dst
            # ends up with, and dst has the file at NEW path from
            # copy_datastore's cross-host tar.
            #
            # NEW side uses ONe's double-slash convention (matches what ONe
            # emits natively and what isofix produces).
            path_rewrites.append(
                f"/var/lib/one/datastores/{vm.dsid}/{vm.id}/"
                f"::/var/lib/one//datastores/{ds.system.id}/{vm.id}/"
            )
            for disk in [d for d in vm.disks() if not d.is_persistent()]:
                updates.append(f"/domain/devices/disk[target/@dev={disk.device}]/source/host/@name::{ds.system.ip}")
                one.update("vm", vm.id, [(f"{disk.xpath}/ISCSI_HOST", ds.system.ip)], ARGS.dry_run)
                # ctldel on source
                cmd = ["ssh", "-A", src.host, "sudo", "ctldel", f"{disk.source_iqn}-{vm.id}-{disk.id}", "1"]
                run_or_log(cmd, f"Failed to delete iscsi target {disk.source_iqn}-{vm.id}-{disk.id} on {src.host}")
                # ctladd on dest
                cmd = ["ssh", "-A", dst.host, "sudo", "ctladd", f"{disk.source_iqn}-{vm.id}-{disk.id}", "1",
                       f"/dev/zvol/{ds.system.zpool}/one/{vm.axc_vmtype}/system/{vm.id}/{disk.id}"]
                run_or_panic(cmd, f"Failed to add iscsi target {disk.source_iqn}-{vm.id}-{disk.id} on {dst.host}")
            # move to to-del on source in one shot
            cmd = ["ssh", "-A", src.host, "sudo", "zfs", "rename", "-p",
                   src.path, f"{srcds.zpool}/one/to-del/{ts}-{vm.id}-system"]
            run_or_log(cmd, f"Failed to move dataset {src.path} into to-del")
            # reload ctld
            cmd = ["ssh", src.host, "sudo", "service", "ctld", "reload"]
            run_or_log(cmd, f"Failed to reload ctld on {src.host}")
            cmd = ["ssh", dst.host, "sudo", "service", "ctld", "reload"]
            run_or_panic(cmd, f"Failed to reload ctld on {dst.host}")

    # Restore VM from the saved state file
    if state_file is not None:
        cmd = ["ssh", vm.host.NAME, ONEAGENT, "restore", state_file]
        for r in path_rewrites:
            cmd += ["-r", r]
        if len(updates) > 0:
            cmd = cmd + ["-u"] + updates
        run_or_log(cmd, f"Failed to restore VM state {vm.title} ({vm.NAME})")


if __name__ == '__main__':
    main()
