"""One-shot, caller-authenticated single-Space CLI. No inference on import.

Examples:
    python client.py presets
    python client.py health
    python client.py plan --input request.json
    python client.py run --input request.json --run-dir runs/unique-name \
        --expected-revision <approved-generator-commit-sha>

Run requires HF_INFERENCE_TOKEN (or --token-env NAME) and the operator's
HF_POOL_ZEROGPU_OVERAGE_DISABLED=1 attestation. This client does not change or
verify the account's billing settings. No SDK credential lookup is allowed.

A claimed directory is never reused, including after errors. An unknown hop
may still be running remotely. Reconcile its event_id/session_hash manually;
there is deliberately no resume, retry, cancellation, or token rotation.
"""
from __future__ import annotations

import argparse
from concurrent.futures import Future
from contextlib import contextmanager, redirect_stderr, redirect_stdout, suppress
from dataclasses import dataclass
import hashlib
import json
import math
import os
from pathlib import Path
import re
import signal
import stat
import sys
import tempfile
import threading
import time
from urllib.request import HTTPRedirectHandler, Request, build_opener

GENERATOR = "https://owenisas-workflow-studio.hf.space"
GENERATOR_REPO = "owenisas/workflow-studio"


class AuthenticationRequired(ValueError):
    pass


class OverageAttestationRequired(ValueError):
    pass


class OwnerTokenForbidden(ValueError):
    pass


class RevisionMismatch(ValueError):
    pass


class RuntimeNotReady(ValueError):
    pass


class IntegrityError(ValueError):
    pass


class InputError(ValueError):
    pass


class CLIArgumentError(ValueError):
    pass


class TerminationRequested(KeyboardInterrupt):
    pass


@contextmanager
def _termination_guard():
    if threading.current_thread() is not threading.main_thread():
        yield
        return
    previous = signal.getsignal(signal.SIGTERM)

    def terminate(signum, frame):
        raise TerminationRequested()

    signal.signal(signal.SIGTERM, terminate)
    try:
        yield
    finally:
        signal.signal(signal.SIGTERM, previous)


def _unique_object(pairs):
    result = {}
    for key, value in pairs:
        if key in result:
            raise InputError()
        result[key] = value
    return result


def _bad_constant(value):
    raise InputError()


def _local_path(value, base):
    if not isinstance(value, str) or not value or "://" in value or "\x00" in value:
        raise InputError()
    path = Path(value)
    if not path.is_absolute():
        path = base / path
    if path.is_symlink() or not path.is_file():
        raise InputError()
    return path.resolve(strict=True)


@dataclass(frozen=True)
class RequestInput:
    preset_id: str
    prompt: str
    canvas: str = ""
    duration: float = 5
    seed: int = 42
    image: Path | None = None
    last_image: Path | None = None

    def plan_args(self):
        return (self.preset_id, self.prompt, self.canvas, self.duration, self.seed,
                self.image is not None, self.last_image is not None)


def load_input(filename):
    path = _local_path(str(filename), Path.cwd())
    with _open_regular(path) as handle:
        raw = handle.read(1024 * 1024 + 1)
    if len(raw) > 1024 * 1024:
        raise InputError()
    data = json.loads(raw, object_pairs_hook=_unique_object, parse_constant=_bad_constant)
    allowed = {"preset_id", "prompt", "canvas", "duration", "seed", "image", "last_image"}
    if not isinstance(data, dict) or set(data) - allowed:
        raise InputError()
    for field in ("preset_id", "prompt"):
        if not isinstance(data.get(field), str) or not data[field].strip():
            raise InputError()
    if (len(data["prompt"]) > 12000 or data["prompt"] != data["prompt"].strip()
            or not isinstance(data.get("canvas", ""), str)):
        raise InputError()
    duration, seed = data.get("duration", 5), data.get("seed", 42)
    if type(duration) not in (int, float) or not 2 <= duration <= 8 or not math.isfinite(duration):
        raise InputError()
    if type(seed) is not int or not 0 <= seed <= 2147483647:
        raise InputError()
    for field in ("image", "last_image"):
        if data.get(field) is not None:
            data[field] = _local_path(data[field], path.parent)
    return RequestInput(**data)


