#!/usr/bin/env python3
"""Keep a release install of the sign detector up to date: a friend's Pi, which gets the detector as signed release
bundles (release/make_release.py on the PC makes and publishes them), not as a git checkout.

A release install is a folder (~/sign-detector by default) holding:
  releases/<version>/  each release unpacked, as published (the code and models, with a RELEASE file)
  current              a link to the release that runs
  data/                what the releases share and keep across updates: captures/, map_cache/, uploads/ and
                       settings.json, linked into each release under those names, so the code finds them where a
                       checkout keeps them
  run.sh               what the desktop launcher and the login autostart run (release/run.sh)
  updater.json         where releases come from: {"server": <url>, "channel": "stable"}
  state.json           which release is current, the one before it, one downloaded and waiting, one on trial

How a release arrives: start.sh runs `updater.py loop` in the background, which every half hour, when there's
internet, reads the channel's latest.json, checks its signature against release/allowed_signers of the release that
runs, and if it names a newer version, downloads its bundle, checks its SHA-256, and unpacks it into releases/. It
doesn't switch to it while the detector runs (that would blank the screen for half a minute, maybe while driving):
run.sh does at the next start (`updater.py apply`), before start.sh. The new release is then on trial until the
detector has written to captures/health.csv a minute after it started (`confirm`); if instead it fails 3 times in a
row, start.sh exits with status 5 and run.sh goes back to the release before (`rollback`), which won't be tried again.

    python3 updater.py install --server URL [--channel stable] [--dir ~/sign-detector]   # the first time (install.sh)
    python3 updater.py status        # which release runs, which is waiting, and the channel's latest
    python3 updater.py check         # look for a newer release now, and download it
    python3 updater.py apply         # switch to a downloaded release (run.sh does, at each start)
    python3 updater.py rollback      # go back to the release before
    python3 updater.py loop          # check now and then, confirm a release on trial (start.sh starts it)

Only Python's standard library and ssh-keygen (for the signatures) are used: the Pi needs nothing installed for it.
"""
import argparse
import fcntl
import hashlib
import json
import os
import shutil
import socket
import subprocess
import sys
import tarfile
import tempfile
import time
import urllib.error
import urllib.request
from datetime import datetime
from pathlib import Path, PurePosixPath

HERE = Path(__file__).resolve().parent  # release/ in the release (or the checkout) this runs from
APP = HERE.parent
RELEASES, DATA, CURRENT, DOWNLOADS = "releases", "data", "current", "downloads"
STATE, CONFIG, LOCK, RELEASE_FILE = "state.json", "updater.json", "updater.lock", "RELEASE"
# Kept in data/ and linked into every release under these names: what the detector saves and what it keeps
LINKED_DIRS = ("captures", "map_cache", "uploads")
LINKED_FILES = ("settings.json",)
CHANNEL_FILE = "latest.json"  # in each channel's folder on the server, signed (latest.json.sig)
SIGNERS = "allowed_signers"  # in release/: the public keys a release must be signed with
IDENTITY, NAMESPACE = "release@sign-detector", "sign-detector-release"  # what make_release.py signs as
CHECK_EVERY_S = 30 * 60  # how often `loop` looks for a release
FIRST_CHECK_S = 60  # ...the first time, after start: the network may not be up at once
CONFIRM_AFTER_S = 90  # a release is good once the detector has written health.csv this long after the switch
TIMEOUT_S = 60  # for each request
KEEP = 3  # releases kept: the current one, the one before, and one downloaded
USER_AGENT = "sign-detector-updater"


class UpdateError(Exception):
    """A release couldn't be got or checked; nothing was changed."""


###############################################################################
# Versions, state and settings
###############################################################################
def version_key(version):
    """A version, e.g. "2026.11.08.2", as a tuple that compares in release order."""
    try:
        return tuple(int(part) for part in version.split("."))
    except (AttributeError, ValueError):
        raise UpdateError(f"{version!r} isn't a release version")


def base_of(app=APP):
    """The release install a release folder is in (the folder holding releases/), or None for a checkout."""
    return app.parent.parent if app.parent.name == RELEASES else None


def read_json(path, default=None):
    try:
        return json.loads(Path(path).read_text())
    except (OSError, ValueError):
        return default


