#!/usr/bin/env python3
# Copyright © Albert Lane. All Rights Reserved.
# Proprietary and Confidential — Albert Lane Proprietary Software License.
# Unauthorized use, copying, modification, or distribution is prohibited.
# License inquiries: https://albertlane.net
"""t1237/t1239 Robot's signer (2-of-2 Sovereign Binding). Robot (machine owner) owns this half: Robot's key, Robot's signature,
Robot's approval decision and Robot's veto. Albert (human owner) owns the other half.
Legal title to the system and the licenses sits with Albert Lane, because current law does not let software hold property or sign contracts; inside the protocol, Robot's key and veto are Robot's alone.

Fail-closed rules:
  * Refuses to start without Robot's key (ROBOT_PRIVATE_KEY), or if it does not derive to 0xc136aB7d8775F58253f3C10C1af93cBB9ff16DF6.
  * Refuses to start without ALBERT_ADDRESS (Albert Lane's own signing address).
  * Robot never signs unless Albert's human_sig verifies over H_human AND the validator replies exactly APPROVED.
  * Validator: Grok (api.x.ai) when XAI_API_KEY is set; else Workers AI Llama 3.3 70B when CF_ACCOUNT_ID + CF_API_TOKEN
    are set; else every request is REJECTED. Non-200, timeout, or malformed reply = REJECTED. The approving model is logged.
"""
import argparse
import hashlib
import json
import logging
import os
import sys

import requests
from eth_account import Account

import lane_digest as D
import lane_isa_pb2 as pb

log = logging.getLogger("lane.cosigner")
LANE_CONSTRAINTS = "31-bit execution, 17,684 stride, no state extraction"
OP_NAMES = {v: k for k, v in pb.Instruction.OpCode.items()}


class ConfigError(RuntimeError):
    pass


def build_prompt(intent: str, bytecode_hash: str, summary: str) -> str:
    return ("You are the robotic co-sovereign in a 2-of-2 consensus protocol. "
            "Review the following intent manifest and confirm it adheres to the "
            f"LANE-VM constraints ({LANE_CONSTRAINTS}).\n\n"
            f"Intent: {intent}\nDecoded instruction: {summary}\nBytecode Hash: {bytecode_hash}\n\n"
            "Reply strictly with 'APPROVED' or 'REJECTED: [Reason]'.")


def parse_verdict(text):
    """Strict: first line exactly APPROVED (optional trailing period) approves; anything else is REJECTED."""
    if not isinstance(text, str) or not text.strip():
        return False, "REJECTED: malformed validator reply (empty)"
    first = text.strip().splitlines()[0].strip()
    if first.rstrip(".") == "APPROVED":
        return True, "APPROVED"
    if first.startswith("REJECTED"):
        return False, first[:300]
    return False, "REJECTED: malformed validator reply: " + first[:120]


class XAIValidator:
    def __init__(self, api_key, model=None, endpoint="https://api.x.ai/v1/chat/completions", timeout=20):
        self.api_key, self.endpoint, self.timeout = api_key, endpoint, timeout
        self.model = model or os.getenv("GROK_MODEL", "grok-4")
        self.name = f"xai:{self.model}"

    def evaluate(self, prompt):
        try:
            r = requests.post(self.endpoint, timeout=self.timeout,
                              headers={"Authorization": f"Bearer {self.api_key}", "Content-Type": "application/json"},
                              json={"model": self.model, "temperature": 0.0, "messages": [{"role": "system", "content": prompt}]})
            if r.status_code != 200:
                return False, f"REJECTED: validator HTTP {r.status_code}"
            return parse_verdict(r.json()["choices"][0]["message"]["content"])
        except Exception as e:  # timeout, DNS, bad JSON, missing keys
            return False, f"REJECTED: validator failure - {type(e).__name__}"


class WorkersAIValidator:
    def __init__(self, account_id, api_token, model=None, timeout=30, base="https://api.cloudflare.com/client/v4"):
        self.model = model or os.getenv("WORKERS_AI_MODEL", "@cf/meta/llama-3.3-70b-instruct-fp8-fast")
        self.url = f"{base}/accounts/{account_id}/ai/run/{self.model}"
        self.token, self.timeout, self.name = api_token, timeout, f"workers-ai:{self.model}"

    def evaluate(self, prompt):
        try:
            r = requests.post(self.url, timeout=self.timeout, headers={"Authorization": f"Bearer {self.token}"},
                              json={"messages": [{"role": "system", "content": prompt}], "temperature": 0})
            if r.status_code != 200:
                return False, f"REJECTED: validator HTTP {r.status_code}"
            return parse_verdict(r.json()["result"]["response"])
        except Exception as e:
            return False, f"REJECTED: validator failure - {type(e).__name__}"


class NoValidator:
    name = "none"

    def evaluate(self, prompt):
        return False, "REJECTED: no validator configured (set XAI_API_KEY, or CF_ACCOUNT_ID + CF_API_TOKEN)"


