"""Prepare offline; run/resume one durable state-rollout request (Python 3.10+).

prepare --input request.json --run-dir NEW_DIRECTORY
run --run-dir EXISTING_DIRECTORY

Input, predictions, receipts and an API-key SHA256 fingerprint are persisted
locally, without the API key. Resume must use the ORIGINAL Key: billing
idempotency is scoped to that Key, even within the same account.
Use a trusted private directory. The SHA detects consistency, not authenticity.
run uses KINEJING_API_KEY and the SDK's fixed production origin; it can consume
trial/paid usage. It never creates a fresh identity when resuming. A stale
active.lock requires manual removal ONLY after ensuring the old process stopped.
Retain request.json after any failure. Inspect result.json.partial locally if
present; do not delete evidence or change identity to retry an unknown outcome.
File fsync/exclusive publication cannot guarantee every OS/disk/power-loss case;
directory fsync is available on POSIX, while Windows durability follows its OS.
"""
import argparse
import hashlib
import json
import math
import os
from pathlib import Path
import re
import stat
import sys
import uuid

if not Path(__file__).with_name("kinejing_client.py").is_file():
    sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "sdk" / "python"))
from kinejing_client import KineJingClient, KineJingError, MODEL, _validate_states

TOOL_VERSION = "1.0.0"
INPUT_LIMIT = 65536
RESULT_LIMIT = 16 * 1024 * 1024
PAYLOAD_KEYS = {"model", "state_schema", "action_schema", "context_states", "actions"}
RECORD_KEYS = {"schema_version", "tool_version", "payload", "payload_sha256", "idempotency_key"}


class LocalError(Exception):
    """Fixed diagnostic code only; never attach input, credentials or raw errors."""


class Parser(argparse.ArgumentParser):
    def error(self, message):
        raise LocalError("invalid_arguments")


def canonical(value):
    return json.dumps(value, allow_nan=False, separators=(",", ":")).encode("utf-8")


def _pairs(pairs):
    out = {}
    for key, value in pairs:
        if key in out:
            raise LocalError("duplicate_json_key")
        out[key] = value
    return out


def _constant(value):
    raise LocalError("nonfinite_json_number")


def read_json(path, limit):
    if not stat.S_ISREG(path.lstat().st_mode):
        raise LocalError("regular_file_required")
    with path.open("rb") as stream:
        raw = stream.read(limit + 1)
    if len(raw) > limit:
        raise LocalError("json_file_too_large")
    return json.loads(raw.decode("utf-8"), object_pairs_hook=_pairs, parse_constant=_constant)


def payload(value):
    if not isinstance(value, dict) or set(value) != PAYLOAD_KEYS:
        raise LocalError("invalid_request_fields")
    if (value["model"] != MODEL or value["state_schema"] != "pusht-512xy-angle-rad-v1"
            or value["action_schema"] != "pusht-absolute-xy-10hz-v1"):
        raise LocalError("unsupported_request_contract")
    _validate_states(value["context_states"], value["actions"])
    # Same key order and JSON encoding as SDK 0.1.2's outgoing request body.
    return {key: value[key] for key in ("model", "state_schema", "action_schema", "context_states", "actions")}


def _sync_directory(directory):
    if os.name == "posix":
        fd = os.open(directory, os.O_RDONLY | os.O_DIRECTORY)
        try:
            os.fsync(fd)
        finally:
            os.close(fd)


def publish(path, value):
    partial = path.with_name(path.name + ".partial")
    if os.path.lexists(path) or os.path.lexists(partial):
        raise LocalError("local_output_conflict")
    fd = os.open(partial, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600)
    with os.fdopen(fd, "wb") as stream:
        stream.write(canonical(value) + b"\n")
        stream.flush()
        os.fsync(stream.fileno())
    # Atomic no-replace publication; unsupported hard links fail closed.
    os.link(partial, path)
    os.unlink(partial)
    _sync_directory(path.parent)


def request_record(directory):
    if os.path.lexists(directory / "request.json.partial"):
        raise LocalError("request_partial_requires_local_review")
    record = read_json(directory / "request.json", INPUT_LIMIT * 2)
    if (not isinstance(record, dict) or set(record) != RECORD_KEYS
            or type(record["schema_version"]) is not int or record["schema_version"] != 1
            or record["tool_version"] != TOOL_VERSION):
        raise LocalError("invalid_request_record")
    request = payload(record["payload"])
    if (not isinstance(record["idempotency_key"], str)
            or not re.fullmatch("[0-9a-f]{32}", record["idempotency_key"])
            or record["payload_sha256"] != hashlib.sha256(canonical(request)).hexdigest()):
        raise LocalError("request_identity_mismatch")
    return record


def _check_result(result, record, binding):
    expected = {"schema_version", "tool_version", "payload_sha256", "idempotency_key", "origin", "api_key_sha256", "response"}
    if (not isinstance(result, dict) or set(result) != expected
            or type(result["schema_version"]) is not int or result["schema_version"] != 1
            or result["tool_version"] != TOOL_VERSION
            or any(result[key] != binding[key] for key in expected - {"response"})):
        raise LocalError("result_identity_mismatch")
    response = result["response"]
    states = response.get("predicted_states") if isinstance(response, dict) else None
    if (not isinstance(response, dict) or response.get("model") != MODEL
            or not isinstance(states, list) or len(states) != len(record["payload"]["actions"])):
        raise LocalError("invalid_saved_result")
    for row in states:
        if (not isinstance(row, list) or len(row) != 5
                or any(type(x) not in (int, float) or not math.isfinite(x) for x in row)):
            raise LocalError("invalid_saved_result")
    return response


