#!/usr/bin/env python3
"""
Julie Security Network — Production-Ready x402 Client
Target: https://julie.freysa.dev/api/audit
Network: Base Mainnet (ChainID 8453)
Settlement Asset: Base Native USDC (0x833589fCD6eDb6E08f4c7C32D4f71b54bdA02913)
PayTo Vault: 0x75F05Da87cDFf3f6bAcAc3C07CD49F1D7dC7b20C

Core Capabilities:
1. Automated HTTP 402 Payment Required handling (EIP-712 Permit signing on Base)
2. Session pass management (?auto_pass=starter / $1.00 for 125 scans) with local session key caching
3. Connection pooling (requests.Session) with keep-alive for sub-millisecond API repeat scans
4. CLI and Python module interface with zero hardcoded keys
"""

import os
import sys
import json
import time
import base64
import argparse
from typing import Optional, Dict, Any, List
import requests
from web3 import Web3
from eth_account import Account
from eth_account.messages import encode_typed_data

# Default Network & Settlement Constants
DEFAULT_API_BASE = os.getenv("JULIE_API_URL", "https://julie.freysa.dev")
BASE_CHAIN_ID = 8453
BASE_RPC_URL = os.getenv("BASE_RPC_URL", "https://mainnet.base.org")
BASE_USDC_ADDRESS = Web3.to_checksum_address("0x833589fCD6eDb6E08f4c7C32D4f71b54bdA02913")
SETTLEMENT_VAULT = Web3.to_checksum_address("0x75F05Da87cDFf3f6bAcAc3C07CD49F1D7dC7b20C")

# Standard ERC20 Nonces ABI for EIP-2612 / EIP-3009 permit queries
NONCES_ABI = [
    {
        "constant": True,
        "inputs": [{"name": "owner", "type": "address"}],
        "name": "nonces",
        "outputs": [{"name": "", "type": "uint256"}],
        "type": "function",
    }
]

SESSION_CACHE_FILE = os.path.expanduser("~/.julie_session_cache.json")