def write_json(path, data):
    """Write JSON through a temporary file, so a power cut leaves the old file or the new one, never half of one."""
    path = Path(path)
    tmp = path.with_name(path.name + ".part")
    with open(tmp, "w") as f:
        f.write(json.dumps(data, indent=1) + "\n")
        f.flush()
        os.fsync(f.fileno())
    os.replace(tmp, path)


def load_state(base):
    state = read_json(base / STATE, {}) or {}
    for key in ("current", "previous", "staged", "trial"):
        state.setdefault(key, None)
    state.setdefault("bad", [])
    return state


def now_iso():
    return datetime.now().isoformat(timespec="seconds")


def say(*words):
    print(f"{datetime.now():%Y-%m-%d %H:%M:%S}  {' '.join(str(w) for w in words)}", flush=True)


class Locked:
    """One updater at a time in a release install: `loop` downloading while run.sh applies would race. With
    wait=False, entering raises BlockingIOError if another holds it."""

    def __init__(self, base, wait=True):
        self.path, self.wait, self.f = base / LOCK, wait, None

    def __enter__(self):
        self.f = open(self.path, "a")
        fcntl.flock(self.f, fcntl.LOCK_EX | (0 if self.wait else fcntl.LOCK_NB))
        return self

    def __exit__(self, *exc):
        self.f.close()


###############################################################################
# Getting and checking a release
###############################################################################
def fetch(url, out=None, token=None):
    """GET url: its bytes, or with out, written to that file. Raises UpdateError when it can't be got."""
    headers = {"User-Agent": USER_AGENT}
    if token:
        headers["Authorization"] = f"token {token}"
    try:
        with urllib.request.urlopen(urllib.request.Request(url, headers=headers), timeout=TIMEOUT_S) as response:
            if out is None:
                return response.read()
            with open(out, "wb") as f:
                shutil.copyfileobj(response, f, 1 << 20)
                f.flush()
                os.fsync(f.fileno())
            return None
    except (OSError, urllib.error.URLError) as e:  # HTTPError is a URLError; timeouts are OSErrors
        raise UpdateError(f"couldn't get {url}: {e}")


def verify(data, signature, signers):
    """Whether signature (an `ssh-keygen -Y sign` signature) is data's, by one of the keys in signers."""
    with tempfile.TemporaryDirectory() as tmp:
        sig = Path(tmp) / "data.sig"
        sig.write_bytes(signature)
        try:
            done = subprocess.run(["ssh-keygen", "-Y", "verify", "-f", str(signers), "-I", IDENTITY,
                                   "-n", NAMESPACE, "-s", str(sig)], input=data, capture_output=True, timeout=30)
        except (OSError, subprocess.TimeoutExpired):
            return False
    return done.returncode == 0


def latest(server, channel, signers):
    """The channel's latest release, as its signed manifest says: {"version", "file", "sha256", "size", ...}."""
    url = f"{server.rstrip('/')}/{channel}/{CHANNEL_FILE}"
    data, signature = fetch(url), fetch(url + ".sig")
    if not verify(data, signature, signers):
        raise UpdateError(f"{url} isn't signed by a key in {signers}: not used")
    try:
        manifest = json.loads(data)
        version_key(manifest["version"])
        if not (isinstance(manifest["file"], str) and isinstance(manifest["sha256"], str)
                and isinstance(manifest["size"], int)):
            raise ValueError
    except (ValueError, KeyError, TypeError):
        raise UpdateError(f"{url} isn't a release manifest")
    if PurePosixPath(manifest["file"]).is_absolute() or ".." in PurePosixPath(manifest["file"]).parts:
        raise UpdateError(f"{url} names a file outside the server's folder: {manifest['file']}")
    return manifest


def sha256(path):
    h = hashlib.sha256()
    with open(path, "rb") as f:
        for chunk in iter(lambda: f.read(1 << 20), b""):
            h.update(chunk)
    return h.hexdigest()


def download(server, manifest, folder):
    """Download the manifest's bundle into folder and check its size and SHA-256. Returns its path."""
    folder.mkdir(parents=True, exist_ok=True)
    path = folder / PurePosixPath(manifest["file"]).name
    part = path.with_name(path.name + ".part")
    fetch(f"{server.rstrip('/')}/{manifest['file']}", out=part)
    if part.stat().st_size != manifest["size"] or sha256(part) != manifest["sha256"]:
        part.unlink(missing_ok=True)
        raise UpdateError(f"{manifest['file']} isn't the file its manifest describes (size or SHA-256): not used")
    os.replace(part, path)
    return path


