#!/usr/bin/env python3
#
# Copyright (c) 2022-2026 openmediavault plugin developers
#
# This file is licensed under the terms of the GNU General Public
# License version 2. This program is licensed "as is" without any
# warranty of any kind, whether express or implied.
#
# version: 0.2.0
#
import argparse
import glob
import os
import re
import shutil
import subprocess
import sys
import time
import xml.etree.ElementTree as ET
from datetime import datetime, timezone, timedelta
from uuid import uuid4

LOG_FILE = "/var/log/omv-backup-vm.log"

def log(message):
    utc_offset = -time.timezone
    if time.localtime().tm_isdst:
        utc_offset += 3600
    tz_info = timezone(timedelta(seconds=utc_offset))
    now = datetime.now(tz_info).strftime('%Y-%m-%d %H:%M:%S%z')
    log_message = f"[{now}] [restore] {message}"
    print(log_message)
    with open(LOG_FILE, "a") as f:
        f.write(log_message + "\n")

def check_requirements():
    if os.geteuid() != 0:
        log("This script must be executed as root or using sudo.")
        sys.exit(99)
    if not shutil.which("virsh"):
        log("virsh command not found!")
        sys.exit(2)

def parse_args():
    parser = argparse.ArgumentParser(description="Restore a virtual machine from backup.")
    parser.add_argument("-d", dest="restore_from_dir", required=True, help="Set restore directory")
    parser.add_argument("-r", dest="restore_vm", required=True, help="VM name to restore TO")
    parser.add_argument("-t", dest="date", required=True, help="Datetime to restore from (YYYY-MM-DD_HH-MM-SS)")
    parser.add_argument("-v", dest="backup_vm", required=True, help="VM name to restore FROM")
    parser.add_argument("-l", dest="restore_to_dir", required=False, help="Optional restore location")
    return parser.parse_args()

def get_qemu_img_format(filename):
    try:
        result = subprocess.run(
            ["qemu-img", "info", "-U", filename],
            capture_output=True,
            text=True,
            check=True
        )
        match = re.search(r"file format:\s+(\w+)", result.stdout)
        if match:
            return match.group(1)
    except subprocess.CalledProcessError as e:
        print(f"Error running qemu-img: {e}")
    except Exception as e:
        print(f"Unexpected error: {e}")

    return "qcow2"

def find_backup(restore_from_dir, backup_vm, date):
    """Locate the backup for a given restore-point date.

    Returns (path, kind, chain_root) where kind is 'incremental' or 'legacy'.
    Incremental restore points are entries inside a chain manifest; the chain
    directory is named after the chain's full (root) backup date. Legacy
    backups are self-contained timestamped directories with no manifest.
    """
    base = os.path.join(restore_from_dir, backup_vm)
    if os.path.isdir(base):
        for entry in sorted(os.listdir(base)):
            chain = os.path.join(base, entry)
            manifest = os.path.join(chain, "manifest")
            if not os.path.isfile(manifest):
                continue
            with open(manifest) as f:
                for line in f:
                    parts = line.strip().split("|")
                    if len(parts) >= 2 and parts[1] == date:
                        return chain, "incremental", entry
    # fall back to a legacy self-contained backup directory
    direct = os.path.join(base, date)
    if os.path.isdir(direct):
        return direct, "legacy", date
    return None, None, None

def load_xml(xml_path, restore_vm, backup_vm):
    """Copy + load the saved domain XML, applying common restore edits."""
    temp_xml = f"/tmp/{uuid4()}.xml"
    shutil.copy(xml_path, temp_xml)
    tree = ET.parse(temp_xml)
    root = tree.getroot()

    # Modify VM name and UUID
    root.find("./name").text = restore_vm
    root.find("./uuid").text = str(uuid4())

    # remove mac address if restoring to a different VM name
    if restore_vm != backup_vm:
        for interface in root.findall(".//interface"):
            mac = interface.find("mac")
            if mac is not None:
                interface.remove(mac)

    tree.write(temp_xml)
    return temp_xml

