#!/usr/bin/env python3
"""Privileged one-shot worker for a previously validated browser OTA package."""

from __future__ import annotations

import hashlib
import json
import os
from pathlib import Path
import shutil
import signal
import subprocess
import sys
import time
from typing import Any

from app.ota.package import extract_payload, inspect_package
from app.ota.state import read_json, utc_now, write_json, write_last_result
from app.ota.versioning import SoftwareVersion

TARGET = Path("/opt/matrix-screen-controller")
RELEASES = Path("/opt/matrix-screen-controller.releases")
DATA_ROOT = Path("/var/lib/matrix-screen-controller")
RUNTIME_ROOT = Path("/run/matrix-screen-controller")
UNIT_PATH = Path("/etc/systemd/system/matrix-screen-controller.service")
SERVICE = "matrix-screen-controller.service"
REQUEST_PATH = RUNTIME_ROOT / "ota-request.json"


class UpdateFailed(RuntimeError):
    pass


def run(command: list[str], *, cwd: Path | None = None, env: dict[str, str] | None = None) -> None:
    subprocess.run(command, cwd=cwd, env=env, check=True)


def snapshot(root: Path, *, ignored: set[str] | None = None) -> dict[str, tuple[int, str]]:
    records: dict[str, tuple[int, str]] = {}
    ignored = ignored or set()
    if not root.exists():
        return records
    for path in sorted((item for item in root.rglob("*") if item.is_file()), key=lambda item: item.as_posix()):
        relative = path.relative_to(root).as_posix()
        if relative in ignored:
            continue
        digest = hashlib.sha256()
        with path.open("rb") as handle:
            for chunk in iter(lambda: handle.read(1024 * 1024), b""):
                digest.update(chunk)
        records[relative] = (path.stat().st_size, digest.hexdigest())
    return records


