#!/usr/bin/env python3

import argparse
import datetime
import os
import platform
import re
import signal
import shutil
import subprocess
import sys
import time

import dbus
from dbus.mainloop.glib import DBusGMainLoop
from gi.repository import GLib


BUS_NAME = "com.kylin.systemupgrade"
OBJECT_PATH = "/com/kylin/systemupgrade"
INTERFACE = "com.kylin.systemupgrade.interface"
DEFAULT_TIMEOUT = 1800
DEFAULT_DETECT_RETRIES = 1
DEFAULT_DETECT_RETRY_DELAY = 10
DEFAULT_DEBUG_URL = "http://ostree.kylin.com/repo/kylin"
DEFAULT_DEBUG_SOURCE_CONFIG = "/opt/kylin-software-properties/ostree-conf/kylin.conf"
DEFAULT_DEBUG_FALLBACK_CONFIG = "/etc/ostree/remotes.d/kylin.conf"
DEFAULT_DEBUG_PUSH_FILE = "/opt/kylin-software-properties/ostree-important.list"

INSTALL_METHOD = "DeployLatestUpdate"
INSTALL_SIGNALS = ("DeployUpdatFinished", "UpdateInstallFinished")
DETECT_METHOD = "UpdateDetect"
DETECT_SIGNALS = ("UpdateDetectFinished",)
DOWNLOAD_METHOD = "UpdateDownloadAll"
DOWNLOAD_SIGNALS = ("UpdateDownloadFinished",)
CANCEL_DOWNLOAD_METHOD = "CancelDownload"
SELF_UPGRADE_ERROR_CODES = {"#0005"}


class CommandError(RuntimeError):
    pass


def build_parser():
    parser = argparse.ArgumentParser(
        description="Trigger system update detect/download/install through D-Bus."
    )
    subparsers = parser.add_subparsers(dest="command", required=True)

    def add_common_arguments(subparser, include_mode=False, include_detect_retry=False):
        subparser.add_argument(
            "--timeout",
            type=int,
            default=DEFAULT_TIMEOUT,
            help=f"Timeout in seconds for each signal wait (default: {DEFAULT_TIMEOUT}).",
        )
        if include_detect_retry:
            subparser.add_argument(
                "--detect-retries",
                type=int,
                default=DEFAULT_DETECT_RETRIES,
                help=(
                    "Extra detect retries when UpdateDetectFinished reports updater "
                    f"self-upgrade completion (default: {DEFAULT_DETECT_RETRIES})."
                ),
            )
            subparser.add_argument(
                "--detect-retry-delay",
                type=int,
                default=DEFAULT_DETECT_RETRY_DELAY,
                help=(
                    "Seconds to wait before retrying detect after updater self-upgrade "
                    f"(default: {DEFAULT_DETECT_RETRY_DELAY})."
                ),
            )
        if include_mode:
            subparser.add_argument(
                "--mode",
                choices=("reboot", "shutdown"),
                default="reboot",
                help="Install mode passed to DeployLatestUpdate (default: reboot).",
            )

    add_common_arguments(
        subparsers.add_parser(
            "detect", help="Run UpdateDetect and wait for UpdateDetectFinished."
        ),
        include_detect_retry=True,
    )
    add_common_arguments(
        subparsers.add_parser(
            "download", help="Run UpdateDownloadAll and wait for UpdateDownloadFinished."
        )
    )
    add_common_arguments(
        subparsers.add_parser(
            "deploy", help="Run DeployLatestUpdate and wait for install finished signal."
        ),
        include_mode=True,
    )
    add_common_arguments(
        subparsers.add_parser("all", help="Run detect, download and deploy in sequence."),
        include_mode=True,
        include_detect_retry=True,
    )
    def add_push_parser(name, help_text):
        push_parser = subparsers.add_parser(name, help=help_text)
        push_parser.add_argument(
            "--url",
            default=DEFAULT_DEBUG_URL,
            help=f"URL written to the debug source config (default: {DEFAULT_DEBUG_URL}).",
        )
        push_parser.add_argument(
            "--content",
            help=(
                "Push content written to ostree-important.list. "
                "When omitted, the latest branch is selected from ostree remote summary."
            ),
        )

    add_push_parser(
        "push",
        "Prepare debug ostree source config and push update content.",
    )
    subparsers.add_parser(
        "push-unlock",
        help="Remove immutable attribute from debug source config and push content file.",
    )
    return parser