def unpack(bundle, dest, version):
    """Unpack a release bundle into dest (which mustn't exist): only folders and plain files, all under the bundle's
    one top folder, sign-detector-<version>/ (a bundle that tries to write elsewhere is refused). Keeps each file's
    time, which says how fresh a map tile is, and whether it's executable."""
    top = f"sign-detector-{version}"
    part = dest.with_name(dest.name + ".part")
    if part.exists():
        shutil.rmtree(part)
    try:
        with tarfile.open(bundle, "r:gz") as tar:
            members = tar.getmembers()
            for m in members:
                name = PurePosixPath(m.name)
                if (name.is_absolute() or ".." in name.parts or not name.parts or name.parts[0] != top
                        or not (m.isfile() or m.isdir())):
                    raise UpdateError(f"{bundle.name} holds {m.name!r}, which a release can't: not used")
            for m in members:
                target = part.joinpath(*PurePosixPath(m.name).parts[1:])
                if m.isdir():
                    target.mkdir(parents=True, exist_ok=True)
                    continue
                target.parent.mkdir(parents=True, exist_ok=True)
                with tar.extractfile(m) as src, open(target, "wb") as out:
                    shutil.copyfileobj(src, out, 1 << 20)
                os.chmod(target, 0o755 if m.mode & 0o111 else 0o644)
                os.utime(target, (m.mtime, m.mtime))
    except (tarfile.TarError, OSError, EOFError) as e:
        shutil.rmtree(part, ignore_errors=True)
        raise UpdateError(f"{bundle.name} couldn't be unpacked: {e}")
    except UpdateError:
        shutil.rmtree(part, ignore_errors=True)
        raise
    info = read_json(part / RELEASE_FILE, {}) or {}
    if info.get("version") != version:
        shutil.rmtree(part, ignore_errors=True)
        raise UpdateError(f"{bundle.name} says it's release {info.get('version')!r}, not {version}: not used")
    return part


def link_data(base, folder):
    """Point the release in folder at data/: its captures/, map_cache/, uploads/ and settings.json become links to
    data/'s. The map tiles it comes with are added to data/map_cache/ first, those it doesn't have yet only: the
    ones there were fetched on this Pi, and may be fresher."""
    data = base / DATA
    for name in LINKED_DIRS:
        (data / name).mkdir(parents=True, exist_ok=True)
    tiles = folder / "map_cache"
    if tiles.is_dir() and not tiles.is_symlink():
        for tile in tiles.iterdir():
            kept = data / "map_cache" / tile.name
            if tile.is_file() and not kept.exists():
                shutil.copy2(tile, kept)
    for name in LINKED_DIRS + LINKED_FILES:
        path = folder / name
        if path.is_symlink() or path.is_file():
            path.unlink()
        elif path.is_dir():
            shutil.rmtree(path)
        # Relative, so the install can be moved: from releases/<version>/ up two to the install
        path.symlink_to(Path("..") / ".." / DATA / name)


def set_link(base, version):
    """Point current at releases/<version>, in one step (a new link put in place of the old)."""
    tmp = base / (CURRENT + ".new")
    if tmp.is_symlink() or tmp.exists():
        tmp.unlink()
    tmp.symlink_to(Path(RELEASES) / version)
    os.replace(tmp, base / CURRENT)


def refresh_launcher(base, version):
    """Copy the release's run.sh into the install, as a new file in place of the old: bash reads a script as it
    runs it, so the run.sh running now goes on reading the old one."""
    src = base / RELEASES / version / "release" / "run.sh"
    if src.is_file():
        tmp = base / "run.sh.new"
        shutil.copy2(src, tmp)
        os.chmod(tmp, 0o755)
        os.replace(tmp, base / "run.sh")


def prune(base, state):
    """Delete releases and downloads that aren't current, the one before or the one waiting."""
    keep = {state["current"], state["previous"], state["staged"]} - {None}
    for path in sorted((base / RELEASES).iterdir()) if (base / RELEASES).is_dir() else []:
        if path.name not in keep:
            shutil.rmtree(path, ignore_errors=True)
    for path in sorted((base / DOWNLOADS).iterdir()) if (base / DOWNLOADS).is_dir() else []:
        path.unlink(missing_ok=True)


