"""One authorized Part A execution, end to end, in ONE process, with a bounded overall deadline (v2.2.16):

  VM start -> package in Blob (hash read back) -> probe session on the VM (action Run Command) -> evidence download
  and manifest check -> Log Analytics events -> VM deallocated in `finally`, retried, and CONFIRMED.

  python -I probe_orchestrator.py --version 2.2.16 --prior-runs 2 --runs-dir "<contest>\\runs" --release-dir "<...>\\release"

Time limits (all monotonic, none of them open-ended):
  --session-deadline (default 3600 s): everything before the deallocation. VM start and every poll are cut at this
    deadline; when it passes, the orchestrator stops waiting and goes straight to the deallocation. A probe still
    running on the VM is ended by the deallocation; its ledger reservation stays counted in full (see the package).
  Deallocation: up to DEALLOC_ATTEMPTS attempts, each polled for at most DEALLOC_POLL_S, then the power state is read
    until it shows "VM deallocated" for at most CONFIRM_S. Worst case after the session deadline: about 45 minutes.
  Every HTTP call has a 120 s timeout and at most four tries (azlite.http).

Exit codes: 0 = session finished and the VM is confirmed deallocated; 2 = the session failed or hit its deadline, VM
confirmed deallocated; 3 = the VM could NOT be confirmed deallocated (printed loudly with the manual command).

No credential passes through this process: the VM reads the Oxylabs API user from Key Vault with its managed identity,
and stops if it cannot (there is no other route).
"""
import argparse
import hashlib
import importlib.util
import json
import os
import re
import sys
import threading
import time
import zipfile

HERE = os.path.dirname(os.path.abspath(__file__))
spec = importlib.util.spec_from_file_location("azlite", os.path.join(HERE, "azlite.py"))
az = importlib.util.module_from_spec(spec)
spec.loader.exec_module(az)
TESTS = {"2.2.15": "86", "2.2.16": "96"}

DEFAULT_SESSION_DEADLINE_S = 3600
VM_START_POLL_S = 900
POLL_EVERY_S = 15
DEALLOC_ATTEMPTS = 3
DEALLOC_POLL_S = 600
CONFIRM_S = 300

clock = time.monotonic       # replaced in tests
sleep = time.sleep


class DeadlineReached(Exception):
    pass


def now():
    return time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())


def say(m):
    print(f"{now()} {m}", flush=True)


def vm_state():
    s, d, _ = az.arm("GET", az.vmpath() + "/instanceView")
    return [x.get("displayStatus") for x in d.get("statuses", [])] if s == 200 else [f"HTTP {s}"]


def poll_operation(op, until):
    """Poll an Azure async operation until it ends or `until` (monotonic) passes. Never unbounded."""
    while clock() < until:
        sleep(min(10, max(0.0, until - clock())))
        try:
            _, d, _ = az.http("GET", op, headers={"Authorization": f"Bearer {az.token(az.ARM)}"})
        except Exception as exc:  # noqa: BLE001 - transient; the deadline still bounds the loop
            say(f"  poll error: {str(exc)[:120]}")
            continue
        if d.get("status") in ("Succeeded", "Failed", "Canceled"):
            return d.get("status")
    return "TimedOut"


def vm_action(action, until):
    s, d, h = az.arm("POST", az.vmpath() + "/" + action)
    if s >= 300:
        return f"HTTP {s}"
    op = h.get("Azure-AsyncOperation") or h.get("azure-asyncoperation")
    return poll_operation(op, until) if op else f"HTTP {s}"


def deallocate_and_confirm():
    """Deallocate with retries and confirm the power state. Returns True only when 'VM deallocated' was read."""
    for attempt in range(1, DEALLOC_ATTEMPTS + 1):
        try:
            say(f"VM deallocate (attempt {attempt}/{DEALLOC_ATTEMPTS}): {vm_action('deallocate', clock() + DEALLOC_POLL_S)}")
        except Exception as exc:  # noqa: BLE001
            say(f"VM deallocate attempt {attempt} raised: {str(exc)[:200]}")
        end = clock() + CONFIRM_S
        while True:
            try:
                state = vm_state()
                say(f"VM state: {state}")
                if "VM deallocated" in state:
                    return True
            except Exception as exc:  # noqa: BLE001
                say(f"VM state read failed: {str(exc)[:200]}")
            if clock() >= end:
                break
            sleep(POLL_EVERY_S)
    return False