def file_disks(root):
    """(index, <disk>, <source>, target dev) for file-backed disks, skipping
    cdroms/floppies, which are not part of a backup."""
    idx = 0
    for disk in root.findall("./devices/disk"):
        if disk.get("device", "disk") != "disk" or disk.get("type") != "file":
            continue
        source = disk.find("source")
        if source is None or "file" not in source.attrib:
            continue
        idx += 1
        target = disk.find("target")
        dev = target.get("dev", f"disk{idx}") if target is not None else f"disk{idx}"
        yield idx, disk, source, dev

def plan_nvram(root, backup_dir, restore_vm, backup_vm, date):
    """Work out where the restored VM's UEFI variable store goes.

    Returns (backup_file or None, new_path or None). The VM never keeps the
    original VM's NVRAM path unless it is restored under the same name.
    """
    nvram = root.find("./os/nvram")
    if nvram is None or not (nvram.text or "").strip():
        return None, None
    old_path = nvram.text.strip()
    if restore_vm == backup_vm:
        new_path = old_path
    else:
        new_path = os.path.join(os.path.dirname(old_path), f"{restore_vm}_VARS.fd")
    nvram.text = new_path
    backup_file = os.path.join(backup_dir, f"nvram_{date}.fd")
    if not os.path.isfile(backup_file):
        # older backups have no copy; libvirt recreates it from the template
        backup_file = None
    return backup_file, new_path

def nvram_targets(src, dst, restore_vm, backup_vm):
    """NVRAM paths that must not already exist. Under the original name (the
    VM itself is gone, see main) a leftover variable store is just replaced."""
    if src and dst and restore_vm != backup_vm:
        return [dst]
    return []

def check_targets(paths):
    """Refuse to overwrite anything that already exists."""
    seen = set()
    for path in paths:
        if path in seen:
            log(f"Two disks would be restored to the same file '{path}'!")
            sys.exit(12)
        seen.add(path)
        if os.path.exists(path):
            log(f"Restore target '{path}' already exists! Choose another VM name or restore location.")
            sys.exit(12)

def find_legacy_image(restore_from_path, old_path, dev, date):
    """Backup image of one disk in a self-contained backup directory."""
    base = os.path.basename(old_path)
    # a disk sharing its file name with an earlier disk is stored with its
    # target device in the name
    for name in (f"{base}_{dev}_{date}.bak", f"{base}_{date}.bak"):
        candidate = os.path.join(restore_from_path, name)
        if os.path.isfile(candidate):
            return candidate
    return None

#############################################################################
# legacy (full, self-contained) restore
#############################################################################
def restore_legacy(restore_from_path, restore_to_dir, backup_vm, restore_vm, date):
    restore_xml = os.path.join(restore_from_path, f"{backup_vm}_{date}.xml")
    temp_xml = load_xml(restore_xml, restore_vm, backup_vm)

    tree = ET.parse(temp_xml)
    root = tree.getroot()
    jobs = []

    for idx, disk, source, dev in file_disks(root):
        old_path = source.attrib["file"]
        backup_path = find_legacy_image(restore_from_path, old_path, dev, date)
        if not backup_path:
            log(f"No backup image found for disk '{dev}' ({old_path}), skipping")
            continue
        new_ext = get_qemu_img_format(backup_path)
        new_path = os.path.join(restore_to_dir if restore_to_dir else os.path.dirname(old_path), f"{restore_vm}_{idx:03d}.{new_ext}")
        source.attrib["file"] = new_path
        # full backups are converted to qcow2 unless made with -s, so the
        # driver must follow the format of the backup image
        driver = disk.find("driver")
        if driver is not None:
            driver.set("type", new_ext)
        jobs.append((backup_path, new_path, new_ext))

    nvram_src, nvram_dst = plan_nvram(root, restore_from_path, restore_vm, backup_vm, date)
    check_targets([j[1] for j in jobs] + nvram_targets(nvram_src, nvram_dst, restore_vm, backup_vm))
    tree.write(temp_xml)

    for backup_path, new_path, fmt in jobs:
        log(f"Restoring {backup_path} to {new_path}")
        if fmt == "qcow2":
            subprocess.run(["qemu-img", "convert", "-O", "qcow2", backup_path, new_path], check=True)
        else:
            subprocess.run(["qemu-img", "dd", f"if={backup_path}", f"of={new_path}", "bs=1M"], check=True)
    restore_nvram(nvram_src, nvram_dst)

    return temp_xml