###############################################################################
# The commands
###############################################################################
def settings(base):
    config = read_json(base / CONFIG, {}) or {}
    if not config.get("server"):
        raise UpdateError(f"{base / CONFIG} doesn't say where releases come from (\"server\")")
    return config["server"], config.get("channel") or "stable"


def signers_of(base, state):
    """The keys a release must be signed with: those of the release that runs."""
    return base / RELEASES / state["current"] / "release" / SIGNERS


def check(base):
    """Look for a newer release on the channel and download and unpack it, to be switched to at the next start.
    Returns the version made ready, or None."""
    server, channel = settings(base)
    state = load_state(base)
    manifest = latest(server, channel, signers_of(base, state))
    version = manifest["version"]
    if version_key(version) <= version_key(state["current"]):
        return None
    if version in state["bad"]:
        say(f"Release {version} is the newest, but it failed here before: staying on {state['current']}.")
        return None
    if state["staged"] == version and (base / RELEASES / version).is_dir():
        return None  # downloaded already, waiting for the next start
    say(f"Release {version} is out (this Pi has {state['current']}): downloading it.")
    bundle = download(server, manifest, base / DOWNLOADS)
    folder = base / RELEASES / version
    if folder.exists():
        shutil.rmtree(folder)
    part = unpack(bundle, folder, version)
    link_data(base, part)
    os.replace(part, folder)
    bundle.unlink(missing_ok=True)
    state = load_state(base)
    state["staged"] = version
    write_json(base / STATE, state)
    prune(base, state)
    say(f"Release {version} is ready: the detector switches to it at its next start.")
    return version


def apply(base):
    """Switch to the release downloaded since the last start, if any, putting it on trial. Returns its version."""
    state = load_state(base)
    version = state["staged"]
    if not version or version == state["current"] or not (base / RELEASES / version / RELEASE_FILE).is_file():
        return None
    state.update(previous=state["current"], current=version, staged=None,
                 trial={"version": version, "since": time.time()})
    set_link(base, version)
    write_json(base / STATE, state)
    refresh_launcher(base, version)
    prune(base, state)
    say(f"Switched to release {version} (from {state['previous']}); on trial until it has run a minute.")
    return version


def confirm(base):
    """End the trial of the current release once the detector has run on it: health.csv written a while after the
    switch. Returns whether it was confirmed now."""
    state = load_state(base)
    trial = state["trial"]
    if not trial or trial.get("version") != state["current"]:
        return False
    try:
        written = (base / DATA / "captures" / "health.csv").stat().st_mtime
    except OSError:
        return False
    if written < trial["since"] + CONFIRM_AFTER_S:
        return False
    state["trial"] = None
    write_json(base / STATE, state)
    say(f"Release {state['current']} runs: kept.")
    return True


def on_trial(base):
    state = load_state(base)
    return bool(state["trial"]) and state["trial"].get("version") == state["current"]


def rollback(base):
    """Go back to the release before, marking the current one bad so it isn't tried again. Returns the version gone
    back to, or None when there's none."""
    state = load_state(base)
    before = state["previous"]
    if not before or not (base / RELEASES / before / RELEASE_FILE).is_file():
        return None
    failed = state["current"]
    state.update(current=before, previous=None, staged=None, trial=None,
                 bad=sorted(set(state["bad"]) | {failed}, key=version_key))
    set_link(base, before)
    write_json(base / STATE, state)
    refresh_launcher(base, before)
    prune(base, state)
    say(f"Went back to release {before}: {failed} didn't run, and won't be tried again.")
    return before


def online(server, timeout=5):
    """Whether the server's host can be reached at all (a cheap test before a check, as most drives have no
    internet)."""
    from urllib.parse import urlsplit
    parts = urlsplit(server)
    if parts.scheme == "file":
        return True
    try:
        socket.create_connection((parts.hostname, parts.port or (443 if parts.scheme == "https" else 80)),
                                 timeout=timeout).close()
        return True
    except OSError:
        return False


