#!/usr/bin/env python3
"""OTA software update for ButtonTask.

The web side receives an uploaded package (tar.gz containing the ButtonTask
binary + webconfig/ + manifest.json), validates it, then hands the actual
swap/health-check/rollback to the privileged, detached helper `bt-update`.
Because the web service itself is part of the package being replaced, the
helper is launched detached (setsid) so it survives the service restart.
Status is reported via run/update-status.json.
"""
import hashlib
import json
import re
import subprocess
import time
from pathlib import Path
from typing import Optional

from flask import Blueprint, abort, jsonify, request

from auth import login_required
from paths import (UPDATE_STATUS, current_link, install_root, run_dir,
                   scripts_dir, versions_dir)

updater_bp = Blueprint("updater", __name__)

STAGING_NAME = "staging"


def _read_manifest(directory):
    mf = Path(directory) / "manifest.json"
    if not mf.exists():
        return {}
    try:
        with mf.open("r", encoding="utf-8") as f:
            return json.load(f)
    except (OSError, ValueError):
        return {}


def current_version():
    link = current_link()
    target = link.resolve() if link.exists() else None
    manifest = _read_manifest(target) if target else {}
    return {
        "version": manifest.get("version", "unknown"),
        "path": str(target) if target else "",
    }


def read_status():
    p = run_dir() / UPDATE_STATUS
    if not p.exists():
        return {"state": "idle"}
    try:
        with p.open("r", encoding="utf-8") as f:
            return json.load(f)
    except (OSError, ValueError):
        return {"state": "unknown"}


def _staging_dir():
    p = install_root() / STAGING_NAME
    p.mkdir(parents=True, exist_ok=True)
    return p


def _dir_size(path: Path) -> int:
    total = 0
    try:
        for f in path.rglob("*"):
            if f.is_file():
                total += f.stat().st_size
    except OSError:
        pass
    return total


def _previous_path() -> Optional[Path]:
    mark = install_root() / ".previous"
    if not mark.is_file():
        return None
    try:
        raw = mark.read_text(encoding="utf-8").strip()
        if not raw:
            return None
        return Path(raw).resolve()
    except OSError:
        return None


def list_versions():
    vd = versions_dir()
    current = current_link().resolve() if current_link().exists() else None
    previous = _previous_path()
    items = []
    if not vd.is_dir():
        return items
    for p in sorted(vd.iterdir(), key=lambda x: x.name, reverse=True):
        if not p.is_dir() or p.name.startswith("."):
            continue
        resolved = p.resolve()
        manifest = _read_manifest(p)
        items.append({
            "name": p.name,
            "version": manifest.get("version", p.name),
            "created": manifest.get("created", ""),
            "path": str(resolved),
            "current": current is not None and resolved == current,
            "previous": previous is not None and resolved == previous,
            "size_bytes": _dir_size(p),
        })
    return items


def _run_bt_update(*args):
    script = str(scripts_dir() / "bt-update")
    proc = subprocess.run(
        ["sudo", "-n", script, *args],
        capture_output=True,
        text=True,
        timeout=120,
        check=False,
    )
    if proc.returncode != 0:
        err = (proc.stderr or proc.stdout or "bt-update failed").strip()
        raise RuntimeError(err)
    return proc.stdout.strip()


@updater_bp.route("/api/update/versions", methods=["GET"])
@login_required
def api_update_versions():
    return jsonify({"versions": list_versions(), "current": current_version()})


@updater_bp.route("/api/update/versions/<name>", methods=["DELETE"])
@login_required
def api_update_version_delete(name):
    if not re.fullmatch(r"[A-Za-z0-9._-]+", name or ""):
        abort(400, "invalid version name")
    try:
        msg = _run_bt_update("remove", name)
    except RuntimeError as e:
        return jsonify({"ok": False, "error": str(e)}), 400
    return jsonify({"ok": True, "message": msg, "versions": list_versions()})


@updater_bp.route("/api/update/status", methods=["GET"])
@login_required
def api_update_status():
    return jsonify({"current": current_version(), "status": read_status()})


@updater_bp.route("/api/update", methods=["POST"])
@login_required
def api_update():
    if "file" not in request.files:
        abort(400, "package file required")
    fobj = request.files["file"]
    fname = Path(fobj.filename or "").name
    if not (fname.endswith(".tar.gz") or fname.endswith(".tgz")):
        abort(400, "package must be a .tar.gz")

    staging = _staging_dir()
    pkg_path = staging / f"upload-{int(time.time())}.tar.gz"
    fobj.save(str(pkg_path))

    # Optional integrity check against a client-provided sha256.
    expected = (request.form.get("sha256") or "").strip().lower()
    if expected:
        h = hashlib.sha256()
        with pkg_path.open("rb") as f:
            for chunk in iter(lambda: f.read(1 << 20), b""):
                h.update(chunk)
        if h.hexdigest() != expected:
            pkg_path.unlink(missing_ok=True)
            abort(400, "sha256 mismatch")

    script = str(scripts_dir() / "bt-update")
    # start_new_session detaches the child (setsid) so it survives the web
    # service restart performed mid-update.
    try:
        subprocess.Popen(
            ["sudo", "-n", script, "stage", str(pkg_path)],
            stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,
            start_new_session=True,
        )
    except FileNotFoundError as e:
        return jsonify({"ok": False, "error": str(e)}), 500
    return jsonify({"ok": True, "state": "started"})


@updater_bp.route("/api/update/rollback", methods=["POST"])
@login_required
def api_update_rollback():
    script = str(scripts_dir() / "bt-update")
    try:
        subprocess.Popen(
            ["sudo", "-n", script, "rollback"],
            stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL,
            start_new_session=True,
        )
    except FileNotFoundError as e:
        return jsonify({"ok": False, "error": str(e)}), 500
    return jsonify({"ok": True, "state": "rollback-started"})