"""Perpetual Compute workload SDK.

Signals workload completion or failure to the Perpetual Compute host so
billing stops at the right moment. Works from any container, any language
runtime — see perpetual-done and perpetual-fail for shell/non-Python use.

Usage
-----
Auto (atexit hook, zero config)::

    import perpetual_compute          # hooks registered on import
    run_my_job()
    # done/failed signaled automatically on process exit

Context manager (recommended)::

    import perpetual_compute
    with perpetual_compute.workload():
        run_my_job()                  # done on clean exit, failed on exception

Decorator::

    import perpetual_compute

    @perpetual_compute.task
    def main():
        run_my_job()

    main()

Explicit::

    import perpetual_compute
    run_my_job()
    perpetual_compute.done(reason="All 1000 items processed")   # writes + exits
    perpetual_compute.fail(reason="OOM at batch 400", exit_code=137)

Metrics (visible in the portal workload detail view)::

    import perpetual_compute as pc
    for epoch in range(100):
        loss = train_epoch()
        pc.log_metric("epoch", epoch)
        pc.log_metric("loss", loss)
        pc.log_metric("throughput_samples_per_sec", 1024.0)
    # Metrics are flushed live to /run/pc-status/metrics.json and included
    # in the final done.json payload so the portal can display them.

How it works
------------
The SDK writes a small JSON file to /run/pc-status/done.json — a directory
mounted by the Perpetual Compute host into every workload container. The host
reads this file after 'docker wait' detects that the container has exited and
calls the orchestrator /stop endpoint with the correct status.

The done secret never enters the container. A client cannot stop billing
without their container actually exiting.
"""

from __future__ import annotations

import contextlib
import sys
from typing import Any, Callable, Optional, TypeVar

from perpetual_compute._hooks import register
from perpetual_compute._status import log_metric, write_status

__version__ = "0.1.0"
__all__ = ["done", "fail", "workload", "task", "log_metric"]

# Register atexit + SIGTERM hooks immediately on import.
register()

F = TypeVar("F", bound=Callable[..., Any])


def done(reason: Optional[str] = None, exit_code: int = 0) -> None:
    """Signal successful workload completion and exit the process.

    Writes /run/pc-status/done.json with status='completed', then calls
    sys.exit(exit_code). This must be the last call your script makes.

    Args:
        reason:    Optional human-readable message shown in the portal.
        exit_code: Process exit code (default 0). Prefer 0 for success.
    """
    write_status("completed", reason=reason, exit_code=exit_code)
    sys.exit(exit_code)


def fail(
    reason: Optional[str] = None,
    exit_code: int = 1,
) -> None:
    """Signal workload failure and exit the process.

    Writes /run/pc-status/done.json with status='failed', then calls
    sys.exit(exit_code). This must be the last call your script makes.

    Args:
        reason:    Optional human-readable message shown in the portal.
        exit_code: Process exit code (default 1). Use the actual error code
                   when available (e.g. 137 for OOM kill, 1 for generic error).
    """
    write_status("failed", reason=reason, exit_code=exit_code)
    sys.exit(exit_code)


@contextlib.contextmanager
def workload(reason_on_success: Optional[str] = None):
    """Context manager: signal done on clean exit, fail on exception.

    The process is NOT exited when the context manager exits — normal Python
    cleanup continues. The atexit hook writes the final status on process exit.

    Example::

        with perpetual_compute.workload():
            train_model()

    Raises:
        Re-raises any exception after writing failed status.
    """
    try:
        yield
    except (KeyboardInterrupt, SystemExit):
        write_status("failed", reason="Interrupted")
        raise
    except Exception as exc:
        write_status("failed", reason=f"{type(exc).__name__}: {exc}", exit_code=1)
        raise
    else:
        write_status("completed", reason=reason_on_success)


def task(fn: F) -> F:
    """Decorator: wrap a function with the workload() context manager.

    Example::

        @perpetual_compute.task
        def main():
            train_model()

        main()
    """
    import functools

    @functools.wraps(fn)
    def wrapper(*args: Any, **kwargs: Any) -> Any:
        with workload():
            return fn(*args, **kwargs)

    return wrapper  # type: ignore[return-value]