def loop(base, every=CHECK_EVERY_S, first=FIRST_CHECK_S):
    """Now and then: confirm a release on trial, and look for a newer one when there's internet."""
    next_check = time.monotonic() + first
    while True:
        try:
            with Locked(base):
                confirm(base)
                if time.monotonic() >= next_check:
                    next_check = time.monotonic() + every
                    server, _ = settings(base)
                    if online(server):
                        check(base)
        except UpdateError as e:
            say(f"Couldn't update: {e}")
        except Exception as e:  # never let the loop die: it's the only way updates arrive
            say(f"Updater error: {e!r}")
        time.sleep(60)


def install(base, server, channel, signers=None):
    """Set up a release install in base with the channel's latest release (install.sh runs this). signers: the keys
    to check that first release with (later ones are checked with the keys of the release that runs)."""
    base.mkdir(parents=True, exist_ok=True)
    if (base / ".git").exists() or ((base / "captures").is_dir() and not (base / RELEASES).is_dir()):
        raise UpdateError(f"{base} is a checkout or a copy of a Pi's captures, not a release install: "
                          "choose another folder with --dir")
    signers = Path(signers) if signers else HERE / SIGNERS
    manifest = latest(server, channel, signers)
    version = manifest["version"]
    with Locked(base):
        say(f"Installing release {version} into {base}.")
        bundle = download(server, manifest, base / DOWNLOADS)
        folder = base / RELEASES / version
        if folder.exists():
            shutil.rmtree(folder)
        part = unpack(bundle, folder, version)
        link_data(base, part)
        os.replace(part, folder)
        bundle.unlink(missing_ok=True)
        state = load_state(base)
        state.update(previous=state["current"] if state["current"] != version else state["previous"],
                     current=version, staged=None, trial=None)
        write_json(base / CONFIG, {"server": server, "channel": channel})
        set_link(base, version)
        write_json(base / STATE, state)
        refresh_launcher(base, version)
        prune(base, state)
    say(f"Release {version} is installed. Start it with {base / 'run.sh'}.")
    return version


def status(base):
    state = load_state(base)
    lines = [f"Release install: {base}", f"Running: {state['current']}"
             + (" (on trial)" if on_trial(base) else ""),
             f"Before it: {state['previous'] or '-'}", f"Downloaded, for the next start: {state['staged'] or '-'}"]
    if state["bad"]:
        lines.append(f"Failed here, not tried again: {', '.join(state['bad'])}")
    try:
        server, channel = settings(base)
        lines.append(f"Channel {channel} at {server}: latest "
                     f"{latest(server, channel, signers_of(base, state))['version']}")
    except UpdateError as e:
        lines.append(f"The channel couldn't be read: {e}")
    print("\n".join(lines))


def main(argv=None):
    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    p.add_argument("command", choices=["install", "status", "check", "apply", "confirm", "rollback", "on-trial",
                                       "loop"])
    p.add_argument("--dir", type=Path, help="the release install (default: the one this runs from; for install, "
                                            "~/sign-detector)")
    p.add_argument("--server", help="install: where releases are published (make_release.py prints it)")
    p.add_argument("--channel", default="stable", help="install: which channel to follow (default stable)")
    p.add_argument("--signers", type=Path, help="install: the keys to check the first release with")
    args = p.parse_args(argv)

    base = args.dir.expanduser().resolve() if args.dir else base_of()
    if args.command == "install":
        if not args.server:
            p.error("install needs --server")
        try:
            install(base or Path.home() / "sign-detector", args.server, args.channel, args.signers)
        except UpdateError as e:
            sys.exit(f"Not installed: {e}")
        return 0
    if base is None:
        p.error(f"{APP} isn't in a release install (releases/<version>/): give --dir")

    if args.command == "on-trial":
        return 0 if on_trial(base) else 1
    if args.command == "loop":
        loop(base)
        return 0
    try:
        if args.command == "status":
            status(base)
            return 0
        if args.command == "confirm":
            return 0 if confirm(base) else 1
        if args.command == "rollback":
            with Locked(base):
                return 0 if rollback(base) else 1
        if args.command == "apply":
            try:
                with Locked(base, wait=False):
                    apply(base)
            except BlockingIOError:
                say("The updater is busy downloading: switching at the next start instead.")
            return 0
        with Locked(base):
            if check(base) is None:
                print(f"Up to date: {load_state(base)['current']}.")
        return 0
    except UpdateError as e:
        say(f"Couldn't update: {e}")
        return 1


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