def summary(record, *, cached=None):
    out = {"status": "prepared" if cached is None else "local_result" if cached else "result_saved",
           "model": MODEL, "idempotency_key": record["idempotency_key"],
           "payload_sha256": record["payload_sha256"], "predicted_steps": len(record["payload"]["actions"])}
    print(json.dumps(out, ensure_ascii=True))


def execution_record(directory, record):
    path = directory / "execution.json"
    if os.path.lexists(path.with_name(path.name + ".partial")):
        raise LocalError("execution_partial_requires_local_review")
    saved = read_json(path, 4096)
    expected = RECORD_KEYS - {"payload"} | {"origin", "api_key_sha256"}
    if (not isinstance(saved, dict) or set(saved) != expected
            or type(saved["schema_version"]) is not int
            or any(saved[key] != record[key] for key in RECORD_KEYS - {"payload"})
            or saved["origin"] != "https://kineworld.com" or not isinstance(saved["api_key_sha256"], str)
            or not re.fullmatch("[0-9a-f]{64}", saved["api_key_sha256"])):
        raise LocalError("invalid_execution_identity")
    return saved


def bind_key(directory, record, client):
    binding = {key: record[key] for key in RECORD_KEYS - {"payload"}}
    binding.update(origin="https://kineworld.com", api_key_sha256=hashlib.sha256(client._api_key.encode("ascii")).hexdigest())
    if client.base_url != binding["origin"]:
        raise LocalError("fixed_production_origin_required")
    path = directory / "execution.json"
    if os.path.lexists(path):
        if execution_record(directory, record) != binding:
            raise LocalError("original_api_key_and_execution_identity_required")
    else:
        publish(path, binding)
    return binding


def prepare(input_path, directory):
    request = payload(read_json(input_path, INPUT_LIMIT))
    record = {"schema_version": 1, "tool_version": TOOL_VERSION, "payload": request,
              "payload_sha256": hashlib.sha256(canonical(request)).hexdigest(), "idempotency_key": uuid.uuid4().hex}
    directory.mkdir(mode=0o700)  # Exclusive; never reuse or repair an old directory.
    _sync_directory(directory.parent)
    publish(directory / "request.json", record)
    summary(record)


def run(directory):
    if not stat.S_ISDIR(directory.lstat().st_mode):
        raise LocalError("regular_directory_required")
    lock = directory / "active.lock"
    try:
        fd = os.open(lock, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o600)
    except FileExistsError:
        raise LocalError("run_locked_check_old_process") from None
    try:
        with os.fdopen(fd, "w", encoding="ascii") as stream:
            stream.write(str(os.getpid()) + "\n"); stream.flush(); os.fsync(stream.fileno())
        record = request_record(directory)
        result_path = directory / "result.json"
        if os.path.lexists(result_path.with_name("result.json.partial")):
            raise LocalError("result_partial_requires_local_review")
        if os.path.lexists(result_path):
            _check_result(read_json(result_path, RESULT_LIMIT), record, execution_record(directory, record))
            summary(record, cached=True)
            return
        request = record["payload"]
        try:
            client = KineJingClient()  # Only the SDK reads/validates the environment Key.
            binding = bind_key(directory, record, client)  # Durable before any HTTP attempt.
            _sync_directory(directory)  # Retry a previously failed publication barrier.
            response = client.rollout(request["context_states"], request["actions"],
                                      idempotency_key=record["idempotency_key"])
        except KineJingError as error:
            status = error.status if type(error.status) is int else None
            print(json.dumps({"error": "api_outcome_not_saved", "http_status": status,
                              "idempotency_key": record["idempotency_key"], "next_step": "resume_same_run_directory"}), file=sys.stderr)
            return 1
        result = dict(binding)
        result["response"] = response
        _check_result(result, record, binding)
        publish(result_path, result)
        summary(record, cached=False)
    finally:
        os.unlink(lock)


def main(argv=None):
    try:
        parser = Parser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
        commands = parser.add_subparsers(dest="command", required=True, parser_class=Parser)
        p = commands.add_parser("prepare", help="Validate and durably prepare offline; no Key/network")
        p.add_argument("--input", required=True, type=Path); p.add_argument("--run-dir", required=True, type=Path)
        r = commands.add_parser("run", help="Send or resume exactly this request; may consume usage")
        r.add_argument("--run-dir", required=True, type=Path)
        args = parser.parse_args(argv)
        return prepare(args.input, args.run_dir) if args.command == "prepare" else run(args.run_dir)
    except LocalError as error:
        print(json.dumps({"error": str(error), "next_step": "retain_run_directory_and_review_locally"}), file=sys.stderr)
    except (OSError, ValueError, TypeError, OverflowError, RecursionError):
        print(json.dumps({"error": "local_validation_or_persistence_failed", "next_step": "retain_run_directory_and_review_locally"}), file=sys.stderr)
    else:
        return 0
    return 1


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