shaiwatcher/wrapper/wrapper.py

import os, sys, time, shutil, subprocess, json, pathlib, re, datetime
from typing import Tuple

REPO_URL    = os.getenv("REPO_URL", "https://git.rolfsvaag.no/frarol96/shaiwatcher.git").strip()
REPO_BRANCH = os.getenv("REPO_BRANCH", "main").strip()
REPO_TOKEN  = os.getenv("REPO_AHTOKEN", os.getenv("REPO_TOKEN", "")).strip()

CHECK_TIME_UTC = os.getenv("CHECK_TIME_UTC", "03:00").strip()
IGNORE_TEST_LEVEL = os.getenv("IGNORE_TEST_LEVEL", "1").strip() == "1"
PIP_INSTALL = os.getenv("PIP_INSTALL_REQUIREMENTS", "1").strip() == "1"

CACHE_DIR = pathlib.Path("/cache/app")
TMP_DIR   = pathlib.Path("/cache/tmp")
PREV_DIR  = pathlib.Path("/cache/prev")
DATA_DIR  = pathlib.Path("/data")

RUN_TIMEOUT = int(os.getenv("WRAPPER_STOP_TIMEOUT", "25"))
ROLLBACK_MAX_FAILS = 3

def log(m): print(f"[wrapper] {m}", flush=True)
def run(*cmd, cwd=None, check=True): 
    log(f"$ {' '.join(cmd)}")
    return subprocess.run(cmd, cwd=cwd, check=check, text=True, capture_output=True)

def ensure_git():
    try: run("git","--version")
    except subprocess.CalledProcessError as e:
        log(f"git missing? {e.stderr}"); sys.exit(1)

def utc_now(): return datetime.datetime.utcnow()
def next_utc(hhmm: str) -> float:
    hh, mm = map(int, hhmm.split(":"))
    now = utc_now()
    tgt = now.replace(hour=hh, minute=mm, second=0, microsecond=0)
    if tgt <= now: tgt += datetime.timedelta(days=1)
    return (tgt - now).total_seconds()

_VER_RE = re.compile(r'^\s*VERSION\s*=\s*[\'"]([^\'"]+)[\'"]', re.M)
def extract_version_from(path: pathlib.Path) -> str:
    try:
        txt = path.read_text(encoding="utf-8", errors="ignore")
        m = _VER_RE.search(txt)
        return m.group(1).strip() if m else "v0.0.0.0"
    except Exception:
        return "v0.0.0.0"

def parse_version(ver: str):
    test = "-T" in ver
    core = ver.split("-T")[0].lstrip("v")
    parts = [int(p or 0) for p in core.split(".")+["0","0","0","0"]][:4]
    return (*parts, test)

def should_update(old, new):
    oM,oE,oN,oP,ot = parse_version(old)
    nM,nE,nN,nP,nt = parse_version(new)
    if (oM,oE,oN,oP) != (nM,nE,nN,nP):
        return True
    return not IGNORE_TEST_LEVEL and (ot != nt)

def clone_or_fetch(target: pathlib.Path):
    if target.exists() and (target/".git").exists():
        try:
            run("git","fetch","--all","-p", cwd=target)
            run("git","reset","--hard", f"origin/{REPO_BRANCH}", cwd=target)
            return
        except Exception as e:
            log(f"fetch failed, recloning: {e}")
            shutil.rmtree(target, ignore_errors=True)

    url = REPO_URL
    if REPO_TOKEN and REPO_URL.startswith("https://"):
        url = REPO_URL.replace("https://", f"https://{REPO_TOKEN}@")
    run("git","clone","--branch", REPO_BRANCH, "--depth","1", url, str(target))

def prime_tmp_then_decide():
    TMP_DIR.mkdir(parents=True, exist_ok=True)
    shutil.rmtree(TMP_DIR, ignore_errors=True)
    clone_or_fetch(TMP_DIR)

    new_ver = extract_version_from(TMP_DIR / "bot.py")
    old_ver = extract_version_from(CACHE_DIR / "bot.py") if (CACHE_DIR / "bot.py").exists() else "v0.0.0.0"
    do_swap = not (CACHE_DIR / "bot.py").exists() or should_update(old_ver, new_ver)
    why = "first fetch" if not (CACHE_DIR / "bot.py").exists() else f"update allowed: {do_swap} (old={old_ver}, new={new_ver})"
    return do_swap, old_ver, new_ver, why