def storage_headers(extra=None):
    h = {"Authorization": f"Bearer {az.token('https://storage.azure.com')}", "x-ms-version": "2021-08-06"}
    h.update(extra or {})
    return h


def blob_get(name):
    s, d, _ = az.http("GET", az.blob_url(name), headers=storage_headers(), raw=True)
    return d if s == 200 else None


def blob_list(prefix):
    s, d, _ = az.http("GET", f"https://{az.ACCOUNT}.blob.core.windows.net/{az.CONTAINER}?restype=container&comp=list&prefix={az.urllib.parse.quote(prefix)}",
                      headers=storage_headers(), raw=True)
    return re.findall(r"<Name>(.*?)</Name>", d.decode()) if s == 200 else []


def run_command(script_text, until):
    """Submit the session script (action Run Command) and poll it until it ends or `until` passes."""
    s, d, h = az.arm("POST", az.vmpath() + "/runCommand", {"commandId": "RunPowerShellScript", "script": script_text.splitlines()})
    if s >= 300:
        raise RuntimeError(f"run command not accepted: HTTP {s} {str(d)[:300]}")
    op = h.get("Azure-AsyncOperation") or h.get("azure-asyncoperation")
    say(f"run command submitted (HTTP {s})")
    while clock() < until:
        sleep(min(POLL_EVERY_S, max(0.0, until - clock())))
        try:
            _, d2, _ = az.http("GET", op, headers={"Authorization": f"Bearer {az.token(az.ARM)}"})
        except Exception as exc:  # noqa: BLE001
            say(f"  poll error: {str(exc)[:120]}")
            continue
        if d2.get("status") in ("Succeeded", "Failed", "Canceled"):
            out = "\n".join(v.get("message", "") for v in ((d2.get("properties") or {}).get("output") or {}).get("value", []))
            return d2.get("status"), out
    raise DeadlineReached("session deadline reached while the probe session was still running on the VM")


def fill_template(version, pkg, prior_runs, max_requests):
    sha = hashlib.sha256(open(pkg, "rb").read()).hexdigest()
    launcher_sha = hashlib.sha256(zipfile.ZipFile(pkg).read("scraper/infra/run_probe.ps1")).hexdigest().upper()
    script = open(os.path.join(HERE, "vm_probe_session.tpl.ps1"), encoding="utf-8").read()
    script = (script.replace("__LAUNCHER_SHA__", launcher_sha).replace("__VERSION__", version).replace("__PKG_SHA__", sha.upper())
              .replace("__TESTS__", TESTS[version]).replace("__PRIOR_RUNS__", str(prior_runs))
              .replace("__MAX_REQUESTS__", str(max_requests)))
    if re.findall(r"__[A-Z_]+__", script):
        raise SystemExit("template not fully filled")
    return script, sha


def session(a, script, sha, until):
    """Everything before the deallocation. Raises DeadlineReached when `until` passes."""
    say(f"VM before: {vm_state()}")
    status = vm_action("start", min(until, clock() + VM_START_POLL_S))
    say(f"VM start: {status}")
    if status != "Succeeded":
        raise RuntimeError(f"VM start {status}")
    started = now()
    name = f"deploy/marianos-scraper-{a.version}.zip"
    remote = blob_get(name)
    if remote is None or hashlib.sha256(remote).hexdigest() != sha:
        az.http("PUT", az.blob_url(name), open(a.pkg, "rb").read(), storage_headers({"x-ms-blob-type": "BlockBlob"}))
        remote = blob_get(name)
    if remote is None or hashlib.sha256(remote).hexdigest() != sha:
        raise RuntimeError("package in Blob does not match by read-back")
    say(f"package in Blob verified by read-back: {name}")
    status, output = run_command(script, until)
    say(f"run command {status}")
    print(output)
    m = re.search(r"evidence\\(oxylabs-probe-\d{8}T\d{6}Z)", output)
    if m:
        download_evidence(a, m.group(1))
    events(started, until)
    return status == "Succeeded"