def _open_regular(path):
    # Nonblocking prevents a swapped FIFO from hanging before fstat().
    fd = os.open(path, os.O_RDONLY | os.O_NOFOLLOW | os.O_NONBLOCK)
    try:
        if not stat.S_ISREG(os.fstat(fd).st_mode):
            raise InputError()
        return os.fdopen(fd, "rb")
    except BaseException:
        os.close(fd)
        raise


def _fsync_dir(path):
    fd = os.open(path, os.O_RDONLY)
    try:
        os.fsync(fd)
    finally:
        os.close(fd)


def _sanitize(value, token=""):
    if isinstance(value, str):
        if token:
            value = value.replace(token, "[REDACTED]")
        return re.sub(r"hf_[A-Za-z0-9_]+", "[REDACTED]", value)
    if isinstance(value, dict):
        return {_sanitize(str(k), token): _sanitize(v, token) for k, v in value.items()
                if not re.search(r"token|authorization|password|secret|signature", str(k), re.I)}
    if isinstance(value, (list, tuple)):
        return [_sanitize(v, token) for v in value]
    if value is None or type(value) in (bool, int):
        return value
    if type(value) is float and math.isfinite(value):
        return value
    raise ValueError()


def _atomic_json(path, data):
    # Both initial publication and replacements have private modes and fsync.
    fd, temporary = tempfile.mkstemp(prefix=".journal-", dir=path.parent)
    try:
        with os.fdopen(fd, "w", encoding="utf-8") as handle:
            os.fchmod(handle.fileno(), 0o600)
            json.dump(data, handle, ensure_ascii=True, allow_nan=False, sort_keys=True)
            handle.write("\n")
            handle.flush()
            os.fsync(handle.fileno())
        os.replace(temporary, path)
        _fsync_dir(path.parent)
    finally:
        with suppress(FileNotFoundError):
            os.unlink(temporary)


class Journal:
    def __init__(self, directory, token, revision):
        # Resolve the parent only: a final symlink must fail the exclusive mkdir.
        directory = Path(directory)
        self.directory = directory.parent.resolve(strict=True) / directory.name
        self.directory.mkdir(mode=0o700, exist_ok=False)
        self.directory.chmod(0o700)
        _fsync_dir(self.directory.parent)
        self.token = token
        self.data = {
            "schema_version": 2, "state": "preflight", "checkpoint": 0,
            "generator": GENERATOR, "generator_revision": revision,
            "overage_disabled_operator_attestation": True,
            "hops": {"workflow": {"state": "not_submitted", "event_id": None, "session_hash": None}},
        }
        self.save()

    def save(self):
        self.data["checkpoint"] += 1
        self.data["updated_at_unix"] = time.time()
        _atomic_json(self.directory / "journal.json", _sanitize(self.data, self.token))


def _copy_private(source, destination, expected=None):
    """Snapshot regular bytes, detect concurrent source writes, then publish."""
    fd, temporary = tempfile.mkstemp(prefix=".copy-", dir=destination.parent)
    hasher = hashlib.sha256()
    size = 0
    try:
        with os.fdopen(fd, "wb") as target, _open_regular(source) as original:
            before = os.fstat(original.fileno())
            while chunk := original.read(1024 * 1024):
                hasher.update(chunk)
                size += len(chunk)
                target.write(chunk)
            after = os.fstat(original.fileno())
            stamp = lambda s: (s.st_dev, s.st_ino, s.st_size, s.st_mtime_ns, s.st_ctime_ns)
            if stamp(before) != stamp(after) or size != before.st_size:
                raise IntegrityError()
            digest = hasher.hexdigest()
            if expected is not None and digest != expected:
                raise IntegrityError()
            os.fchmod(target.fileno(), 0o600)
            target.flush()
            os.fsync(target.fileno())
        os.replace(temporary, destination)
        _fsync_dir(destination.parent)
        return {"path": destination.name, "sha256": digest, "bytes": size}
    finally:
        with suppress(FileNotFoundError):
            os.unlink(temporary)