#############################################################################
# incremental (chain) restore
#############################################################################
def restore_incremental(chain_dir, chain_root, restore_to_dir, backup_vm, restore_vm, date):
    # the domain XML was saved against the chain's full (root) backup
    restore_xml = os.path.join(chain_dir, f"{backup_vm}_{chain_root}.xml")
    temp_xml = load_xml(restore_xml, restore_vm, backup_vm)

    tree = ET.parse(temp_xml)
    root = tree.getroot()
    jobs = []

    for idx, disk, source, dev in file_disks(root):
        old_path = source.attrib["file"]

        # the restore-point image for this disk is the full/incr file in the
        # chain dir named for this date and target device. qemu-img convert
        # follows its backing chain to flatten state at that point.
        matches = glob.glob(os.path.join(chain_dir, f"*_{date}_{dev}.qcow2"))
        if not matches:
            log(f"No backup image found for disk '{dev}' at {date}, skipping")
            continue
        restore_point = matches[0]

        new_dir = restore_to_dir if restore_to_dir else os.path.dirname(old_path)
        new_path = os.path.join(new_dir, f"{restore_vm}_{idx:03d}.qcow2")
        source.attrib["file"] = new_path
        driver = disk.find("driver")
        if driver is not None:
            driver.set("type", "qcow2")
        jobs.append((restore_point, new_path))

    nvram_src, nvram_dst = plan_nvram(root, chain_dir, restore_vm, backup_vm, date)
    check_targets([j[1] for j in jobs] + nvram_targets(nvram_src, nvram_dst, restore_vm, backup_vm))
    tree.write(temp_xml)

    for restore_point, new_path in jobs:
        log(f"Restoring {os.path.basename(restore_point)} (chain) to {new_path}")
        subprocess.run(["qemu-img", "convert", "-O", "qcow2", restore_point, new_path], check=True)
    restore_nvram(nvram_src, nvram_dst)

    return temp_xml

def restore_nvram(src, dst):
    if src and dst:
        log(f"Restoring NVRAM {os.path.basename(src)} to {dst}")
        shutil.copy(src, dst)

def vm_exists(name):
    return subprocess.run(["virsh", "dominfo", name], capture_output=True).returncode == 0

def define_vm(temp_xml):
    subprocess.run(["virsh", "define", temp_xml], check=True)

def main():
    check_requirements()
    args = parse_args()
    log(f"Restore FROM dir :: {args.restore_from_dir}")
    restore_to_dir = None
    if args.restore_to_dir:
        log(f"Restore TO dir :: {args.restore_to_dir}")
        restore_to_dir = args.restore_to_dir
    log(f"Date :: {args.date}")
    log(f"Backup VM :: {args.backup_vm}")
    log(f"Restore VM :: {args.restore_vm}")

    if restore_to_dir and not os.path.isdir(restore_to_dir):
        log(f"No restore directory '{restore_to_dir}' found!")
        sys.exit(11)

    if vm_exists(args.restore_vm):
        log(f"A VM named '{args.restore_vm}' already exists! Choose another name or delete it first.")
        sys.exit(13)

    chain_dir, kind, chain_root = find_backup(args.restore_from_dir, args.backup_vm, args.date)
    if not chain_dir:
        log(f"No backup found for '{args.backup_vm}' at {args.date}!")
        sys.exit(10)

    if kind == "incremental":
        log(f"Restoring from incremental chain :: {chain_dir} (point {args.date})")
        temp_xml = restore_incremental(chain_dir, chain_root, restore_to_dir,
                                       args.backup_vm, args.restore_vm, args.date)
    else:
        log(f"Restoring from full backup :: {chain_dir}")
        temp_xml = restore_legacy(chain_dir, restore_to_dir,
                                  args.backup_vm, args.restore_vm, args.date)

    define_vm(temp_xml)
    os.remove(temp_xml)
    log("Done.")

if __name__ == "__main__":
    main()