def download_evidence(a, folder):
    local = os.path.join(a.runs_dir, f"azure_probe_{folder[len('oxylabs-probe-'):len('oxylabs-probe-') + 8]}")
    for n in blob_list(f"probe/{folder}"):
        data = blob_get(n)
        rel = n[len("probe/"):].replace(f"{folder}-session.log", f"{folder}/session.log")
        path = os.path.join(local, *rel.split("/")[1:]) if rel.startswith(folder + "/") else os.path.join(local, rel)
        os.makedirs(os.path.dirname(path), exist_ok=True)
        open(path, "wb").write(data)
    bad = []
    man = os.path.join(local, "manifest.sha256")
    if os.path.exists(man):
        for line in open(man, encoding="utf-8").read().splitlines():
            if line.strip():
                h, f = line.split("  ", 1)
                p = os.path.join(local, *f.split("/"))
                if not os.path.exists(p) or hashlib.sha256(open(p, "rb").read()).hexdigest() != h:
                    bad.append(f)
    say(f"evidence downloaded to {local}; manifest check: {'OK' if not bad else bad}")
    led = blob_get("probe/cumulative-usage.json")          # read only: the probe itself keeps the ledger
    if led:
        open(os.path.join(local, "cumulative-usage.blob.json"), "wb").write(led)
        say(f"cumulative usage (Blob): {json.loads(led).get('total_jobs_counted')} of {json.loads(led).get('ceiling')} jobs")


def events(started, until):
    ws = "3068ae2d-60a8-4bdd-b4f5-3396f9e9d4d8"
    q = (f'Event | where Source == "ShifullahScraper" | where TimeGenerated > datetime({started}) '
         '| project TimeGenerated, Computer, EventID, RenderedDescription | order by TimeGenerated asc')
    rows = []
    for i in range(12):                                     # ingestion can lag a few minutes; bounded by `until`
        s, d, _ = az.http("POST", f"https://api.loganalytics.io/v1/workspaces/{ws}/query", {"query": q},
                          {"Authorization": f"Bearer {az.token('https://api.loganalytics.io')}"})
        rows = d["tables"][0]["rows"] if s == 200 else []
        if len(rows) >= 3 or clock() + 30 >= until:
            break
        sleep(30)
    say(f"Log Analytics events: {len(rows)}")
    for r in rows:
        print("  ", r[0], r[2], str(r[3])[:160])


def main(argv=None):
    ap = argparse.ArgumentParser()
    ap.add_argument("--version", required=True, choices=sorted(TESTS))
    ap.add_argument("--prior-runs", type=int, required=True)
    ap.add_argument("--max-requests", type=int, required=True, help="cap asked for; the probe lowers it to what the ledger leaves")
    ap.add_argument("--runs-dir", required=True)
    ap.add_argument("--release-dir", required=True)
    ap.add_argument("--session-deadline", type=int, default=DEFAULT_SESSION_DEADLINE_S,
                    help="seconds for everything before the deallocation (default 3600)")
    a = ap.parse_args(argv)
    a.pkg = os.path.join(a.release_dir, f"marianos-scraper-{a.version}.zip")
    script, sha = fill_template(a.version, a.pkg, a.prior_runs, a.max_requests)
    say(f"package {a.version} sha256 {sha}; session deadline {a.session_deadline} s")
    until = clock() + a.session_deadline
    ok = False
    try:
        ok = session(a, script, sha, until)
    except DeadlineReached as exc:
        say(f"STOP: {exc}")
    except Exception as exc:  # noqa: BLE001 - reported, then the VM is deallocated
        say(f"STOP: session failed: {type(exc).__name__}: {str(exc)[:300]}")
    finally:
        confirmed = deallocate_and_confirm()
        if confirmed:
            say("VM deallocation CONFIRMED")
        else:
            say("!!! VM NOT CONFIRMED DEALLOCATED. Deallocate it by hand now: "
                f"az vm deallocate -g {az.RG} -n {az.VM}  (then: az vm get-instance-view -g {az.RG} -n {az.VM})")
    code = 3 if not confirmed else (0 if ok else 2)
    say(f"orchestrator exit {code}")
    return code


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