def _verify_file(path, expected):
    with _open_regular(path) as handle:
        actual = hashlib.file_digest(handle, "sha256").hexdigest()
    if actual != expected:
        raise IntegrityError()


class _NoRedirect(HTTPRedirectHandler):
    def redirect_request(self, req, fp, code, msg, headers, newurl):
        # Never forward caller authorization to an unexpected destination.
        return None


def _hub_json(route, token):
    request = Request("https://huggingface.co/api/" + route,
                      headers={"Authorization": "Bearer " + token, "Accept": "application/json"})
    with build_opener(_NoRedirect).open(request, timeout=20) as response:
        return json.load(response)


def fetch_identity(token):
    return _hub_json("whoami-v2", token)


def fetch_metadata(repo, token):
    if repo != GENERATOR_REPO:
        raise ValueError()
    return _hub_json("spaces/" + repo, token)


def _client_factory(src, **options):
    from gradio_client import Client
    # 2.5.0 uses an OR with this env even when analytics_enabled=False.
    os.environ["GRADIO_ANALYTICS_ENABLED"] = "False"
    return Client(src, **options)


def _file_wrapper(path):
    from gradio_client import handle_file
    return handle_file(str(path))


def _new_client(factory, src, token, downloads: bool | Path = False):
    # token=False is intentional: None permits saved SDK owner credentials.
    return factory(src, token=token, max_workers=8, verbose=False,
                   download_files=downloads, analytics_enabled=False,
                   httpx_kwargs={"timeout": 20})


def _within_deadline(call, deadline, clock, cleanup=None):
    """Bound a read-only preflight operation, never an inference submission.

    SDK predict/initialization may wait beyond an HTTP request timeout. The
    daemon performs only this one read; it cannot advance to a GPU submission
    or update the journal after the caller times out. Late-created clients close.
    """
    if clock() >= deadline:
        raise TimeoutError()
    future = Future()

    def invoke():
        try:
            future.set_result(call())
        except BaseException as error:
            future.set_exception(error)

    threading.Thread(target=invoke, name="workflow-preflight", daemon=True).start()
    try:
        remaining = deadline - clock()
        if remaining <= 0:
            raise TimeoutError()
        return future.result(timeout=remaining)
    except BaseException:
        if cleanup is not None:
            def close_late(done):
                with suppress(Exception):
                    if not done.cancelled() and done.exception() is None:
                        cleanup(done.result())
            future.add_done_callback(close_late)
        raise


def _capture(hop, client, job=None):
    session = getattr(client, "session_hash", None)
    if isinstance(session, str):
        hop["session_hash"] = session
    event = getattr(getattr(job, "communicator", None), "event_id", None)
    if isinstance(event, str):
        hop["event_id"] = event


def _submit_once(client, name, args, api, journal, deadline, clock):
    if clock() >= deadline:
        raise TimeoutError()
    hop = journal.data["hops"][name]
    # An abrupt process death after this fsync is ambiguous, never resubmittable.
    hop["state"] = "unknown"
    journal.data["state"] = "unknown"
    journal.data["active_hop"] = name
    _capture(hop, client)
    journal.save()
    job = None
    try:
        job = client.submit(*args, api_name=api)
        while True:
            _capture(hop, client, job)
            journal.save()
            remaining = deadline - clock()
            if remaining <= 0:
                raise TimeoutError()
            try:
                result = job.future.result(timeout=min(2.0, remaining))
            except TimeoutError:
                if job.future.done():
                    raise
                continue
            _capture(hop, client, job)
            hop["state"] = "complete"
            journal.data["state"] = name + "_complete"
            journal.save()
            return result
    finally:
        # Event IDs often arrive while waiting, including just before errors.
        _capture(hop, client, job)
        journal.save()