def to_python(value):
    if isinstance(value, dbus.Boolean):
        return bool(value)
    if isinstance(value, (dbus.Int16, dbus.Int32, dbus.Int64, dbus.UInt16, dbus.UInt32, dbus.UInt64)):
        return int(value)
    if isinstance(value, dbus.String):
        return str(value)
    if isinstance(value, (dbus.Array, list, tuple)):
        return [to_python(item) for item in value]
    return value


class UpdateTrigger:
    def __init__(self, timeout):
        DBusGMainLoop(set_as_default=True)
        self.timeout = timeout
        self.bus = dbus.SystemBus()
        self.object = None
        self.interface = None
        self.loop = None
        self.timeout_id = None
        self.signal_handlers = []
        self.signal_result = None
        self.signal_error = None
        self.refresh_proxy()

    def refresh_proxy(self):
        try:
            self.object = self.bus.get_object(BUS_NAME, OBJECT_PATH)
            self.interface = dbus.Interface(self.object, dbus_interface=INTERFACE)
        except dbus.DBusException as exc:
            raise convert_dbus_exception(exc) from exc

    def call_method(self, method_name, *args):
        self.refresh_proxy()
        method = getattr(self.interface, method_name)
        try:
            status_code, message = method(*args)
        except dbus.DBusException as exc:
            raise CommandError(f"{method_name} call failed: {exc.get_dbus_message()}") from exc

        status_code = int(status_code)
        message = str(message)
        if status_code != 0:
            raise CommandError(
                f"{method_name} returned status={status_code}, message={message}"
            )

        print(f"{method_name} accepted")

    def call_and_wait(
        self, method_name, method_args, signal_names, cancel_method_name=None
    ):
        self.signal_result = None
        self.signal_error = None
        self.loop = GLib.MainLoop()
        self.timeout_id = GLib.timeout_add_seconds(self.timeout, self._on_timeout, signal_names)
        original_sigint_handler = None

        for signal_name in signal_names:
            handler = self._build_signal_handler(signal_name)
            self.signal_handlers.append((handler, signal_name))
            self.bus.add_signal_receiver(
                handler,
                signal_name=signal_name,
                dbus_interface=INTERFACE,
                bus_name=BUS_NAME,
                path=OBJECT_PATH,
            )

        if cancel_method_name is not None:
            original_sigint_handler = signal.getsignal(signal.SIGINT)
            signal.signal(
                signal.SIGINT,
                self._build_cancel_handler(cancel_method_name),
            )

        try:
            self.call_method(method_name, *method_args)
            self.loop.run()
        finally:
            if original_sigint_handler is not None:
                signal.signal(signal.SIGINT, original_sigint_handler)
            self._clear_wait_state()

        if self.signal_error:
            raise self.signal_error

        return self.signal_result

    def wait_for_service(self, wait_seconds):
        deadline = time.monotonic() + wait_seconds
        while time.monotonic() < deadline:
            if self.bus.name_has_owner(BUS_NAME):
                self.refresh_proxy()
                return
            time.sleep(1)
        raise CommandError(
            f"Updater service did not come back on D-Bus within {wait_seconds}s. "
            "If this is the self-upgrade test flow, restart kylin-system-updater manually and retry detect."
        )

    def _build_signal_handler(self, signal_name):
        def _handler(*args):
            self.signal_result = (signal_name, tuple(to_python(arg) for arg in args))
            if self.loop is not None and self.loop.is_running():
                self.loop.quit()

        return _handler

    def _build_cancel_handler(self, cancel_method_name):
        def _handler(_signum, _frame):
            print(f"Interrupted, calling {cancel_method_name} ...", file=sys.stderr)
            try:
                self.call_method(cancel_method_name)
            except CommandError as exc:
                self.signal_error = CommandError(
                    f"Interrupted, but {cancel_method_name} failed: {exc}"
                )
            else:
                self.signal_error = CommandError(
                    f"Interrupted; {cancel_method_name} accepted"
                )

            if self.loop is not None and self.loop.is_running():
                self.loop.quit()

        return _handler

    def _on_timeout(self, signal_names):
        self.signal_error = CommandError(
            f"Timed out after {self.timeout}s while waiting for: {', '.join(signal_names)}"
        )
        if self.loop is not None and self.loop.is_running():
            self.loop.quit()
        return False

    def _clear_wait_state(self):
        if self.timeout_id is not None:
            try:
                GLib.source_remove(self.timeout_id)
            except Exception:
                pass
            self.timeout_id = None

        for handler, signal_name in self.signal_handlers:
            self.bus.remove_signal_receiver(
                handler,
                signal_name=signal_name,
                dbus_interface=INTERFACE,
                bus_name=BUS_NAME,
                path=OBJECT_PATH,
            )
        self.signal_handlers = []
        self.loop = None


