Onboard a machine on Solana

"""Onboard one machine terminal-first on Solana from the public Python API, resuming from a journal.

The Python twin of the JS SDK's ``onboard-solana.mjs``: the same stages, the same journal
discipline, the same output lines. It does whatever the SDK's
status read says is next until the machine is ``complete``, so a second invocation after any
stop continues where the first ended and never re-sends::

    python onboard_solana.py --dir <machine dir> --serial <serial> --pay-in PEAQ \\
        --owner <base58 owner> --operator <0x operator> --owner-wallet <OWS wallet> \\
        --manufacturer <base58> --machine-type <type> [--dry-run] [--wait-minutes 60] \\
        [--tier 1] [--compute-unit-price-micro-lamports 10000]

The machine dir holds ``.env`` (the CLI's keys), ``did.json`` (the DID document) and, after the
first write, ``peaqos-py.log`` (one serialized journal row per line). Every Solana write is signed
by the OWS vault wallet named by ``--owner-wallet``, unlocked with ``OWS_PASSPHRASE``. Nothing on
peaq is signed: ``--owner`` must already be registered to ``--operator`` (``[1/7]``, once per owner
wallet: ``peaqos activate --phase registration`` or ``register_solana_operator``).

Reads are keyless. ``--dry-run`` stops after them: the registration, the plan, the quote and the
link push quote (route, fee). No endpoint URL is ever printed.
"""

from __future__ import annotations

import argparse
import json
import os
import sys
import time
from dataclasses import replace
from datetime import datetime, timezone
from pathlib import Path
from typing import Any

from peaq_os_sdk import (
    ActivateSolanaMachineParams,
    ActivateSolanaTerminalParams,
    Address,
    PeaqosClient,
    PreviewSolanaMachineActivationParams,
    ServiceEndpointInput,
    SolanaMachineActivationStatusParams,
    SolanaMachineIdentity,
    SolanaOnboardingInputs,
    SolanaOnboardingRecoveryParams,
    SolanaPubkey,
    SolanaVerificationMethodInput,
    Tokenomics20Config,
    TokenomicsPendingTransactionError,
    normalize_solana_controller,
)
from peaq_os_sdk.tokenomics.solana_onboarding_attempts import (
    SolanaOnboardingAttempt,
    SolanaOnboardingAttemptRecord,
    TerminalActivationAttempt,
    decode_solana_onboarding_attempt,
    serialize_solana_onboarding_attempt,
)

# ---- the caps and budgets (the owner, operator, manufacturer and machine type are arguments) ----
MAX_NATIVE_FEE_LAMPORTS = 50_000
MAX_NATIVE_RENT_LAMPORTS = 6_000_000
COMPUTE_UNIT_LIMIT = 600_000  # the mint's budget; the link push sizes its own by route
DEFAULT_COMPUTE_UNIT_PRICE = 10_000
"""Micro-lamports per compute unit. At 0, Solana's public endpoint can drop a write under load,
and a journaled write is reconciled, never re-sent; 10,000 costs 6,000 lamports on the mint and
1,200 on the push."""
CALL_TIMEOUT_SECONDS = 120.0
POLL_SECONDS = 15
HISTORY_TIMEOUT_SECONDS = 600.0
"""The history reads both chains; right after a run the public Solana endpoint still
rate-limits, and the SDK's 120 s default can leave sources unread. The CLI's `machine history`
allows up to 900."""
STAGE = {
    "registration": 1,
    "request": 2,
    "credit": 3,
    "finalise": 4,
    "bond": 5,
    "native_onboarding": 6,
    "linkage": 7,
}
TERMINAL_WAIT_SECONDS = {"credit": 180.0, "bond": 360.0}
CLOSED_REQUEST = ("cancelled_or_refunded", "closed")
"""Request states the terminal reports once a request can no longer settle."""


def log(line: str) -> None:
    print(f"{datetime.now(timezone.utc).strftime('%H:%M:%S')}  {line}", flush=True)


def short(value: object) -> str:
    text = str(value)
    return f"{text[:10]}…{text[-5:]}" if len(text) > 20 else text


def header(phase: str, text: str) -> None:
    log(f"[{STAGE[phase]}/7] {phase}: {text}")


def read_env(path: Path) -> dict[str, str]:
    """The machine dir's ``.env``, which wins over the shell (an exported endpoint once did not)."""
    values: dict[str, str] = {}
    for raw in path.read_text(encoding="utf-8").splitlines():
        line = raw.strip()
        if not line or line.startswith("#") or "=" not in line:
            continue
        key, value = line.split("=", 1)
        values[key.strip()] = value.strip().strip('"').strip("'")
    return values


