1"""
2Hyperliquid L1 Action Signing — phantom agent EIP-712 pattern.
3 
4Hyperliquid trading actions (orders, cancels, leverage, modify) require an EIP-712
5"Agent" signature over a keccak256 hash of the msgpack-serialized action.
6 
7Since we don't have direct private key access (Privy server wallet), we delegate
8EIP-712 signing to the wallet service's /agent/sign-typed-data endpoint.
9 
10Two signing schemes:
111. L1 Action (orders, cancels, leverage) — msgpack → keccak → Agent struct
122. User-Signed (USDC transfers, withdrawals) — direct EIP-712 over action fields
13 
14The connectionId field (bytes32) encoding must match what the Privy wallet service
15expects. We try multiple encodings and verify each with local ecrecover.
16"""
17 
18import logging
19from typing import Optional
20 
21import msgpack
22from eth_utils import keccak
23 
24from core.wallet_runtime import wallet_request as _wallet_request
25 
26logger = logging.getLogger(__name__)
27 
28# Hyperliquid chain constants
29MAINNET_SOURCE = "a"
30 
31# EIP-712 domains
32L1_DOMAIN = {
33    "name": "Exchange",
34    "version": "1",
35    "chainId": 1337,
36    "verifyingContract": "0x0000000000000000000000000000000000000000",
37}
38 
39USER_SIGNED_DOMAIN = {
40    "name": "HyperliquidSignTransaction",
41    "version": "1",
42    "chainId": 42161,  # Arbitrum mainnet
43    "verifyingContract": "0x0000000000000000000000000000000000000000",
44}
45 
46# EIP-712 types for Agent struct
47AGENT_TYPES = {
48    "Agent": [
49        {"name": "source", "type": "string"},
50        {"name": "connectionId", "type": "bytes32"},
51    ]
52}
53 
54 
55# ── Wallet address cache ─────────────────────────────────────────────────
56 
57_cached_wallet_address: Optional[str] = None
58 
59 
60async def _get_wallet_address() -> Optional[str]:
61    """Get the Privy wallet address for signature verification (cached)."""
62    global _cached_wallet_address
63    if _cached_wallet_address:
64        return _cached_wallet_address
65 
66    try:
67        from core.wallet_runtime import is_fly_machine as _is_fly_machine
68        if not _is_fly_machine():
69            return None
70 
71        data = await _wallet_request("GET", "/agent/wallet")
72        wallets = data if isinstance(data, list) else data.get("wallets", [])
73        for w in wallets:
74            if w.get("chain_type") == "ethereum":
75                _cached_wallet_address = w["wallet_address"].lower()
76                return _cached_wallet_address
77    except Exception as e:
78        logger.debug(f"Could not fetch wallet address for verification: {e}")
79 
80    return None
81 
82 
83# ── EIP-712 hash computation (manual, no encode_typed_data dependency) ────
84 
85def _eip712_domain_separator(domain: dict) -> bytes:
86    """Compute EIP-712 domain separator hash."""
87    type_hash = keccak(
88        b"EIP712Domain(string name,string version,uint256 chainId,address verifyingContract)"
89    )
90    encoded = type_hash
91    encoded += keccak(domain["name"].encode())
92    encoded += keccak(domain["version"].encode())
93    encoded += domain["chainId"].to_bytes(32, "big")
94    addr = domain["verifyingContract"].replace("0x", "")
95    encoded += bytes.fromhex(addr).rjust(32, b"\x00")
96    return keccak(encoded)
97 
98 
99def _eip712_encode_field(field_type: str, value) -> bytes:
100    """ABI-encode a single EIP-712 field value to 32 bytes."""
101    if field_type == "string":
102        return keccak(value.encode() if isinstance(value, str) else value)
103    elif field_type == "bytes32":
104        if isinstance(value, bytes):
105            return value.ljust(32, b"\x00")
106        elif isinstance(value, str):
107            return bytes.fromhex(value.replace("0x", "")).ljust(32, b"\x00")
108        elif isinstance(value, list):
109            return bytes(value).ljust(32, b"\x00")
110    elif field_type in ("uint256", "uint64"):
111        return int(value).to_bytes(32, "big")
112    elif field_type == "bool":
113        return int(bool(value)).to_bytes(32, "big")
114    elif field_type == "address":
115        return bytes.fromhex(value.replace("0x", "")).rjust(32, b"\x00")
116    raise ValueError(f"Unsupported EIP-712 type: {field_type}")
117 
118 
119def _eip712_hash_struct(primary_type: str, types: dict, message: dict) -> bytes:
120    """Compute hashStruct for an EIP-712 message."""
121    fields = types[primary_type]
122    type_string = (
123        primary_type
124        + "("
125        + ",".join(f"{f['type']} {f['name']}" for f in fields)
126        + ")"
127    )
128    encoded = keccak(type_string.encode())
129    for field in fields:
130        encoded += _eip712_encode_field(field["type"], message[field["name"]])
131    return keccak(encoded)
132 
133 
134def _eip712_digest(domain: dict, types: dict, primary_type: str, message: dict) -> bytes:
135    """Compute the full EIP-712 digest: keccak256(0x1901 || domainSep || structHash)."""
136    domain_sep = _eip712_domain_separator(domain)
137    struct_hash = _eip712_hash_struct(primary_type, types, message)
138    return keccak(b"\x19\x01" + domain_sep + struct_hash)
139 
140 
141# ── Local ecrecover verification ──────────────────────────────────────────
142 
143def _verify_signature_locally(
144    domain: dict,
145    types: dict,
146    primary_type: str,
147    message: dict,
148    signature_hex: str,
149    expected_address: Optional[str],
150) -> bool:
151    """
152    Verify an EIP-712 signature locally using ecrecover.
153 
154    Computes the EIP-712 hash manually (no encode_typed_data dependency)
155    and recovers the signer address from the signature.
156 
157    Returns True if the recovered address matches expected_address.
158    """
159    if not expected_address:
160        return True  # Can't verify without address, assume OK
161 
162    try:
163        from eth_keys import keys
164 
165        digest = _eip712_digest(domain, types, primary_type, message)
166 
167        sig_bytes = bytes.fromhex(signature_hex.replace("0x", ""))
168        r = int.from_bytes(sig_bytes[:32], "big")
169        s = int.from_bytes(sig_bytes[32:64], "big")
170        v = sig_bytes[64]
171        if v >= 27:
172            v -= 27
173 
174        sig = keys.Signature(vrs=(v, r, s))
175        public_key = sig.recover_public_key_from_msg_hash(digest)
176        recovered = public_key.to_checksum_address()
177        recovered_lower = recovered.lower()
178        expected_lower = expected_address.lower()
179 
180        if recovered_lower == expected_lower:
181            logger.info(f"EIP-712 ecrecover OK: {recovered}")
182            return True
183        else:
184            logger.warning(
185                f"EIP-712 ecrecover MISMATCH: recovered={recovered}, expected={expected_address}"
186            )
187            return False
188 
189    except Exception as e:
190        logger.warning(f"Local ecrecover verification failed: {e}")
191        return False
192 
193 
194def _signature_to_hex(sig: dict) -> str:
195    """Convert {r, s, v} dict to flat hex signature string."""
196    r = sig["r"].replace("0x", "").zfill(64)
197    s = sig["s"].replace("0x", "").zfill(64)
198    v = sig["v"]
199    return "0x" + r + s + format(v, "02x")
200 
201 
202# ── Core hashing ──────────────────────────────────────────────────────────
203 
204def action_hash(action: dict, vault_address: Optional[str], nonce: int) -> bytes:
205    """
206    Compute keccak256 hash of a Hyperliquid L1 action.
207 
208    Steps:
209    1. msgpack serialize the action dict
210    2. Append nonce as 8-byte big-endian
211    3. Append vault flag: 0x01 if vault_address else 0x00
212    4. If vault, append 20 bytes of vault address
213    5. keccak256 the whole blob
214    """
215    data = msgpack.packb(action)
216    data += nonce.to_bytes(8, "big")
217 
218    if vault_address:
219        data += b"\x01"
220        # Remove 0x prefix if present, decode hex to bytes
221        addr_bytes = bytes.fromhex(vault_address.replace("0x", ""))
222        data += addr_bytes
223    else:
224        data += b"\x00"
225 
226    return keccak(data)
227 
228 
229# ── L1 Action Signing (orders, cancels, leverage) ────────────────────────
230 
231async def sign_l1_action(
232    action: dict,
233    nonce: int,
234    vault_address: Optional[str] = None,
235) -> dict:
236    """
237    Sign an L1 action (order, cancel, leverage, modify) via wallet service.
238 
239    Tries multiple connectionId encodings and verifies each with local ecrecover.
240    Uses the first encoding that produces a signature matching the wallet address.
241 
242    Returns: {"r": "0x...", "s": "0x...", "v": 27|28}
243    """
244    hash_bytes = action_hash(action, vault_address, nonce)
245    source = MAINNET_SOURCE
246    expected_address = await _get_wallet_address()
247 
248    # Multiple connectionId encodings to try
249    encodings = [
250        ("hex_with_0x", "0x" + hash_bytes.hex()),          # current: hex string with prefix
251        ("byte_array",  list(hash_bytes)),                  # byte array as list of ints
252        ("hex_no_prefix", hash_bytes.hex()),                # hex string without 0x prefix
253    ]
254 
255    last_error = None
256    for encoding_name, connection_id_value in encodings:
257        try:
258            payload = {
259                "domain": L1_DOMAIN,
260                "types": AGENT_TYPES,
261                "primaryType": "Agent",
262                "message": {
263                    "source": source,
264                    "connectionId": connection_id_value,
265                },
266            }
267 
268            logger.info(
269                f"sign_l1_action: trying encoding={encoding_name}, "
270                f"connectionId type={type(connection_id_value).__name__}"
271            )
272 
273            result = await _wallet_request("POST", "/agent/sign-typed-data", payload)
274            sig = _parse_signature(result)
275 
276            # Verify locally with ecrecover
277            # For verification, connectionId must be the raw bytes32
278            verify_message = {
279                "source": source,
280                "connectionId": hash_bytes,
281            }
282 
283            sig_hex = _signature_to_hex(sig)
284            match = _verify_signature_locally(
285                domain=L1_DOMAIN,
286                types=AGENT_TYPES,
287                primary_type="Agent",
288                message=verify_message,
289                signature_hex=sig_hex,
290                expected_address=expected_address,
291            )
292 
293            if match:
294                logger.info(f"sign_l1_action: encoding={encoding_name} VERIFIED")
295                return sig
296            else:
297                logger.warning(
298                    f"sign_l1_action: encoding={encoding_name} signature mismatch, trying next"
299                )
300 
301        except Exception as e:
302            logger.warning(f"sign_l1_action: encoding={encoding_name} failed: {e}")
303            last_error = e
304 
305    raise RuntimeError(
306        f"sign_l1_action: No encoding produced a verified signature. "
307        f"Expected address: {expected_address}. Last error: {last_error}"
308    )
309 
310 
311# ── User-Signed Action (transfers, withdrawals, abstraction) ──────────────────────────
312 
313async def sign_user_set_abstraction(
314    user: str,
315    abstraction: str,
316    nonce: int,
317) -> dict:
318    """Sign a user set abstraction action.
319 
320    Uses HyperliquidSignTransaction domain for user-signed actions.
321    Matches SDK's sign_user_set_abstraction_action exactly.
322 
323    Args:
324        user: User address (lowercase)
325        abstraction: Abstraction mode ("unifiedAccount", "portfolioMargin", "disabled")
326        nonce: Timestamp in milliseconds
327 
328    Returns: {"r": "0x...", "s": "0x...", "v": 27|28}
329    """
330    types = {
331        "HyperliquidTransaction:UserSetAbstraction": [
332            {"name": "hyperliquidChain", "type": "string"},
333            {"name": "user", "type": "address"},  # ADDRESS, not string!
334            {"name": "abstraction", "type": "string"},
335            {"name": "nonce", "type": "uint64"},
336        ]
337    }
338 
339    action = {
340        "hyperliquidChain": "Mainnet",
341        "user": user.lower(),
342        "abstraction": abstraction,
343        "nonce": nonce,
344    }
345 
346    return await sign_user_action(
347        action=action,
348        types=types,
349        primary_type="HyperliquidTransaction:UserSetAbstraction",
350    )
351 
352 
353async def sign_user_action(
354    action: dict,
355    types: dict,
356    primary_type: str,
357) -> dict:
358    """
359    Sign a user-level action (USDC transfer, withdrawal) via wallet service.
360    Uses the HyperliquidSignTransaction domain (no msgpack/keccak).
361 
362    Verifies the signature locally with ecrecover before returning.
363 
364    Returns: {"r": "0x...", "s": "0x...", "v": 27|28}
365    """
366    expected_address = await _get_wallet_address()
367 
368    payload = {
369        "domain": USER_SIGNED_DOMAIN,
370        "types": types,
371        "primaryType": primary_type,
372        "message": action,
373    }
374 
375    result = await _wallet_request("POST", "/agent/sign-typed-data", payload)
376    sig = _parse_signature(result)
377 
378    # Verify locally
379    sig_hex = _signature_to_hex(sig)
380    match = _verify_signature_locally(
381        domain=USER_SIGNED_DOMAIN,
382        types=types,
383        primary_type=primary_type,
384        message=action,
385        signature_hex=sig_hex,
386        expected_address=expected_address,
387    )
388 
389    if not match:
390        logger.error(
391            f"sign_user_action: signature mismatch! "
392            f"Expected address: {expected_address}, action type: {primary_type}"
393        )
394 
395    return sig
396 
397 
398# ── Signature parsing ─────────────────────────────────────────────────────
399 
400def _parse_signature(result: dict) -> dict:
401    """
402    Parse signature from wallet service into {r, s, v} format.
403 
404    The wallet service may return either:
405    - A flat hex signature string in result["signature"]
406    - Already-split {r, s, v} components
407    """
408    sig = result.get("signature", result)
409 
410    if isinstance(sig, str):
411        # Flat hex signature — split into r, s, v
412        sig_hex = sig.replace("0x", "")
413        if len(sig_hex) == 130:
414            r = "0x" + sig_hex[:64]
415            s = "0x" + sig_hex[64:128]
416            v = int(sig_hex[128:130], 16)
417            # Normalize v to 27/28
418            if v < 27:
419                v += 27
420            return {"r": r, "s": s, "v": v}
421        raise ValueError(f"Unexpected signature length: {len(sig_hex)}")
422 
423    if isinstance(sig, dict):
424        # Already has components
425        r = sig.get("r", "")
426        s = sig.get("s", "")
427        v = sig.get("v", 0)
428        if isinstance(v, str):
429            v = int(v, 16) if v.startswith("0x") else int(v)
430        if v < 27:
431            v += 27
432        return {"r": r, "s": s, "v": v}
433 
434    raise ValueError(f"Cannot parse signature from wallet response: {result}")
435