def parse_signal_result(signal_name, signal_args):
    success = bool(signal_args[0]) if signal_args else False
    groups = signal_args[1] if len(signal_args) > 1 else []
    error_code = signal_args[2] if len(signal_args) > 2 else ""
    error_desc = signal_args[3] if len(signal_args) > 3 else ""
    return success, groups, error_code, error_desc


def convert_dbus_exception(exc):
    dbus_name = exc.get_dbus_name() or ""
    dbus_message = exc.get_dbus_message() or str(exc)

    if dbus_name in (
        "org.freedesktop.DBus.Error.ServiceUnknown",
        "org.freedesktop.DBus.Error.NameHasNoOwner",
    ):
        return CommandError(
            "Updater service is not running. "
            "Start com.kylin.systemupgrade first, then retry."
        )

    return CommandError(dbus_message)


def run_command(command):
    try:
        return subprocess.run(
            command,
            check=True,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
        )
    except FileNotFoundError as exc:
        raise CommandError(f"Command not found: {command[0]}") from exc
    except subprocess.CalledProcessError as exc:
        stderr = exc.stderr.strip()
        stdout = exc.stdout.strip()
        detail = stderr or stdout or f"exit status {exc.returncode}"
        raise CommandError(f"{' '.join(command)} failed: {detail}") from exc


def chattr(path, flag):
    command = ["chattr", flag, path]
    if hasattr(os, "geteuid") and os.geteuid() != 0:
        command.insert(0, "sudo")
    run_command(command)


def is_immutable(path):
    if not os.path.exists(path):
        return False

    try:
        result = run_command(["lsattr", path])
    except CommandError:
        return False

    attrs = result.stdout.split(None, 1)[0] if result.stdout.split(None, 1) else ""
    return "i" in attrs


def update_url_config(config_path, fallback_config_path, url):
    config_dir = os.path.dirname(config_path)
    os.makedirs(config_dir, exist_ok=True)
    was_immutable = is_immutable(config_path)

    if not os.path.exists(config_path):
        if not os.path.exists(fallback_config_path):
            raise CommandError(
                f"Source config does not exist and fallback is missing: {fallback_config_path}"
            )
        shutil.copy2(fallback_config_path, config_path)

    if was_immutable:
        chattr(config_path, "-i")

    with open(config_path, "r", encoding="utf-8") as config_file:
        lines = config_file.readlines()

    url_line = f"url={url}\n"
    updated = False
    for index, line in enumerate(lines):
        if line.lstrip().startswith("url="):
            prefix = line[: len(line) - len(line.lstrip())]
            lines[index] = f"{prefix}{url_line}"
            updated = True
            break

    if not updated:
        if lines and not lines[-1].endswith("\n"):
            lines[-1] += "\n"
        lines.append(url_line)

    with open(config_path, "w", encoding="utf-8") as config_file:
        config_file.writelines(lines)

    chattr(config_path, "+i")


def select_remote():
    result = run_command(["ostree", "remote", "list"])
    remotes = [line.strip() for line in result.stdout.splitlines() if line.strip()]

    if not remotes:
        raise CommandError("No ostree remotes found.")

    return remotes[0]


def parse_ostree_timestamp(value):
    normalized = value.strip()
    if normalized.endswith("Z"):
        normalized = normalized[:-1] + "+00:00"
    normalized = re.sub(r"([+-]\d{2})$", r"\1:00", normalized)

    try:
        return datetime.datetime.fromisoformat(normalized)
    except ValueError as exc:
        raise CommandError(f"Invalid ostree commit timestamp: {value}") from exc