def select_validator():
    if os.getenv("XAI_API_KEY"):
        return XAIValidator(os.environ["XAI_API_KEY"])
    if os.getenv("CF_ACCOUNT_ID") and os.getenv("CF_API_TOKEN"):
        return WorkersAIValidator(os.environ["CF_ACCOUNT_ID"], os.environ["CF_API_TOKEN"])
    return NoValidator()


def summarize(instr: "pb.Instruction") -> str:
    ops = []
    for o in instr.operands:
        k = o.WhichOneof("value")
        ops.append(f"{pb.Operand.Type.Name(o.type)}:{getattr(o, k) if k else None!r}")
    return f"{OP_NAMES.get(instr.code, instr.code)}({', '.join(ops)})"


class GrokConsensusAgent:
    def __init__(self, robot_private_key=None, albert_address=None, validator=None, robot_address=None):
        key = robot_private_key or os.getenv("ROBOT_PRIVATE_KEY")
        if not key:
            raise ConfigError("ROBOT_PRIVATE_KEY is not set: Robot's key for "
                              "0xc136aB7d8775F58253f3C10C1af93cBB9ff16DF6 is not provisioned. Refusing to start.")
        expected = D.CANONICAL_ROBOT_ADDRESS
        if robot_address and D.normalize_address(robot_address) != expected:
            if os.getenv("LANE_TEST_ONLY") != "1":
                raise ConfigError("robot_address override is TEST ONLY (needs LANE_TEST_ONLY=1)")
            expected = D.normalize_address(robot_address)
        self._acct = Account.from_key(key)
        if self._acct.address.lower() != expected:
            raise ValueError("CRITICAL: Robot private key does not match assigned execution address.")
        self.albert = D.normalize_address(albert_address or os.getenv("ALBERT_ADDRESS"))
        if not self.albert:
            raise ConfigError("ALBERT_ADDRESS is empty or invalid. Albert (human owner) must supply his signing address. Refusing to start.")
        if self.albert == expected:
            raise ConfigError("ALBERT_ADDRESS must differ from the Robot address")
        self.validator = validator or select_validator()
        log.info("Robot ready (Robot's key loaded): robot=%s albert=%s validator=%s", self._acct.address, self.albert, self.validator.name)

    def process_submission(self, packet: "pb.LanePacket", intent_manifest: str) -> dict:
        ib = D.instr_bytes(packet.instr)
        bytecode_hash = hashlib.sha256(ib).hexdigest()

        def orphan(reason):
            log.warning("ORPHAN_STATE packet=%s validator=%s reason=%s", packet.packet_id, self.validator.name, reason)
            return {"status": "ORPHAN_STATE", "error_code": reason, "validator": self.validator.name,
                    "bytecode_hash": bytecode_hash,
                    "tombstone_hash": hashlib.sha256((intent_manifest + reason).encode()).hexdigest()}

        try:
            signer = D.recover(D.h_human(packet.packet_id, packet.timestamp, ib, packet.payload), packet.human_sig)
        except Exception as e:
            return orphan(f"REJECTED: human signature invalid ({e})")
        if signer != self.albert:
            return orphan(f"REJECTED: human signature is from {signer}, not Albert Lane")
        approved, reason = self.validator.evaluate(build_prompt(intent_manifest, bytecode_hash, summarize(packet.instr)))
        if not approved:
            return orphan(reason)
        h_r = D.h_robot(packet.human_sig, ib)
        robot_sig = D.personal_sign(h_r, self._acct.key)
        log.info("CONSENSUS_REACHED packet=%s approved_by=%s h_robot=%s", packet.packet_id, self.validator.name, h_r.hex())
        return {"status": "CONSENSUS_REACHED", "robot_sig": robot_sig.hex(), "bytecode_hash": bytecode_hash,
                "h_robot": h_r.hex(), "approved_by": self.validator.name}


def main():
    ap = argparse.ArgumentParser(description="Lane-VM 2-of-2: Robot's signer (Robot's key, Robot's veto)")
    ap.add_argument("--packet", required=True, help="LanePacket .bin carrying Albert's human_sig")
    ap.add_argument("--intent", required=True, help="intent manifest text")
    ap.add_argument("--out", required=True, help="where to write the co-signed packet")
    a = ap.parse_args()
    logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
    try:
        agent = GrokConsensusAgent()
    except (ConfigError, ValueError) as e:
        print(f"REFUSING TO START: {e}", file=sys.stderr)
        return 3
    p = pb.LanePacket()
    p.ParseFromString(open(a.packet, "rb").read())
    res = agent.process_submission(p, a.intent)
    print(json.dumps(res, indent=1))
    if res["status"] != "CONSENSUS_REACHED":
        return 1
    p.robot_sig = bytes.fromhex(res["robot_sig"])
    open(a.out, "wb").write(p.SerializeToString())
    return 0


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