#!/usr/bin/env python3
"""Vicena Compute native compute-bundle runner.

Copy this runner unchanged, write ``vicena_compute_bundle_spec.json`` beside
it, create the declared source directory, and run the runner.
"""

from __future__ import annotations

import hashlib
import json
from pathlib import Path

import vicena_compute


EXAMPLE_SPEC = {
    "environment": "openfoam-13",
    "title": "OpenFOAM compute bundle",
    "source_dir": "vicena_bundle",
    "run_label": "initial",
    "argv": ["bash", "run.sh"],
    "outputs": ["results"],
    "limits": {"cpu_cores": 8, "memory_mib": 16384, "wall_seconds": 1800},
}
SPEC = EXAMPLE_SPEC
RUNNER_DIR = Path(__file__).resolve().parent
SPEC_PATH = RUNNER_DIR / "vicena_compute_bundle_spec.json"
STATE_PATH = RUNNER_DIR / "vicena_compute_bundle.json"
LEGACY_STATE_PATH = RUNNER_DIR / "vicena_remote_project.json"
RESULTS_DIR = RUNNER_DIR / "vicena_compute_results"
MAX_WAIT_SECONDS = 90
MAX_SPEC_BYTES = 1024 * 1024


def main() -> int:
    global SPEC
    require_session_local_runner()
    SPEC = load_spec()
    state = load_state()
    client = vicena_compute.BundleClient()
    safe_preflight(client)
    source_sha256 = source_tree_sha256()
    key = stable_key(source_sha256)
    jobs = state.setdefault("jobs", {})
    if "job_id" in state and key not in jobs:
        jobs[key] = {"job_id": state.pop("job_id"), "run_label": SPEC["run_label"]}
        save_state(state)
    if key not in jobs:
        existing = client.job_for_idempotency(key)
        if existing is not None:
            jobs[key] = {
                "job_id": existing.id,
                "run_label": SPEC["run_label"],
            }
            save_state(state)
    if key not in jobs:
        project = source_bundle(client, state, source_sha256)
        job = project.run(
            SPEC["argv"],
            SPEC["outputs"],
            idempotency_key=key,
            title=SPEC["title"],
            limits=SPEC["limits"],
        )
        jobs[key] = {"job_id": job.id, "run_label": SPEC["run_label"]}
        save_state(state)
    job = vicena_compute.Job(client, jobs[key]["job_id"])
    status = job.wait(timeout=MAX_WAIT_SECONDS)
    failed = False
    result = None
    if status["status"] in vicena_compute.TERMINAL_STATUSES:
        RESULTS_DIR.mkdir(exist_ok=True)
        try:
            result = job.result()
        except vicena_compute.VicenaComputeError as error:
            result = status
            result["result_error"] = f"{type(error).__name__}: {error}"
            failed = True
        (RESULTS_DIR / f"{job.id}.json").write_text(
            json.dumps(result, indent=2, sort_keys=True) + "\n"
        )
        jobs[key]["status"] = status["status"]
        jobs[key]["workspace_output_folder"] = status.get("workspace_output_folder")
        save_state(state)
        if status.get("workspace_output_folder"):
            print("Workspace results:", status["workspace_output_folder"])
        failed = failed or status["status"] != "completed"
    print_receipt(job.id, status, result)
    return 1 if failed else 0


def require_session_local_runner() -> None:
    sessions_root = Path("/root/sessions")
    if RUNNER_DIR.is_relative_to(Path("/root")) and (
        RUNNER_DIR == sessions_root or not RUNNER_DIR.is_relative_to(sessions_root)
    ):
        raise SystemExit(
            "Vicena Compute submission runners must be staged inside the active "
            "Compute session; /root is shared across conversations."
        )


def load_spec() -> dict:
    if not SPEC_PATH.is_file():
        raise SystemExit(
            f"Missing {SPEC_PATH.name}. Write the documented bundle JSON object "
            "beside this unchanged runner."
        )
    if SPEC_PATH.stat().st_size > MAX_SPEC_BYTES:
        raise ValueError(f"{SPEC_PATH.name} exceeds {MAX_SPEC_BYTES} bytes")
    value = json.loads(SPEC_PATH.read_text(encoding="utf-8"))
    if not isinstance(value, dict):
        raise ValueError(f"{SPEC_PATH.name} must contain one JSON object")
    required = {
        "environment",
        "title",
        "source_dir",
        "run_label",
        "argv",
        "outputs",
        "limits",
    }
    if set(value) != required:
        missing = sorted(required - set(value))
        unexpected = sorted(set(value) - required)
        raise ValueError(
            f"invalid bundle spec fields; missing={missing}, unexpected={unexpected}"
        )
    for field in ("environment", "title", "source_dir", "run_label"):
        if (
            not isinstance(value[field], str)
            or not value[field].strip()
            or len(value[field]) > 4096
            or "\x00" in value[field]
        ):
            raise ValueError(f"{field} must be a non-empty string")
    validate_string_list(value["argv"], "argv")
    validate_string_list(value["outputs"], "outputs")
    validate_output_paths(value["outputs"])
    validate_limits(value["limits"])
    return value


def validate_string_list(value: object, field: str) -> None:
    if (
        not isinstance(value, list)
        or not value
        or len(value) > 64
        or not all(
            isinstance(item, str)
            and item
            and len(item) <= 4096
            and "\x00" not in item
            for item in value
        )
    ):
        raise ValueError(f"{field} must contain 1 to 64 non-empty strings")