def get_system_arch():
    try:
        result = run_command(["dpkg", "--print-architecture"])
        arch = result.stdout.strip()
        if arch:
            return arch
    except CommandError:
        pass

    machine = platform.machine().lower()
    arch_map = {
        "x86_64": "amd64",
        "aarch64": "arm64",
        "loongarch64": "loong64",
        "sw_64": "sw64",
    }
    return arch_map.get(machine, machine)


def ref_matches_arch(ref, arch):
    return ref == arch or ref.endswith(f"/{arch}")


def get_current_origin_refspec():
    result = run_command(["ostree", "admin", "status"])
    for line in result.stdout.splitlines():
        match = re.search(r"origin refspec:\s*(\S+)", line)
        if match:
            return match.group(1)

    raise CommandError("No origin refspec found in ostree admin status.")


def split_refspec(refspec):
    remote, separator, branch = refspec.partition(":")
    if not separator or not remote or not branch:
        raise CommandError(f"Invalid origin refspec: {refspec}")
    return remote, branch


def branch_field(branch, arch):
    parts = branch.split("/")
    if len(parts) < 4 or parts[-1] != arch:
        raise CommandError(f"Invalid {arch} branch format: {branch}")
    return parts[-4]


def ref_matches_current_field(ref, arch, current_field):
    parts = ref.split("/")
    return (
        len(parts) >= 4
        and parts[-1] == arch
        and parts[-4] == current_field
    )


def latest_branch_from_summary(summary_output, arch, current_field):
    latest_ref = None
    latest_commit = None
    latest_timestamp = None
    current_ref = None
    current_commit = None

    for line in summary_output.splitlines():
        ref_match = re.match(r"^\s*\*\s+(\S+)", line)
        if ref_match:
            ref = ref_match.group(1)
            current_ref = (
                ref if ref_matches_current_field(ref, arch, current_field) else None
            )
            current_commit = None
            continue

        commit_match = re.match(r"^\s+([0-9a-f]{64})\s*$", line)
        if commit_match and current_ref:
            current_commit = commit_match.group(1)
            continue

        timestamp_match = re.search(
            r"Timestamp \(ostree\.commit\.timestamp\):\s*(\S+)",
            line,
        )
        if timestamp_match and current_ref:
            timestamp = parse_ostree_timestamp(timestamp_match.group(1))
            if latest_timestamp is None or timestamp > latest_timestamp:
                latest_ref = current_ref
                latest_commit = current_commit
                latest_timestamp = timestamp

    if latest_ref is None:
        raise CommandError(
            f"No {current_field} {arch} branch timestamp found in ostree remote summary."
        )

    return latest_ref, latest_commit, latest_timestamp


def build_auto_push_content():
    arch = get_system_arch()
    origin_remote, origin_branch = split_refspec(get_current_origin_refspec())
    current_field = branch_field(origin_branch, arch)
    remote = origin_remote or select_remote()
    result = run_command(["ostree", "remote", "summary", remote])
    branch, commit, timestamp = latest_branch_from_summary(
        result.stdout, arch, current_field
    )
    return f"{remote}:{branch}", commit


def write_push_content(push_file, content, lock_file):
    if not content:
        raise CommandError("Push content is empty.")

    os.makedirs(os.path.dirname(push_file), exist_ok=True)

    if lock_file and os.path.exists(push_file):
        try:
            chattr(push_file, "-i")
        except CommandError as exc:
            print(f"Warning: failed to unlock {push_file}: {exc}", file=sys.stderr)

    with open(push_file, "w", encoding="utf-8") as output_file:
        output_file.write(content.strip() + "\n")

    if lock_file:
        chattr(push_file, "+i")


def print_debug_push_summary(url, content, commit=None):
    print("Push configuration completed.")
    print(f"URL: {url}")
    print(f"Push content: {content.strip()}")
    if commit:
        print(f"Latest commit: {commit}")
    print("Locked: source config, push file")