def build_inputs(args: argparse.Namespace, did: dict[str, Any]) -> SolanaOnboardingInputs:
    """The machine's native inputs: never change them between invocations of one machine."""
    return SolanaOnboardingInputs(
        solana_owner=SolanaPubkey(args.owner),
        controller=normalize_solana_controller(),  # owner-only, as the CLI
        manufacturer=SolanaPubkey(args.manufacturer),
        machine_type=args.machine_type,
        credential_subject=args.serial.encode("utf-8"),
        verification_methods=tuple(
            SolanaVerificationMethodInput(
                id=m["id"],
                method_type=m["methodType"],
                controller=SolanaPubkey(m["controller"]),
                public_key_multibase=m["publicKeyMultibase"],
            )
            for m in did["verificationMethods"]
        ),
        authentication=tuple(int(i) for i in did["authentication"]),
        service_endpoints=tuple(
            ServiceEndpointInput(
                id=s["id"], service_type=s["serviceType"], service_endpoint=s["serviceEndpoint"]
            )
            for s in did["serviceEndpoints"]
        ),
        max_native_fee_lamports=MAX_NATIVE_FEE_LAMPORTS,
        max_native_rent_lamports=MAX_NATIVE_RENT_LAMPORTS,
        evm_operator=Address(args.operator),
    )


class Journal:
    """``peaqos-py.log``: one serialized attempt per line, written before broadcast."""

    def __init__(self, path: Path, client: PeaqosClient) -> None:
        self.path = path
        self.client = client

    def records(self) -> tuple[SolanaOnboardingAttemptRecord, ...]:
        if not self.path.exists():
            return ()
        lines = self.path.read_text(encoding="utf-8").splitlines()
        return tuple(decode_solana_onboarding_attempt(line) for line in lines if line.strip())

    def terminal(self) -> tuple[TerminalActivationAttempt, ...]:
        return tuple(
            r.attempt for r in self.records() if isinstance(r.attempt, TerminalActivationAttempt)
        )

    def for_request(self, state: Any) -> tuple[TerminalActivationAttempt, ...]:
        """The terminal rows a new request may reconcile: none of a request already closed.

        A request the terminal has cancelled, refunded or closed is history; the next request
        opens the following nonce. Handed on, its row would be reconciled again instead of a new
        request being opened.
        """
        rows = self.terminal()
        terminal = state.observations.terminal if state.observations is not None else None
        last = terminal.terminal.last_nonce if terminal is not None else None
        if last is None or terminal is None or terminal.request_state not in CLOSED_REQUEST:
            return rows
        current = tuple(r for r in rows if r.nonce > last)
        if len(current) < len(rows):
            log(
                f"      journal: request nonce {last} is {terminal.request_state} on the "
                f"terminal; its {len(rows) - len(current)} row(s) stay in the journal and the "
                f"request opens nonce {last + 1}"
            )
        return current

    def persist(self, attempt: SolanaOnboardingAttempt) -> None:
        deployment = self.client.tokenomics20
        assert deployment is not None
        with self.path.open("a", encoding="utf-8") as handle:
            handle.write(serialize_solana_onboarding_attempt(attempt, deployment) + "\n")
            handle.flush()
            os.fsync(handle.fileno())
        reference = getattr(attempt, "signature", None) or getattr(
            attempt, "transaction_hash", None
        )
        log(f"      {attempt.phase} signed and journaled: {short(reference)}")