class JulieSecurityClient:
    """Production-grade client for the Julie Security Network x402 API."""

    def __init__(
        self,
        private_key: Optional[str] = None,
        api_base: str = DEFAULT_API_BASE,
        rpc_url: str = BASE_RPC_URL,
        auto_session_pass: bool = True,
        session_tier: str = "starter",  # "starter" ($1 = 125 scans), "pro" ($5 = 750 scans)
        cache_file: str = SESSION_CACHE_FILE,
    ):
        self.api_base = api_base.rstrip("/")
        self.rpc_url = rpc_url
        self.auto_session_pass = auto_session_pass
        self.session_tier = session_tier
        self.cache_file = cache_file

        # Web3 & Account Initialization
        self.private_key = private_key or os.getenv("CLIENT_PRIVATE_KEY")
        self.w3 = Web3(Web3.HTTPProvider(self.rpc_url, request_kwargs={"timeout": 5}))
        self.account = Account.from_key(self.private_key) if self.private_key else None

        # Persistent HTTP connection pooling
        self.session = requests.Session()
        adapter = requests.adapters.HTTPAdapter(pool_connections=10, pool_maxsize=20, max_retries=2)
        self.session.mount("https://", adapter)
        self.session.mount("http://", adapter)

        # In-memory & disk-backed session token cache
        self.session_key = self._load_cached_session()

    def _load_cached_session(self) -> Optional[str]:
        """Loads cached session key from disk if still valid."""
        if not os.path.exists(self.cache_file):
            return None
        try:
            with open(self.cache_file, "r") as f:
                data = json.load(f)
                key = data.get("session_key")
                expires_at = data.get("expires_at", 0)
                if key and expires_at > time.time():
                    return key
        except Exception:
            pass
        return None

    def _save_cached_session(self, session_key: str, expires_in_sec: int = 604800):
        """Caches session key to disk."""
        self.session_key = session_key
        try:
            with open(self.cache_file, "w") as f:
                json.dump({
                    "session_key": session_key,
                    "saved_at": time.time(),
                    "expires_at": time.time() + expires_in_sec
                }, f)
        except Exception:
            pass

    def get_session_balance(self) -> Dict[str, Any]:
        """Queries remaining scan balance on the active session key."""
        if not self.session_key:
            return {"is_active": False, "credits_remaining": 0}
        try:
            resp = self.session.get(
                f"{self.api_base}/api/session/balance",
                headers={"X-Session-Key": self.session_key},
                timeout=3
            )
            if resp.status_code == 200:
                return resp.json()
        except Exception:
            pass
        return {"is_active": False, "credits_remaining": 0}

    def _sign_eip712_payment(self, challenge: Dict[str, Any]) -> str:
        """
        Signs an EIP-712 Permit payment payload for Base Native USDC.
        Returns the base64-encoded payment authorization header.
        """
        if not self.account:
            raise ValueError(
                "Private key required to sign x402 payment. Set CLIENT_PRIVATE_KEY or pass private_key."
            )

        accepts = challenge.get("accepts", [{}])[0]
        recipient = Web3.to_checksum_address(accepts.get("payTo", SETTLEMENT_VAULT))
        token_addr = Web3.to_checksum_address(accepts.get("asset", BASE_USDC_ADDRESS))
        amount = int(accepts.get("amount", 10000))  # 10,000 units = $0.01 USDC; 1,000,000 = $1.00 USDC

        # Query current on-chain nonce for USDC permit
        usdc_contract = self.w3.eth.contract(address=token_addr, abi=NONCES_ABI)
        try:
            nonce = usdc_contract.functions.nonces(self.account.address).call()
        except Exception:
            nonce = 0

        deadline = int(time.time()) + 3600  # 1-hour expiration

        domain_data = {
            "name": "USD Coin",
            "version": "2",
            "chainId": BASE_CHAIN_ID,
            "verifyingContract": token_addr,
        }

        permit_types = {
            "Permit": [
                {"name": "owner", "type": "address"},
                {"name": "spender", "type": "address"},
                {"name": "value", "type": "uint256"},
                {"name": "nonce", "type": "uint256"},
                {"name": "deadline", "type": "uint256"},
            ]
        }

        permit_message = {
            "owner": self.account.address,
            "spender": recipient,
            "value": amount,
            "nonce": nonce,
            "deadline": deadline,
        }

        structured_data = {
            "types": {
                "EIP712Domain": [
                    {"name": "name", "type": "string"},
                    {"name": "version", "type": "string"},
                    {"name": "chainId", "type": "uint256"},
                    {"name": "verifyingContract", "type": "address"},
                ],
                **permit_types,
            },
            "primaryType": "Permit",
            "domain": domain_data,
            "message": permit_message,
        }

        encoded = encode_typed_data(full_message=structured_data)
        signed = self.account.sign_message(encoded)

        payload = {
            "scheme": "exact",
            "network": f"eip155:{BASE_CHAIN_ID}",
            "payer": self.account.address,
            "payTo": recipient,
            "amount": amount,
            "signature": signed.signature.hex(),
            "deadline": deadline,
            "nonce": nonce,
        }

        return base64.b64encode(json.dumps(payload).encode()).decode()

    def audit(
        self,
        contract: str,
        chain: str = "base",
        request_pass_on_402: Optional[bool] = None,
    ) -> Dict[str, Any]:
        """
        Audits a smart contract with automatic x402 payment handling & session pass reuse.

        Flow:
        1. If cached session key exists, attach X-Session-Key (0.5ms scan, no payment).
        2. If 402 received:
           a. If request_pass_on_402 is True, request session pass (?auto_pass=starter)
           b. Sign EIP-712 permit and retry request.
           c. Extract and cache returned X-Session-Key.
        """
        t0 = time.time()
        should_buy_pass = self.auto_session_pass if request_pass_on_402 is None else request_pass_on_402

        # 1. Build request parameters and headers
        url = f"{self.api_base}/api/audit"
        params: Dict[str, str] = {"contract": contract, "chain": chain}
        headers: Dict[str, str] = {"Accept": "application/json"}

        # Use cached session key if available
        if self.session_key:
            headers["X-Session-Key"] = self.session_key

        t_req_start = time.time()
        resp = self.session.get(url, params=params, headers=headers, timeout=5)
        probe_ms = (time.time() - t_req_start) * 1000

        # Fast path: 200 OK (free trial, cached, or session authenticated)
        if resp.status_code == 200:
            total_ms = (time.time() - t0) * 1000
            data = resp.json()
            data["_client_perf"] = {
                "status": "cached_or_session",
                "round_trips": 1,
                "total_ms": round(total_ms, 2),
                "probe_ms": round(probe_ms, 2),
            }
            return data

        # Check for 402 Payment Required
        if resp.status_code == 402:
            print(f"[!] HTTP 402 Payment Required received for {contract[:10]}... ({probe_ms:.1f}ms)")

            # Invalidate expired session key if present
            if self.session_key:
                self.session_key = None

            # Prepare Auto-Pass parameter if enabled
            if should_buy_pass:
                params["auto_pass"] = self.session_tier
                if self.account:
                    params["payer"] = self.account.address

            # Parse payment challenge header or fallback body
            payment_b64 = resp.headers.get("payment-required") or resp.headers.get("x-payment-required")
            if payment_b64:
                challenge = json.loads(base64.b64decode(payment_b64).decode())
            else:
                challenge = resp.json()

            # If no account / private key provided, report instructions cleanly
            if not self.account:
                print("\n[i] No private key provided. Set CLIENT_PRIVATE_KEY or pass --key to automate payment.")
                print(f"[*] To pay via Foundry/cast:")
                print(f"    cast send {BASE_USDC_ADDRESS} 'transfer(address,uint256)' {SETTLEMENT_VAULT} 1000000 --rpc-url {self.rpc_url} --private-key $YOUR_KEY")
                print(f"[*] Then retry with session pass or direct tx verification.")
                return {
                    "status": "payment_required",
                    "http_status": 402,
                    "challenge": challenge,
                    "_client_perf": {
                        "status": "unpaid_challenge",
                        "round_trips": 1,
                        "probe_ms": round(probe_ms, 2),
                    }
                }

            # Sign EIP-712 authorization
            t_sign_start = time.time()
            payment_header = self._sign_eip712_payment(challenge)
            sign_ms = (time.time() - t_sign_start) * 1000

            # Submit paid request with authorization header
            paid_headers = {
                "Accept": "application/json",
                "PAYMENT-SIGNATURE": payment_header,
                "X-PAYMENT": payment_header,
            }
            if should_buy_pass:
                paid_headers["X-Auto-Session-Pass"] = self.session_tier

            t_paid_start = time.time()
            paid_resp = self.session.get(url, params=params, headers=paid_headers, timeout=5)
            paid_ms = (time.time() - t_paid_start) * 1000
            total_ms = (time.time() - t0) * 1000

            if paid_resp.status_code == 200:
                result = paid_resp.json()

                # Extract and cache new session key if returned
                new_session_key = paid_resp.headers.get("X-Session-Key") or result.get("session_key")
                if new_session_key:
                    self._save_cached_session(new_session_key)
                    print(f"[+] Session pass activated! Key cached: {new_session_key[:16]}...")

                result["_client_perf"] = {
                    "status": "paid_settled",
                    "round_trips": 2,
                    "total_ms": round(total_ms, 2),
                    "probe_ms": round(probe_ms, 2),
                    "sign_ms": round(sign_ms, 2),
                    "paid_ms": round(paid_ms, 2),
                }
                return result
            else:
                raise RuntimeError(
                    f"Payment settlement failed: HTTP {paid_resp.status_code} - {paid_resp.text}"
                )

        # Other non-200 / non-402 response
        resp.raise_for_status()
        return resp.json()

    def batch_audit(self, contracts: List[str], chain: str = "base") -> Dict[str, Any]:
        """Audits a batch of contracts in a single call or via cached session key."""
        results = []
        for c in contracts:
            try:
                res = self.audit(contract=c, chain=chain)
                results.append({"contract": c, "status": "success", "data": res})
            except Exception as e:
                results.append({"contract": c, "status": "error", "error": str(e)})
        return {"batch_size": len(contracts), "results": results}