def unlock_debug_push_files():
    unlocked_files = []
    for path in (DEFAULT_DEBUG_SOURCE_CONFIG, DEFAULT_DEBUG_PUSH_FILE):
        if not os.path.exists(path):
            print(f"Skipped missing file: {path}")
            continue

        chattr(path, "-i")
        unlocked_files.append(path)

    if unlocked_files:
        print("Push files unlocked.")
    else:
        print("No push files found to unlock.")


def run_debug_push(args):
    update_url_config(
        DEFAULT_DEBUG_SOURCE_CONFIG,
        DEFAULT_DEBUG_FALLBACK_CONFIG,
        args.url,
    )

    content_arg = args.content
    content = content_arg.strip() if content_arg else ""
    commit = None
    if not content:
        content, commit = build_auto_push_content()

    write_push_content(DEFAULT_DEBUG_PUSH_FILE, content, True)
    print_debug_push_summary(args.url, content, commit)


def print_signal_result(signal_name, signal_args):
    success, groups, error_code, error_desc = parse_signal_result(signal_name, signal_args)

    print(
        f"{signal_name}: success={success}, groups={groups}, "
        f"error_code={error_code}, error_desc={error_desc}"
    )

    return success, groups, error_code, error_desc


def assert_success(signal_name, signal_args):
    success, groups, error_code, error_desc = print_signal_result(signal_name, signal_args)

    if not success:
        raise CommandError(
            f"{signal_name} failed: error_code={error_code}, error_desc={error_desc}"
        )

    return groups


def should_retry_detect(success, error_code):
    return not success and error_code in SELF_UPGRADE_ERROR_CODES


def run_detect(trigger, detect_retries, detect_retry_delay):
    remaining_retries = detect_retries

    while True:
        print(f"Calling {DETECT_METHOD} ...")
        signal_name, signal_args = trigger.call_and_wait(
            DETECT_METHOD, (), DETECT_SIGNALS
        )
        success, groups, error_code, error_desc = print_signal_result(signal_name, signal_args)

        if should_retry_detect(success, error_code) and remaining_retries > 0:
            print(
                f"Retrying {DETECT_METHOD} after updater self-upgrade "
                f"in {detect_retry_delay}s ({remaining_retries} retry left) ..."
            )
            time.sleep(detect_retry_delay)
            trigger.wait_for_service(detect_retry_delay)
            remaining_retries -= 1
            continue

        if not success:
            raise CommandError(
                f"{signal_name} failed: error_code={error_code}, error_desc={error_desc}"
            )

        if not groups:
            print("No updates available.")
        return groups


def run_download(trigger):
    print(f"Calling {DOWNLOAD_METHOD} ...")
    signal_name, signal_args = trigger.call_and_wait(
        DOWNLOAD_METHOD, (), DOWNLOAD_SIGNALS, cancel_method_name=CANCEL_DOWNLOAD_METHOD
    )
    return assert_success(signal_name, signal_args)


def run_deploy(trigger, mode):
    print(f"Calling {INSTALL_METHOD}({mode}) ...")
    signal_name, signal_args = trigger.call_and_wait(
        INSTALL_METHOD, (mode,), INSTALL_SIGNALS
    )
    assert_success(signal_name, signal_args)


def main():
    args = build_parser().parse_args()
    try:
        if args.command == "push":
            run_debug_push(args)
            return 0

        if args.command == "push-unlock":
            unlock_debug_push_files()
            return 0

        trigger = UpdateTrigger(args.timeout)

        def cleanup(*_args):
            raise SystemExit(1)

        signal.signal(signal.SIGINT, cleanup)
        signal.signal(signal.SIGTERM, cleanup)

        if args.command == "detect":
            run_detect(trigger, args.detect_retries, args.detect_retry_delay)
            return 0

        if args.command == "download":
            run_download(trigger)
            return 0

        if args.command == "deploy":
            run_deploy(trigger, args.mode)
            return 0

        groups = run_detect(trigger, args.detect_retries, args.detect_retry_delay)
        if not groups:
            print("Stopping after detect.")
            return 0

        run_download(trigger)
        run_deploy(trigger, args.mode)
        return 0
    except CommandError as exc:
        print(str(exc), file=sys.stderr)
        return 1


if __name__ == "__main__":
    sys.exit(main())