def swap_cache_to_prev():
    PREV_DIR.mkdir(parents=True, exist_ok=True)
    shutil.rmtree(PREV_DIR, ignore_errors=True)
    if CACHE_DIR.exists():
        shutil.move(str(CACHE_DIR), str(PREV_DIR))

def copy_tmp_to_cache():
    shutil.rmtree(CACHE_DIR, ignore_errors=True)
    shutil.copytree(TMP_DIR, CACHE_DIR, dirs_exist_ok=False)

def pip_install(cwd: pathlib.Path):
    req = cwd / "requirements.txt"
    if PIP_INSTALL and req.exists():
        try: run(sys.executable, "-m", "pip", "install", "-r", str(req))
        except subprocess.CalledProcessError as e:
            log("pip install failed (continuing)"); log(e.stdout + "\n" + e.stderr)

def set_boot_env(status: str, old_ver: str, new_ver: str, commit: str = "", subject: str = ""):
    os.environ["BOOT_STATUS"]   = status
    os.environ["BOOT_OLDVER"]   = old_ver
    os.environ["BOOT_NEWVER"]   = new_ver
    os.environ["BUILD_COMMIT"]  = commit
    os.environ["BUILD_SUBJECT"] = subject

def get_head_info(cwd: pathlib.Path):
    try:
        sha = run("git","rev-parse","--short","HEAD", cwd=cwd).stdout.strip()
        subj = run("git","log","-1","--pretty=%s", cwd=cwd).stdout.strip()
        return (sha, subj)
    except Exception:
        return ("","")

def start_bot(cwd: pathlib.Path) -> subprocess.Popen:
    env = os.environ.copy()
    env.setdefault("DATA_FILE", "/data/data.json")  # new name
    return subprocess.Popen([sys.executable, "-u", "bot.py"], cwd=cwd, env=env)

def graceful_stop(proc: subprocess.Popen) -> bool:
    try:
        proc.terminate()
        try:
            proc.wait(timeout=RUN_TIMEOUT); return True
        except subprocess.TimeoutExpired:
            proc.kill(); proc.wait(timeout=10); return True
    except Exception:
        return False

def run_loop():
    ensure_git()
    updated, old_ver, new_ver, why = prime_tmp_then_decide()
    if updated: log(f"updating cache: {why}"); swap_cache_to_prev(); copy_tmp_to_cache()
    else: log(f"no update: {why}")

    pip_install(CACHE_DIR)
    sha, subj = get_head_info(CACHE_DIR)
    set_boot_env(
        "Successfully fetched, cached, and booted new version" if updated else "Successfully booted from cached version",
        old_ver, new_ver, sha, subj
    )
    proc = start_bot(CACHE_DIR)
    log(f"bot started pid={proc.pid}")

    consecutive_failures = 0
    while True:
        delay = next_utc(CHECK_TIME_UTC)
        log(f"sleeping {int(delay)}s until {CHECK_TIME_UTC} UTC for update check")
        time.sleep(delay)

        try:
            upd, cur_ver, remote_ver, why = prime_tmp_then_decide()
            log(f"nightly check: {why}")
            if not upd:
                continue

            log("updating to new version at nightly window")
            graceful_stop(proc)
            swap_cache_to_prev()
            copy_tmp_to_cache()
            pip_install(CACHE_DIR)
            sha, subj = get_head_info(CACHE_DIR)
            set_boot_env("Successfully fetched, cached, and booted new version", cur_ver, remote_ver, sha, subj)
            proc = start_bot(CACHE_DIR)
            log(f"bot restarted on new version pid={proc.pid}")
            consecutive_failures = 0

        except Exception as e:
            log(f"nightly update failed: {e}")
            consecutive_failures += 1
            if consecutive_failures < ROLLBACK_MAX_FAILS and PREV_DIR.exists() and (PREV_DIR / "bot.py").exists():
                log("attempting rollback to previous cached version")
                if proc.poll() is None: graceful_stop(proc)
                shutil.rmtree(CACHE_DIR, ignore_errors=True)
                shutil.copytree(PREV_DIR, CACHE_DIR, dirs_exist_ok=False)
                pip_install(CACHE_DIR)
                set_boot_env("Rolled back to last known working version", "-", extract_version_from(CACHE_DIR / "bot.py"))
                proc = start_bot(CACHE_DIR)
            elif consecutive_failures >= ROLLBACK_MAX_FAILS:
                log("critical: failed 3 times to update/restart; entering freeze mode")
                try:
                    if proc.poll() is None: graceful_stop(proc)
                except Exception: pass
                while True: time.sleep(3600)

if __name__ == "__main__":
    run_loop()