def validate_limits(value: object) -> None:
    if not isinstance(value, dict):
        raise ValueError("limits must be an object")
    allowed = {"cpu_cores", "memory_mib", "wall_seconds"}
    if set(value) - allowed:
        raise ValueError("limits contains unsupported fields")
    if not all(
        isinstance(item, int) and not isinstance(item, bool) and item > 0
        for item in value.values()
    ):
        raise ValueError("limits values must be positive integers")


def validate_output_paths(value: list[str]) -> None:
    for item in value:
        path = Path(item)
        if path.is_absolute() or ".." in path.parts or item in {".", ""}:
            raise ValueError(f"output path must be bundle-relative: {item!r}")


def safe_preflight(client) -> None:
    environments = {
        value.get("id"): value
        for value in client.environments()
        if isinstance(value, dict)
    }
    environment = environments.get(SPEC["environment"])
    if environment is None:
        raise RuntimeError(f"environment is not registered: {SPEC['environment']}")
    resources = environment.get("resources")
    if not isinstance(resources, dict):
        raise RuntimeError("registered environment has no resource contract")
    for field, maximum_field in (
        ("cpu_cores", "max_cpu_cores"),
        ("memory_mib", "max_memory_mib"),
        ("wall_seconds", "max_wall_seconds"),
    ):
        requested = SPEC["limits"].get(field)
        maximum = resources.get(maximum_field)
        if requested is not None and (
            not isinstance(maximum, int) or requested > maximum
        ):
            raise ValueError(
                f"{field}={requested} exceeds registered {maximum_field}={maximum}"
            )
    print(
        "Vicena Compute bundle profile:",
        SPEC["environment"],
        environment.get("tool"),
        environment.get("version"),
    )


def source_directory() -> Path:
    path = Path(SPEC["source_dir"])
    runner_dir = RUNNER_DIR.resolve()
    resolved = (path if path.is_absolute() else runner_dir / path).resolve()
    if not resolved.is_relative_to(runner_dir):
        raise ValueError("source_dir must stay inside the active Compute session")
    return resolved


def stable_key(source_sha256: str | None = None) -> str:
    value = {key: SPEC[key] for key in sorted(SPEC) if key != "title"}
    value["source_sha256"] = source_sha256 or source_tree_sha256()
    canonical = json.dumps(value, sort_keys=True, separators=(",", ":")).encode()
    return "remote-project-" + hashlib.sha256(canonical).hexdigest()[:32]


def source_tree_sha256() -> str:
    root = source_directory()
    if not root.is_dir():
        raise ValueError(f"source directory does not exist: {root}")
    digest = hashlib.sha256()
    file_count = 0
    for path in sorted(root.rglob("*")):
        relative = path.relative_to(root).as_posix()
        if path.is_symlink():
            raise ValueError(f"source tree contains a symlink: {relative}")
        if not path.is_file():
            continue
        file_digest = hashlib.sha256()
        with path.open("rb") as stream:
            while chunk := stream.read(1024 * 1024):
                file_digest.update(chunk)
        digest.update(relative.encode())
        digest.update(b"\0")
        digest.update(file_digest.hexdigest().encode())
        digest.update(b"\0")
        file_count += 1
    if file_count == 0:
        raise ValueError("source directory contains no files")
    return digest.hexdigest()


def source_bundle(client, state: dict, source_sha256: str):
    source_key = project_state_key(source_sha256)
    projects = state.setdefault("projects", {})
    entry = projects.get(source_key)
    if entry is None:
        legacy_entry = projects.get(source_sha256)
        if (
            isinstance(legacy_entry, dict)
            and legacy_entry.get("environment") == SPEC["environment"]
        ):
            entry = legacy_entry
            projects[source_key] = projects.pop(source_sha256)
            save_state(state)
    if entry is None:
        project = client.create_bundle(
            SPEC["environment"],
            title=SPEC["title"],
            ephemeral_key=source_key,
        )
        entry = {
            "project_id": project.id,
            "environment": SPEC["environment"],
            "source_uploaded": False,
        }
        projects[source_key] = entry
        save_state(state)
    else:
        project = client.open_bundle(
            entry["project_id"], SPEC["environment"], title=SPEC["title"]
        )
    if not entry.get("source_uploaded"):
        project.upload_tree(source_directory())
        entry["source_uploaded"] = True
        save_state(state)
    return project


def project_state_key(source_sha256: str) -> str:
    source = {
        "environment": SPEC["environment"],
        "source_sha256": source_sha256,
    }
    return hashlib.sha256(
        json.dumps(source, sort_keys=True, separators=(",", ":")).encode()
    ).hexdigest()


def load_state() -> dict:
    path = STATE_PATH if STATE_PATH.is_file() else LEGACY_STATE_PATH
    if not path.is_file():
        return {}
    value = json.loads(path.read_text())
    return value if isinstance(value, dict) else {}


def save_state(value: dict) -> None:
    temporary = STATE_PATH.with_suffix(".tmp")
    temporary.write_text(json.dumps(value, indent=2, sort_keys=True) + "\n")
    temporary.replace(STATE_PATH)


def print_receipt(job_id: str, status: dict, result: dict | None) -> None:
    receipt = {
        "job_id": job_id,
        "environment": SPEC["environment"],
        "run_label": SPEC["run_label"],
        "status": status.get("status"),
        "backend_job_id": status.get("backend_job_id"),
        "failure_phase": status.get("failure_phase"),
        "failure_kind": status.get("failure_kind"),
        "error": status.get("error"),
        "workspace_output_folder": status.get("workspace_output_folder"),
        "result": result,
    }
    print("Vicena Compute receipt:")
    print(json.dumps(receipt, indent=2, sort_keys=True))


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