def _artifact(result, downloads):
    if not isinstance(result, (tuple, list)) or len(result) != 2 or not isinstance(result[1], dict):
        raise IntegrityError()
    value = result[0]
    # gr.Video may deserialize into (video_path, subtitle_path).
    if isinstance(value, (list, tuple)) and len(value) == 2:
        value = value[0]
    if isinstance(value, dict):
        value = value.get("path") or value.get("video")
        if isinstance(value, dict):
            value = value.get("path")
    if isinstance(value, Path):
        value = str(value)
    path = _local_path(value, downloads)
    if not path.is_relative_to(downloads.resolve()):
        raise IntegrityError()
    return path, result[1]


def _sha(value):
    if not isinstance(value, str) or re.fullmatch(r"[0-9a-f]{64}", value) is None:
        raise IntegrityError()
    return value


def _plan(client, request):
    result = client.predict(*request.plan_args(), api_name="/plan")
    if (not isinstance(result, dict) or not isinstance(result.get("canvas"), str)
            or not result["canvas"] or not isinstance(result.get("adapter"), str)):
        raise InputError()
    prompt = result.get("prompt")
    if not isinstance(prompt, str) or not prompt or prompt != prompt.strip() or len(prompt) > 12000:
        raise InputError()
    for field in ("height", "width", "num_frames", "steps"):
        if type(result.get(field)) is not int or result[field] <= 0:
            raise InputError()
    return result


def _run(args, request, token, factory, metadata, identity, wrap, clock):
    journal = Journal(args.run_dir, token, args.expected_revision)
    clients = []
    deadline = clock() + args.max_wall_seconds
    try:
        staged = {}
        for name in ("image", "last_image"):
            source = getattr(request, name)
            if source is not None:
                suffix = source.suffix if re.fullmatch(r"\.[A-Za-z0-9]{1,10}", source.suffix) else ".bin"
                target = journal.directory / (name + suffix)
                staged[name] = _copy_private(source, target)
        journal.data["inputs"] = staged
        journal.data["request"] = {"preset_id": request.preset_id, "prompt": request.prompt,
                                   "canvas": request.canvas, "duration": request.duration, "seed": request.seed}
        journal.save()
        caller = _within_deadline(lambda: identity(token), deadline, clock).get("name")
        if not isinstance(caller, str) or not caller:
            raise AuthenticationRequired()
        if caller.casefold() == "owenisas":
            raise OwnerTokenForbidden()
        journal.data["caller"] = caller
        if _within_deadline(lambda: metadata(GENERATOR_REPO, token), deadline, clock).get("sha") != args.expected_revision:
            raise RevisionMismatch()
        downloads = journal.directory / "workflow-downloads"
        downloads.mkdir(mode=0o700)
        generator = _within_deadline(
            lambda: _new_client(factory, GENERATOR, token, downloads),
            deadline, clock, cleanup=lambda client: client.close())
        clients.append(generator)
        health = _within_deadline(lambda: generator.predict(api_name="/health"), deadline, clock)
        # Models are staged and loaded inside GPU calls, not held resident here.
        if (not isinstance(health, dict) or type(health.get("schema_version")) is not int
                or health["schema_version"] != 2 or health.get("status") != "ready"
                or health.get("blockers") != [] or health.get("runtime_initialized") is not True
                or health.get("conditioning_transport") != "in_process"):
            raise RuntimeNotReady()
        journal.data["health"] = health
        plan = _within_deadline(lambda: _plan(generator, request), deadline, clock)
        journal.data["plan"] = plan
        journal.save()

        def image(name):
            if name not in staged:
                return None
            path = journal.directory / staged[name]["path"]
            _verify_file(path, staged[name]["sha256"])
            return wrap(path)

        first, last = image("image"), image("last_image")
        result = _submit_once(generator, "workflow",
                              (request.preset_id, plan["prompt"], first, last, plan["canvas"],
                               request.duration, request.seed),
                              "/generate", journal, deadline, clock)
        path, report = _artifact(result, downloads)
        expected = _sha(report.get("output_sha256"))
        output = _copy_private(path, journal.directory / "final.mp4", expected)
        journal.data["result"] = _sanitize(report, token)
        journal.data["output"] = output
        journal.data["state"] = "succeeded"
        journal.save()
        return {"state": "succeeded", "output": str(journal.directory / "final.mp4"),
                "result": journal.data["result"]}
    except BaseException as error:
        journal.data["state"] = ("unknown" if any(h["state"] == "unknown"
                                  for h in journal.data["hops"].values()) else "failed")
        journal.data["error"] = type(error).__name__
        journal.save()
        raise
    finally:
        for client in clients:
            with suppress(Exception):
                client.close()