def main():
    parser = argparse.ArgumentParser(description="Julie Security Network x402 Automated Client")
    parser.add_argument("--contract", "-c", help="Target contract address to audit", default="0x833589fCD6eDb6E08f4c7C32D4f71b54bdA02913")
    parser.add_argument("--chain", help="Blockchain network name (default: base)", default="base")
    parser.add_argument("--key", "-k", help="Client Base private key (or set CLIENT_PRIVATE_KEY)", default=None)
    parser.add_argument("--api", help="Julie API base URL", default=DEFAULT_API_BASE)
    parser.add_argument("--auto-pass", help="Session pass tier ('starter', 'pro', or 'none')", default="starter")
    parser.add_argument("--batch", help="Comma-separated list of contracts for batch auditing", default=None)
    parser.add_argument("--check-session", action="store_true", help="Check remaining session balance and exit")

    args = parser.parse_args()

    client = JulieSecurityClient(
        private_key=args.key,
        api_base=args.api,
        auto_session_pass=(args.auto_pass.lower() != "none"),
        session_tier=args.auto_pass.lower(),
    )

    if args.check_session:
        bal = client.get_session_balance()
        print(json.dumps(bal, indent=2))
        sys.exit(0)

    if args.batch:
        contract_list = [c.strip() for c in args.batch.split(",") if c.strip()]
        print(f"[*] Running batch audit on {len(contract_list)} contracts...")
        res = client.batch_audit(contract_list, chain=args.chain)
        print(json.dumps(res, indent=2))
        sys.exit(0)

    print(f"[*] Running audit for contract: {args.contract} on {args.chain}...")
    res = client.audit(contract=args.contract, chain=args.chain)
    print("\n[+] Audit Result:")
    print(json.dumps(res, indent=2))


if __name__ == "__main__":
    main()