#!/usr/bin/env python3
"""Verify an immutable bundle or launch its fixed-duration Linux diagnostic."""

import argparse
import json
from pathlib import Path
import re
import stat
import subprocess
import sys
import time
import uuid
import xml.etree.ElementTree as ET

sys.dont_write_bytecode = True

from analyze import CONTAINER_SECONDS, DOCKER_LABEL, PROTOCOL, container_command, evidence_hashes, inspect_run, require
from bundle import context_identity, image_identity, verify_bundle


def remove_container(context, run_id):
    docker = ["docker", "--host", context["endpoint"]]

    def containers():
        result = subprocess.check_output([*docker, "ps", "--all", "--quiet", "--no-trunc",
                                          "--filter", f"label={DOCKER_LABEL}={run_id}"], text=True, timeout=10)
        ids = result.split()
        require(all(re.fullmatch(r"[0-9a-f]{64}", value) for value in ids), "Unexpected container ID")
        return ids

    ids = containers()
    if ids:
        subprocess.run([*docker, "rm", "--force", *ids], stdout=subprocess.DEVNULL,
                       stderr=subprocess.DEVNULL, timeout=15)
    require(not containers(), "The task-owned container is still present")


def write(path, value):
    with path.open("x", encoding="utf-8") as stream:
        json.dump(value, stream, indent=2, allow_nan=False)
        stream.write("\n")


def verify_local_image(bundle):
    context = bundle["dockerContext"]
    require(context_identity(context["name"]) == context, "Reviewed Docker context endpoint changed")
    require(image_identity(bundle["image"], context) == bundle["imageIdentity"], "Wrong cached image identity")


def create_task_output(output, bundle_root):
    require(not output.is_symlink(), "Run output cannot be a symbolic link")
    output = output.resolve()
    require(not output.is_relative_to(bundle_root), "Run output must be outside the immutable bundle")
    output.mkdir(mode=0o700, parents=True, exist_ok=False)
    temporary = output / "tmp"
    temporary.mkdir(mode=0o700, exist_ok=False)
    for directory in (output, temporary):
        mode = directory.lstat().st_mode
        require(stat.S_ISDIR(mode) and stat.S_IMODE(mode) == 0o700, f"Expected private directory mode 0700 for {directory}")
    return output


def launch(root, bundle, bundle_hash, output):
    require(Path(__file__).resolve() == root / "tools/launch.py", "Invoke the reviewed bundled launcher")
    output = create_task_output(output, root)
    run_id = uuid.uuid4().hex
    invocation = container_command(root, output, bundle, run_id)
    provenance = {"kind": "launch", "protocol": PROTOCOL, "acceptance": False, "runId": run_id,
                  "bundleReceiptSha256": bundle_hash, "sourceRevision": bundle["sourceRevision"], "image": bundle["image"],
                  "coreClrSha256": bundle["files"]["app/libcoreclr.so"], "command": invocation,
                  "dockerContext": bundle["dockerContext"], "bundlePath": str(root), "outputPath": str(output),
                  "containerTimeoutSeconds": CONTAINER_SECONDS}
    write(output / "launch.json", provenance)
    errors = []
    removed = False
    started = time.monotonic()
    try:
        with (output / "container.log").open("x", encoding="utf-8") as log:
            process = subprocess.run(invocation, stdout=log, stderr=subprocess.STDOUT, timeout=CONTAINER_SECONDS)
        status = process.returncode
    except subprocess.TimeoutExpired:
        status = 124
        errors.append("The container reached the fixed outer deadline")
    except KeyboardInterrupt:
        status = 130
        errors.append("Container execution was interrupted")
    except OSError as error:
        status = 125
        errors.append(str(error))
    finally:
        wall = time.monotonic() - started
        try:
            remove_container(bundle["dockerContext"], run_id)
            removed = True
        except (OSError, subprocess.SubprocessError, ValueError) as error:
            errors.append(f"Container cleanup failed: {error}")
    (output / "container-exit-status.txt").write_text(f"{status}\n", encoding="ascii")
    receipt = {**provenance, "kind": "run", "containerExitStatus": status, "wallSeconds": wall,
               "containerRemoved": removed, "diagnosticCompleted": False}
    try:
        verify_bundle(root, bundle_hash)
        verify_local_image(bundle)
        receipt["evidenceSha256"] = evidence_hashes(output)
        inspected = inspect_run(output, provenance, wall)
        for field in ("diagnosticOutcome", "testExitStatus", "descriptorTests", "experimentTests"):
            receipt[field] = inspected[field]
    except (OSError, KeyError, TypeError, ValueError, ET.ParseError, subprocess.SubprocessError) as error:
        errors.append(f"Evidence validation failed: {error}")
    receipt["diagnosticCompleted"] = not errors
    if errors:
        receipt["validationErrors"] = errors
    # Return the original container status when the test rethrows its first resource cap.
    receipt["launcherExitStatus"] = (128 - status if status < 0 else status) or int(bool(errors))
    write(output / "run-receipt.json", receipt)
    print(json.dumps(receipt, indent=2, allow_nan=False))
    return receipt["launcherExitStatus"]


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--bundle", type=Path, required=True)
    parser.add_argument("--receipt-sha256", required=True)
    action = parser.add_mutually_exclusive_group(required=True)
    action.add_argument("--verify-only", action="store_true")
    action.add_argument("--output", type=Path)
    options = parser.parse_args()
    try:
        root = options.bundle.resolve(strict=True)
        bundle = verify_bundle(root, options.receipt_sha256)
        verify_local_image(bundle)
        if options.verify_only:
            print(json.dumps({"protocol": PROTOCOL, "sourceRevision": bundle["sourceRevision"],
                              "receiptSha256": options.receipt_sha256, "verified": True,
                              "acceptance": False, "workloadLaunched": False}, indent=2))
            return 0
        return launch(root, bundle, options.receipt_sha256, options.output)
    except (OSError, KeyError, TypeError, ValueError, subprocess.SubprocessError) as error:
        print(f"Late-slope launcher rejected the request: {error}", file=sys.stderr)
        return 1


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