class _Parser(argparse.ArgumentParser):
    def error(self, message):
        raise CLIArgumentError()


def _parser():
    parser = _Parser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    commands = parser.add_subparsers(dest="command", required=True)
    commands.add_parser("presets", help="List public presets without credentials")
    commands.add_parser("health", help="Read public runtime readiness without credentials")
    plan = commands.add_parser("plan", help="Validate and plan; no files uploaded")
    plan.add_argument("--input", required=True, type=Path)
    run = commands.add_parser("run", help="One-shot authenticated single-Space inference")
    run.add_argument("--input", required=True, type=Path)
    run.add_argument("--run-dir", required=True, type=Path)
    run.add_argument("--expected-revision", required=True)
    run.add_argument("--token-env", default="HF_INFERENCE_TOKEN")
    run.add_argument("--max-wall-seconds", default=900.0, type=float)
    return parser


def main(argv=None, *, env=None, client_factory=None, metadata_fetcher=None,
         identity_fetcher=None, file_wrapper=None, clock=time.monotonic):
    """Dependencies are injectable for offline tests; return an exit status."""
    env = os.environ if env is None else env
    token = ""
    try:
        args = _parser().parse_args(argv)
        request = load_input(args.input) if args.command in ("plan", "run") else None
        factory = client_factory or _client_factory
        if args.command == "run":
            if re.fullmatch(r"[0-9a-f]{40}", args.expected_revision) is None:
                raise InputError()
            if not math.isfinite(args.max_wall_seconds) or not 0 < args.max_wall_seconds <= 86400:
                raise InputError()
            if re.fullmatch(r"[A-Za-z_][A-Za-z0-9_]*", args.token_env) is None:
                raise InputError()
            token = env.get(args.token_env, "")
            if not isinstance(token, str) or not token or any(c.isspace() for c in token):
                raise AuthenticationRequired()
            if env.get("HF_POOL_ZEROGPU_OVERAGE_DISABLED") != "1":
                raise OverageAttestationRequired()
        # Do not relay SDK prints, provider errors, or warnings to CLI output.
        with _termination_guard(), open(os.devnull, "w") as sink, redirect_stdout(sink), redirect_stderr(sink):
            if args.command == "run":
                result = _run(args, request, token, factory, metadata_fetcher or fetch_metadata,
                              identity_fetcher or fetch_identity, file_wrapper or _file_wrapper, clock)
            else:
                client = _new_client(factory, GENERATOR, False)
                try:
                    result = (_plan(client, request) if args.command == "plan"
                              else client.predict(api_name="/" + args.command))
                finally:
                    with suppress(Exception):
                        client.close()
        print(json.dumps(_sanitize(result, token), sort_keys=True, allow_nan=False))
        return 0
    except SystemExit as error:
        # argparse help is the only non-JSON successful output.
        return int(error.code or 0)
    except (Exception, KeyboardInterrupt) as error:
        print(json.dumps({"state": "error", "error": _sanitize(type(error).__name__, token)}), file=sys.stderr)
        if isinstance(error, TerminationRequested):
            return 143
        return 130 if isinstance(error, KeyboardInterrupt) else 1


if __name__ == "__main__":
    code = main()
    sys.stdout.flush()
    sys.stderr.flush()
    # ThreadPoolExecutor's atexit join can outlive an unknown remote job. The
    # CLI must exit within its polling bound without cancelling that job.
    # All durable writes above are fsynced before returning from main().
    os._exit(code)