class Transaction:
    def __init__(self, request: dict[str, Any]) -> None:
        self.request = request
        self.job_id = checked_token(request.get("job_id"), "job_id")
        self.target_version = SoftwareVersion.parse(str(request.get("target_version")))
        self.current_version = SoftwareVersion.parse(str(request.get("current_version")))
        if self.target_version <= self.current_version:
            raise UpdateFailed("target version is not newer than the installed version")
        self.package_path = Path(str(request.get("package_path")))
        self.status_path = Path(str(request.get("status_path")))
        if self.status_path != RUNTIME_ROOT / "ota-status.json":
            raise UpdateFailed("unexpected OTA status path")
        self.work_root = Path(f"/opt/matrix-screen-controller-ota/work.{self.job_id}")
        self.extract_root = self.work_root / "extracted"
        self.release = RELEASES / f"{self.target_version}-{self.job_id}"
        self.data_candidate = DATA_ROOT.with_name(f"matrix-screen-controller.ota.{self.job_id}")
        self.data_backup = DATA_ROOT.with_name(f"matrix-screen-controller.rollback.{self.job_id}")
        self.unit_backup = self.work_root / "matrix-screen-controller.service.old"
        self.previous_target: Path | None = None
        self.previous_directory: Path | None = None
        self.program_swapped = False
        self.data_swapped = False
        self.visual: subprocess.Popen[bytes] | None = None
        self.job = {
            "id": self.job_id,
            "target_version": str(self.target_version),
            "packaged_at": str(request.get("packaged_at") or ""),
            "phase": "preparing",
            "percent": 8,
            "message": "正在准备更新",
            "started_at": str(request.get("accepted_at") or utc_now()),
            "finished_at": None,
            "error": None,
        }

    def update(self, phase: str, percent: int, message: str) -> None:
        self.job.update(phase=phase, percent=max(0, min(100, percent)), message=message)
        write_json(self.status_path, {"schema_version": 1, "active": True, "job": self.job})

    def prepare(self) -> None:
        if any(path.exists() for path in (self.work_root, self.release, self.data_candidate, self.data_backup)):
            raise UpdateFailed("OTA transaction paths already exist")
        self.work_root.mkdir(parents=True, mode=0o700)
        package = inspect_package(self.package_path, current_version=self.current_version)
        if package.target_version != self.target_version:
            raise UpdateFailed("request and package versions do not match")
        stat = shutil.disk_usage("/opt")
        required = max(package.expanded_bytes * 4, 512 * 1024 * 1024)
        if stat.free < required:
            raise UpdateFailed("not enough free space to stage and validate the OTA release")

        self.update("extracting", 12, "正在解压全量更新包")
        extract_payload(package, self.extract_root)
        software = self.extract_root / "software"
        wheelhouse = self.extract_root / "wheelhouse"
        if SoftwareVersion.parse((software / "VERSION").read_text(encoding="utf-8")) != self.target_version:
            raise UpdateFailed("extracted software VERSION does not match the package")

        self.update("dependencies", 22, "正在校验离线依赖并创建运行环境")
        run(["sha256sum", "-c", "SHA256SUMS"], cwd=wheelhouse)
        shutil.move(str(software), self.release)
        run(["python3", "-m", "venv", str(self.release / ".venv")])
        run([
            str(self.release / ".venv/bin/pip"),
            "install",
            "--no-index",
            "--disable-pip-version-check",
            "--find-links",
            str(wheelhouse),
            "-r",
            str(self.release / "requirements-dev.txt"),
        ])

        self.update("building", 42, "正在编译并测试屏幕驱动")
        run(["make", "-C", str(self.release / "app/display/native"), "clean", "all", "test"])
        run([str(self.release / ".venv/bin/python"), "-m", "compileall", "-q", str(self.release / "app")])

        self.update("testing", 55, "正在运行新版本自动化测试")
        test_env = os.environ.copy()
        test_env.update(
            MATRIX_DRIVER="mock",
            MATRIX_DATA_DIR=str(self.work_root / "test-data"),
            MATRIX_RUNTIME_DIR=str(self.work_root / "test-runtime"),
            MATRIX_SOURCE_ONLY_UPDATE_TESTS="1",
        )
        run([str(self.release / ".venv/bin/python"), "-m", "pytest", "-q"], cwd=self.release, env=test_env)
        run([str(self.release / ".venv/bin/python"), str(self.release / "scripts/dedicated_host.py"), "check"])

        self.update("migrating", 68, "正在迁移用户数据副本")
        shutil.copytree(DATA_ROOT, self.data_candidate, symlinks=True)
        runtime_candidate = self.work_root / "migration-runtime"
        migration_env = os.environ.copy()
        migration_env["PYTHONPATH"] = str(self.release)
        run([
            str(self.release / ".venv/bin/python"),
            "-m",
            "scripts.prepare_data_root",
            "--data-root",
            str(self.data_candidate),
            "--runtime-root",
            str(runtime_candidate),
        ], cwd=self.release, env=migration_env)

    def switch(self) -> None:
        self.update("switching", 78, "正在切换到新版本")
        shutil.copy2(UNIT_PATH, self.unit_backup)
        run(["systemctl", "stop", SERVICE])
        self.start_visual()

        RELEASES.mkdir(parents=True, exist_ok=True)
        if TARGET.is_symlink():
            self.previous_target = Path(os.readlink(TARGET))
            TARGET.unlink()
        elif TARGET.is_dir():
            self.previous_directory = RELEASES / f"{self.current_version}-pre-ota-{self.job_id}"
            TARGET.rename(self.previous_directory)
        else:
            raise UpdateFailed("production target is neither a release symlink nor a directory")
        os.symlink(self.release, TARGET, target_is_directory=True)
        self.program_swapped = True

        DATA_ROOT.rename(self.data_backup)
        self.data_candidate.rename(DATA_ROOT)
        self.data_swapped = True
        shutil.copy2(self.release / "systemd/matrix-screen-controller.service", UNIT_PATH)
        shutil.copy2(self.release / "systemd/matrix-screen-controller-ota.service", Path("/etc/systemd/system/matrix-screen-controller-ota.service"))
        run(["systemctl", "daemon-reload"])
        self.stop_visual()
        run(["systemctl", "start", SERVICE])

    def verify(self) -> None:
        self.update("verifying", 90, "正在验证新版本和真实屏幕驱动")
        deadline = time.monotonic() + 45
        last_error = "service did not respond"
        while time.monotonic() < deadline:
            try:
                output = subprocess.check_output(
                    ["curl", "--fail", "--silent", "http://127.0.0.1:8080/api/status"],
                    timeout=5,
                )
                status = json.loads(output.decode("utf-8"))
                screen = status.get("screen", {})
                driver = screen.get("driver_status") or {}
                if status.get("service", {}).get("software_version") != str(self.target_version):
                    raise UpdateFailed("new service reports the wrong software version")
                if screen.get("driver") != "walnutpi-h618-hub75":
                    raise UpdateFailed("new service did not start the production HUB75 driver")
                if screen.get("hardware_mapping") != "walnutpi-pi-bank-pwm-oe-v2":
                    raise UpdateFailed("new service reports the wrong hardware mapping")
                if driver.get("oe_timing_backend") != "h618-pwm4" or driver.get("oe_pulse_faults") != 0:
                    raise UpdateFailed("new service did not pass the PWM4 OE health check")
                return
            except (subprocess.SubprocessError, OSError, UnicodeError, json.JSONDecodeError) as exc:
                last_error = str(exc)
                time.sleep(1)
        raise UpdateFailed(f"new release health check timed out: {last_error}")

    def succeed(self) -> None:
        result = {
            "status": "success",
            "target_version": str(self.target_version),
            "packaged_at": str(self.request.get("packaged_at") or ""),
            "installed_at": utc_now(),
            "release_notes": str(self.request.get("release_notes") or ""),
            "error": None,
        }
        write_last_result(DATA_ROOT, result)
        self.job.update(
            phase="complete",
            percent=100,
            message="更新完成",
            finished_at=result["installed_at"],
            error=None,
        )
        write_json(self.status_path, {"schema_version": 1, "active": False, "job": self.job})
        self.package_path.unlink(missing_ok=True)
        self.request_path().unlink(missing_ok=True)
        shutil.rmtree(self.data_backup, ignore_errors=True)
        if self.previous_directory is not None:
            shutil.rmtree(self.previous_directory, ignore_errors=True)
        if self.previous_target is not None:
            previous_release = self.previous_target if self.previous_target.is_absolute() else TARGET.parent / self.previous_target
            if previous_release.parent == RELEASES:
                shutil.rmtree(previous_release, ignore_errors=True)
        shutil.rmtree(self.work_root, ignore_errors=True)
        prune_empty(self.package_path.parent)

    def rollback(self, error: BaseException) -> None:
        self.stop_visual()
        if self.program_swapped or self.data_swapped:
            subprocess.run(["systemctl", "stop", SERVICE], check=False)
        if self.data_swapped:
            failed_data = DATA_ROOT.with_name(f"matrix-screen-controller.failed.{self.job_id}")
            if DATA_ROOT.exists():
                DATA_ROOT.rename(failed_data)
            if self.data_backup.exists():
                self.data_backup.rename(DATA_ROOT)
            shutil.rmtree(failed_data, ignore_errors=True)
        if self.program_swapped:
            if TARGET.is_symlink():
                TARGET.unlink()
            if self.previous_directory is not None and self.previous_directory.exists():
                self.previous_directory.rename(TARGET)
            elif self.previous_target is not None:
                os.symlink(self.previous_target, TARGET, target_is_directory=True)
        if self.unit_backup.exists():
            shutil.copy2(self.unit_backup, UNIT_PATH)
            subprocess.run(["systemctl", "daemon-reload"], check=False)
        if self.program_swapped or self.data_swapped:
            subprocess.run(["systemctl", "start", SERVICE], check=False)
        shutil.rmtree(self.release, ignore_errors=True)
        shutil.rmtree(self.data_candidate, ignore_errors=True)
        result = {
            "status": "failed",
            "target_version": str(self.target_version),
            "packaged_at": str(self.request.get("packaged_at") or ""),
            "installed_at": utc_now(),
            "release_notes": str(self.request.get("release_notes") or ""),
            "error": str(error),
        }
        try:
            write_last_result(DATA_ROOT, result)
        except OSError:
            pass
        self.job.update(
            phase="failed",
            percent=0,
            message="更新失败，已恢复原版本",
            finished_at=result["installed_at"],
            error=str(error),
        )
        write_json(self.status_path, {"schema_version": 1, "active": False, "job": self.job})
        self.package_path.unlink(missing_ok=True)
        self.request_path().unlink(missing_ok=True)
        shutil.rmtree(self.work_root, ignore_errors=True)

    def start_visual(self) -> None:
        python = TARGET / ".venv/bin/python"
        self.visual = subprocess.Popen([
            str(python), "-m", "scripts.hub75_visual_hold",
            "--mode", "ota",
            "--brightness", "40",
            "--progress-file", str(self.status_path),
            "--orientation", str(int(self.request.get("orientation", 0))),
        ])
        time.sleep(1)
        if self.visual.poll() is not None:
            raise UpdateFailed("independent OTA display process failed to start")

    def stop_visual(self) -> None:
        if self.visual is None:
            return
        if self.visual.poll() is None:
            self.visual.send_signal(signal.SIGTERM)
            try:
                self.visual.wait(timeout=5)
            except subprocess.TimeoutExpired:
                self.visual.kill()
                self.visual.wait(timeout=2)
        self.visual = None

    def request_path(self) -> Path:
        return Path(str(self.request.get("request_path") or REQUEST_PATH))


def checked_token(value: Any, name: str) -> str:
    candidate = str(value or "")
    if not candidate or any(character not in "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789._-" for character in candidate):
        raise UpdateFailed(f"invalid {name}")
    return candidate


def prune_empty(path: Path) -> None:
    try:
        path.rmdir()
    except OSError:
        pass


def main() -> int:
    if os.geteuid() != 0 or os.uname().machine != "aarch64":
        print("OTA worker must run as root on the WalnutPi AArch64 host", file=sys.stderr)
        return 2
    request = read_json(REQUEST_PATH)
    if request is None or request.get("schema_version") != 1:
        print("OTA request is missing or invalid", file=sys.stderr)
        return 2
    transaction: Transaction | None = None
    try:
        transaction = Transaction(request)
        transaction.prepare()
        transaction.switch()
        transaction.verify()
        transaction.succeed()
        return 0
    except BaseException as exc:
        if transaction is not None:
            transaction.rollback(exc)
        print(f"OTA update failed: {exc}", file=sys.stderr)
        return 1


if __name__ == "__main__":
    raise SystemExit(main())