class Runner:
    def __init__(self, args: argparse.Namespace) -> None:
        self.dir = Path(args.dir)
        self.serial: str = args.serial
        self.pay_in: str = args.pay_in
        self.tier: int = args.tier
        self.wallet: str = args.owner_wallet
        self.owner = SolanaPubkey(args.owner)
        self.operator = Address(args.operator)
        self.price: int = args.compute_unit_price_micro_lamports
        self.deadline = time.monotonic() + args.wait_minutes * 60
        env = read_env(self.dir / ".env")

        def need(key: str) -> str:
            value = env.get(key)
            if not value:
                raise SystemExit(f"{key} is not set in {self.dir / '.env'}")
            return value

        self.deployment_id = need("TOKENOMICS_DEPLOYMENT_ID")
        # Keyless: no private key. Nothing on peaq is signed by this script.
        self.client = PeaqosClient(
            rpc_url=need("PEAQOS_RPC_URL"),
            identity_registry=Address(need("IDENTITY_REGISTRY_ADDRESS")),
            identity_staking=Address(need("IDENTITY_STAKING_ADDRESS")),
            event_registry=Address(need("EVENT_REGISTRY_ADDRESS")),
            machine_nft=Address(need("MACHINE_NFT_ADDRESS")),
            did_registry=Address(need("DID_REGISTRY_ADDRESS")),
            batch_precompile=Address(need("BATCH_PRECOMPILE_ADDRESS")),
            tokenomics20=Tokenomics20Config(
                deployment_id=self.deployment_id,
                creation_home="solana",
                solana_rpc_url=need("PEAQOS_SVM_RPC_URL"),
            ),
        )
        did = json.loads((self.dir / "did.json").read_text(encoding="utf-8"))
        self.identity = SolanaMachineIdentity(args.machine_type, self.serial.encode("utf-8"))
        self.inputs = build_inputs(args, did)
        self.preview = PreviewSolanaMachineActivationParams(
            inputs=self.inputs,
            compute_unit_limit=COMPUTE_UNIT_LIMIT,
            compute_unit_price_micro_lamports=self.price,
        )
        self.journal = Journal(self.dir / "peaqos-py.log", self.client)
        self.signer: Any = None
        self.machine_id = 0

    # ---- reads ---------------------------------------------------------------------------
    def budget_left(self) -> int:
        return max(0, int(self.deadline - time.monotonic()))

    def registration(self) -> None:
        state = self.client.get_solana_operator_registration(
            self.owner, peaq_operator=self.operator
        )
        bound = state.bound_operator
        if state.state != "bound" or bound is None or bound.lower() != self.operator.lower():
            raise SystemExit(
                f"registration is {state.state} ({bound}); this script signs nothing on peaq"
            )
        header("registration", f"reconciled (bound to {self.operator}; nothing was signed)")

    def quote_request(self) -> int:
        quote = self.client.quote_solana_activation(
            self.identity,
            tier=self.tier,
            pay_in=self.pay_in,  # type: ignore[arg-type]
            wallet=self.owner,
        )
        log(
            f"quote: bond {quote.bond.base} PEAQ base units · suggested escrow {quote.max_in} "
            f"({self.pay_in} base units) · wallet balance {quote.wallet_balance} · "
            f"stamp age {quote.price_age_secs} s"
        )
        return quote.max_in

    def status(self) -> Any:
        records = self.journal.records()
        return self.client.get_machine_activation_state(
            self.machine_id,
            solana=SolanaMachineActivationStatusParams(inputs=self.inputs, records=records or None),
        )

    def link_quote(self) -> Any:
        quote = self.client.quote_solana_link_push(self.machine_id, payer=self.owner)
        log(
            f"      link push quote: {quote.status} · transport {quote.transport} · messaging fee "
            f"{quote.native_fee_lamports} · rent {quote.push_account_rent_lamports} · "
            f"compute {quote.compute_unit_limit} · total {quote.total_lamports} lamports"
        )
        return quote

    def get_signer(self) -> Any:
        if self.signer is None:
            plan = self.client.preview_machine_activation(self.preview)
            required = self.client.get_solana_onboarding_signer_requirements(
                self.inputs, plan, "native_onboarding", candidate_address=self.owner
            )
            self.signer = self.client.solana_signer_from_wallet(
                self.wallet, os.environ.get("OWS_PASSPHRASE"), chain_id=required.chain_id
            )
            self.client.validate_solana_onboarding_signer(
                self.inputs, plan, "native_onboarding", solana_signer=self.signer
            )
            log(
                f"      signer {self.wallet} ({short(self.owner)}) unlocked for chain "
                f"{required.chain_id}"
            )
        return self.signer

    # ---- writes --------------------------------------------------------------------------
    def terminal(
        self,
        phase: str,
        max_in: int | None = None,
        records: tuple[TerminalActivationAttempt, ...] | None = None,
    ) -> None:
        started = time.monotonic()
        header(phase, "…")
        observe = phase in TERMINAL_WAIT_SECONDS
        params = ActivateSolanaTerminalParams(
            self.identity,
            phase,  # type: ignore[arg-type]
            tier=self.tier,
            pay_in=self.pay_in,  # type: ignore[arg-type]
            solana_signer=None if observe else self.get_signer(),
            max_in_amount=max_in,
            max_native_fee_lamports=MAX_NATIVE_FEE_LAMPORTS,
            max_native_rent_lamports=MAX_NATIVE_RENT_LAMPORTS,
            compute_unit_price_micro_lamports=self.price,
            records=self.journal.terminal() if records is None else records,
            wait_seconds=(
                min(float(self.budget_left()), TERMINAL_WAIT_SECONDS[phase]) if observe else None
            ),
            timeout_seconds=CALL_TIMEOUT_SECONDS,
            on_transaction_submitted=None if observe else self.journal.persist,
        )
        result = self.client.activate_machine(params)
        submitted = result.confirmation.submitted if result.confirmation else result.submitted
        reference = getattr(submitted, "signature", None)
        request = result.request.state if result.request is not None else "-"
        header(
            phase,
            f"{result.status}{f' ({result.exit})' if result.exit else ''}"
            f"{f' {short(reference)}' if reference else ''} · request {request} "
            f"nonce {result.nonce} · {int(time.monotonic() - started)} s",
        )
        swap = result.settlement
        if swap is not None:
            log(
                f"      settlement: swap simulated {swap.simulated_in_amount} · cap sent "
                f"{swap.max_in_amount} · escrow {swap.escrow_amount} ({self.pay_in} base units) · "
                f"slippage {swap.slippage_bps} bps · lock in table {swap.lock_in_table}"
            )
        if result.status == "exited":
            raise SystemExit(f"{phase} exited: {result.exit} {result.advice or ''}")

    def native(self, phase: str) -> None:
        header(phase, "…")
        recovery = self.client.get_machine_activation_state(
            self.machine_id,
            solana=SolanaOnboardingRecoveryParams(self.inputs, self.journal.records()),
        )
        if recovery.error is not None:
            raise recovery.error
        plan = self.client.preview_machine_activation(self.preview)
        signer = self.get_signer()
        if phase == "linkage":
            self.link_quote()
        params = ActivateSolanaMachineParams(
            inputs=self.inputs,
            accepted_plan=plan,
            phase=phase,  # type: ignore[arg-type]
            expected_machine_id=self.machine_id,
            records=recovery.history.records,
            compute_unit_limit=COMPUTE_UNIT_LIMIT,
            compute_unit_price_micro_lamports=self.price,
            timeout_seconds=CALL_TIMEOUT_SECONDS,
            on_transaction_submitted=self.journal.persist,
        )
        # No max_link_push_fee_lamports: the Trust Validator node's route has no fee. On a
        # LayerZero route the SDK refuses before signing and names the cap to pass.
        params = replace(params, solana_signer=signer)
        try:
            result = self.client.activate_machine(params)
        except TokenomicsPendingTransactionError as error:
            log(f"      pending after signing: {error}; the row is journaled, rerun to reconcile")
            return
        confirmation = result.confirmation
        signature = getattr(getattr(confirmation, "submitted", None), "signature", None)
        header(
            phase,
            f"{result.status} ({result.reason}, {result.origin})"
            f"{f' {short(signature)}' if signature else ''}",
        )
        if phase == "linkage" and confirmation is not None:
            pushed = getattr(confirmation, "submitted", None)
            transport = getattr(pushed, "transport", None) or "layerzero"
            tv_nonce = getattr(pushed, "tv_nonce", None) or "-"
            log(
                f"      transport {transport} · TV nonce {tv_nonce} · "
                f"seq {getattr(pushed, 'sequence', '-')}"
            )
        if result.error is not None:
            log(f"      error: {result.error.code}: {result.error}")
        if result.status != "confirmed":
            raise SystemExit(f"{phase} ended {result.status}: {result.reason}")

    # ---- the run -------------------------------------------------------------------------
    def dry_run(self, plan: Any) -> None:
        try:
            self.quote_request()
        except Exception as error:  # a refusal is the answer a dry run reports
            log(f"quote: {getattr(error, 'code', type(error).__name__)}: {error}")
        for stage in plan.stages:
            log(
                f"  stage {stage.phase}: {stage.action} (signer {stage.signer or '-'}, "
                f"unresolved {list(stage.unresolved)})"
            )
        try:
            self.link_quote()
        except Exception as error:
            log(
                f"      link push quote: {getattr(error, 'code', type(error).__name__)}: {error} "
                "(expected before the mint: MACHINE_NOT_ONBOARDED)"
            )
        step = self.status().next_step
        log(f"next step: {step.phase} {step.action}")

    def run(self, dry_run: bool) -> None:
        mode = "dry run" if dry_run else "write"
        log(
            f"Solana onboarding (Python) · {self.deployment_id} · serial {self.serial} · "
            f"{self.pay_in} · {mode}"
        )
        self.registration()
        # The machine ID comes from the plan (keyless), never from the quote: once the request
        # is bonded the terminal rightly refuses a fresh quote, and a resume still needs the ID.
        plan = self.client.preview_machine_activation(self.preview)
        self.machine_id = plan.binding.machine_id
        log(f"machine {self.machine_id}")
        log(f"journal: {len(self.journal.records())} row(s) in {self.journal.path.name}")
        if dry_run:
            self.dry_run(plan)
            return
        for _ in range(60):
            state = self.status()
            if state.error is not None:
                raise state.error
            step = state.next_step
            if step.phase == "complete":
                break
            if step.phase == "cancel_request":
                raise SystemExit(
                    "the request expired before the node credited it; nothing was settled. "
                    "Wait for the node to cancel it (the escrow and rents come back), or cancel "
                    "it yourself with `peaqos activate ... --phase cancel_request`; then rerun"
                )
            if step.phase == "registration":
                raise SystemExit("the status read wants a registration; this script signs none")
            if step.phase == "request":
                self.terminal(
                    "request", max_in=self.quote_request(), records=self.journal.for_request(state)
                )
            elif step.phase in ("credit", "bond", "finalise"):
                self.terminal(step.phase)
            elif step.phase == "native_onboarding":
                self.native("native_onboarding")
            elif step.phase == "linkage" and step.action == "observe":
                if self.budget_left() == 0:
                    raise SystemExit(
                        "PENDING: the link is not applied within --wait-minutes; rerun to keep "
                        "waiting (nothing is resent)"
                    )
                log(
                    "      waiting for the node to deliver the link push on peaq "
                    f"({self.budget_left()} s left)"
                )
                time.sleep(POLL_SECONDS)
            elif step.phase == "linkage":
                self.native("linkage")
            else:
                raise SystemExit(f"unexpected next step {step.phase} {step.action}")
        else:
            raise SystemExit("PENDING: too many steps without completing; rerun to continue")
        self.finish()

    def finish(self) -> None:
        recovery = self.client.get_machine_activation_state(
            self.machine_id,
            solana=SolanaOnboardingRecoveryParams(
                self.inputs, self.journal.records(), include_linkage=True
            ),
        )
        relays = recovery.tv_relays
        log(
            f"complete: link applied; messages {len(recovery.messages)} (LayerZero), "
            f"tv relays {len(relays)}"
        )
        for relay in relays:
            log(
                f"      tv relay: {short(relay.signature)} dst {relay.dst_eid} nonce "
                f"{relay.nonce} applied {relay.applied}"
            )
        history = self.client.get_solana_onboarding_history(
            self.machine_id, timeout_seconds=HISTORY_TIMEOUT_SECONDS
        )
        log(f"history: {len(history.rows)} rows, unread {list(history.unread)}")
        for row in history.rows:
            log(
                f"  {row.step:<18} {row.actor:<9} {row.chain:<6} {short(row.reference)}"
                f"{' · reused' if row.reused else ''}"
            )
        for total in history.totals:
            log(
                f"  total {total.chain} {short(total.payer)} lamports {total.lamports_delta} "
                f"gas {total.gas_wei} tokens {dict(total.token_deltas)}"
            )
        log(f"done · machine {self.machine_id}")


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0] if __doc__ else None)
    parser.add_argument("--dir", required=True, help="the machine dir (.env, did.json, journal)")
    parser.add_argument("--serial", required=True, help="the credential subject, as text")
    parser.add_argument("--pay-in", required=True, choices=("PEAQ", "USDC"))
    parser.add_argument("--owner", required=True, help="the Solana owner's base58 public key")
    parser.add_argument("--operator", required=True, help="the peaq operator it is registered to")
    parser.add_argument("--owner-wallet", required=True, help="the owner's OWS vault wallet")
    parser.add_argument("--manufacturer", required=True, help="base58, recorded unverified")
    parser.add_argument("--machine-type", required=True, help="with the serial, fixes the ID")
    parser.add_argument("--dry-run", action="store_true", help="keyless reads only")
    parser.add_argument("--wait-minutes", type=float, default=60.0)
    parser.add_argument("--tier", type=int, default=1)
    parser.add_argument(
        "--compute-unit-price-micro-lamports", type=int, default=DEFAULT_COMPUTE_UNIT_PRICE
    )
    args = parser.parse_args(argv)
    Runner(args).run(args.dry_run)
    return 0


if __name__ == "__main__":
    sys.exit(main())