#!/usr/bin/env python3
"""
HDGL Multi-Chain Bridge Daemon - Production Version
Bridges Bitcoin ↔ Solana ↔ Ethereum using HDGL lattice and ZK proofs
"""

import json
import subprocess
import time
import logging
import os
import struct
import base64
from decimal import Decimal, getcontext
from typing import Dict, List, Optional, Tuple
from dataclasses import dataclass
from concurrent.futures import ThreadPoolExecutor, as_completed

# Web3 and Solana imports
from web3 import Web3
from web3.middleware import geth_poa_middleware
from solana.rpc.api import Client as SolanaClient
from solana.publickey import PublicKey
from solana.transaction import Transaction
from solana.system_program import create_account, CreateAccountParams
from solana.account import Account as SolanaAccount

# HDGL operators (your custom module)
from hdgl_operators import D_n, blend_slot, STRANDS, WAVES, SCALE, lattice_commitment

# Configuration
getcontext().prec = 50
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler('hdgl_daemon.log'),
        logging.StreamHandler()
    ]
)
logger = logging.getLogger(__name__)

@dataclass
class HDGLConfig:
    # Bitcoin Covenant Parameters
    SEED_VECTORS: int = 2
    FOLD_STEPS: int = 28
    SEED_PARAMS: List[str] = None
    FIBS: List[int] = None
    CTV_PUBKEY: str = ""
    
    # Solana Configuration
    SOLANA_RPC_URL: str = "https://api.mainnet-beta.solana.com"
    SOLANA_PROGRAM_ID: str = ""
    SOLANA_STATE_ACCOUNT: str = ""
    SOLANA_PRIVATE_KEY: str = ""
    
    # Ethereum Configuration
    ETH_RPC_URL: str = ""
    ETH_CONTRACT_ADDRESS: str = ""
    ETH_PRIVATE_KEY: str = ""
    ETH_CHAIN_ID: int = 1
    
    # ZK Circuit Paths
    CIRCUIT_WASM: str = "hdgl_blending_js/hdgl_blending.wasm"
    CIRCUIT_ZKEY: str = "hdgl_blending_final.zkey"
    WITNESS_GEN_JS: str = "hdgl_blending_js/generate_witness.js"
    
    # Runtime Parameters
    POLL_INTERVAL: int = 30
    MAX_GAS_PRICE: int = 100_000_000_000  # 100 gwei
    PROOF_BATCH_SIZE: int = 5
    
    def __post_init__(self):
        if self.SEED_PARAMS is None:
            self.SEED_PARAMS = ["0x812", "0x502", "0x310", "0x192"]
        if self.FIBS is None:
            self.FIBS = [1, 1, 2, 3, 5, 8, 13, 21, 34, 55, 89]

@dataclass
class HDGLState:
    lattice: List[int]
    r_dim_scaled: int
    omega_scaled: int
    fold_step: int
    strand_id: int
    last_updated: int = 0

class SolanaManager:
    def __init__(self, config: HDGLConfig):
        self.config = config
        self.client = SolanaClient(config.SOLANA_RPC_URL)
        self.program_id = PublicKey(config.SOLANA_PROGRAM_ID)
        self.account = SolanaAccount(base64.b64decode(config.SOLANA_PRIVATE_KEY))
        
    def fetch_hdgl_state(self, state_pubkey: str) -> Optional[HDGLState]:
        """Fetch HDGLState from Solana account with error handling"""
        try:
            pubkey = PublicKey(state_pubkey)
            resp = self.client.get_account_info(pubkey)
            
            if not resp['result']['value']:
                logger.error(f"State account not found: {state_pubkey}")
                return None
            
            data_b64 = resp['result']['value']['data'][0]
            data_bytes = base64.b64decode(data_b64)
            
            # Deserialize HDGLState (256 u64s + metadata)
            if len(data_bytes) < 256 * 8 + 21:
                logger.error(f"Invalid state data length: {len(data_bytes)}")
                return None
                
            lattice = list(struct.unpack("<256Q", data_bytes[:256*8]))
            r_dim_scaled, omega_scaled = struct.unpack("<QQ", data_bytes[256*8:256*8+16])
            fold_step = struct.unpack("<I", data_bytes[256*8+16:256*8+20])[0]
            strand_id = struct.unpack("<B", data_bytes[256*8+20:256*8+21])[0]
            
            return HDGLState(
                lattice=lattice,
                r_dim_scaled=r_dim_scaled,
                omega_scaled=omega_scaled,
                fold_step=fold_step,
                strand_id=strand_id,
                last_updated=int(time.time())
            )
            
        except Exception as e:
            logger.error(f"Failed to fetch HDGL state: {e}")
            return None
    
    def trigger_fold_evolve(self, state_pubkey: str) -> bool:
        """Trigger FoldEvolve instruction on Solana"""
        try:
            # Create FoldEvolve instruction
            instruction_data = struct.pack("<B", 1)  # FoldEvolve = 1
            
            # Build transaction (simplified - use actual Solana transaction building)
            logger.info("Triggering fold evolution on Solana")
            # TODO: Implement actual transaction building
            return True
            
        except Exception as e:
            logger.error(f"Failed to trigger fold evolution: {e}")
            return False

class EthereumManager:
    def __init__(self, config: HDGLConfig):
        self.config = config
        self.w3 = Web3(Web3.HTTPProvider(config.ETH_RPC_URL))
        
        # Add PoA middleware for testnets
        if config.ETH_CHAIN_ID != 1:
            self.w3.middleware_onion.inject(geth_poa_middleware, layer=0)
            
        self.account = self.w3.eth.account.from_key(config.ETH_PRIVATE_KEY)
        self.w3.eth.default_account = self.account.address
        
        # Load contract ABI
        with open('hdgl_bridge_abi.json', 'r') as f:
            abi = json.load(f)
        
        self.contract = self.w3.eth.contract(
            address=config.ETH_CONTRACT_ADDRESS,
            abi=abi
        )
        
    def submit_zk_proof(self, proof: Dict, zk_input: Dict) -> Optional[str]:
        """Submit ZK proof to Ethereum contract"""
        try:
            # Check gas price
            gas_price = self.w3.eth.gas_price
            if gas_price > self.config.MAX_GAS_PRICE:
                logger.warning(f"Gas price too high: {gas_price} > {self.config.MAX_GAS_PRICE}")
                return None
            
            # Prepare proof data
            pi_a = proof["proof"]["pi_a"][:2]
            pi_b = proof["proof"]["pi_b"][:2][::-1]  # Reverse for Ethereum
            pi_c = proof["proof"]["pi_c"][:2]
            proof_bytes = pi_a + pi_b + pi_c
            
            # Build transaction
            tx = self.contract.functions.submitZKProof(
                proof_bytes,
                bytes.fromhex(zk_input["commitment"][2:]),
                int(zk_input["r_dim_scaled"]),
                int(zk_input["Omega_scaled"])
            ).build_transaction({
                "from": self.account.address,
                "gas": 800000,
                "gasPrice": min(gas_price, self.config.MAX_GAS_PRICE),
                "nonce": self.w3.eth.get_transaction_count(self.account.address)
            })
            
            # Sign and send
            signed_tx = self.w3.eth.account.sign_transaction(tx, private_key=self.config.ETH_PRIVATE_KEY)
            tx_hash = self.w3.eth.send_raw_transaction(signed_tx.rawTransaction)
            
            logger.info(f"ZK proof submitted: {tx_hash.hex()}")
            return tx_hash.hex()
            
        except Exception as e:
            logger.error(f"Failed to submit ZK proof: {e}")
            return None
    
    def get_latest_commitment(self) -> Optional[str]:
        """Get latest lattice commitment from contract"""
        try:
            commitment = self.contract.functions.getLatestCommitment().call()
            return commitment.hex()
        except Exception as e:
            logger.error(f"Failed to get latest commitment: {e}")
            return None

class ZKProofGenerator:
    def __init__(self, config: HDGLConfig):
        self.config = config
        self.executor = ThreadPoolExecutor(max_workers=2)
        
    def solana_to_zk_input(self, state: HDGLState) -> Dict:
        """Convert Solana HDGLState to ZK circuit input"""
        try:
            # Convert scaled values back to decimals
            r_dim = Decimal(state.r_dim_scaled) / Decimal(1_000_000_000)
            omega = Decimal(state.omega_scaled) / Decimal(1_000_000_000)
            
            # Process first 8 lattice slots
            raw_slots = [Decimal(x) / Decimal(1_000_000_000) for x in state.lattice[:8]]
            D_slots = [blend_slot(raw_slots, i, r_dim) for i in range(8)]
            
            # Generate commitment
            commitment_data = b''.join([
                int(slot * SCALE).to_bytes(32, 'big') for slot in D_slots
            ] + [
                int(r_dim * SCALE).to_bytes(32, 'big'),
                int(omega * SCALE).to_bytes(32, 'big')
            ])
            commitment = Web3.keccak(commitment_data).hex()
            
            return {
                "D_slots": [str(int(slot * SCALE)) for slot in D_slots],
                "WAVES": WAVES,
                "STRANDS": STRANDS,
                "r_dim_scaled": str(int(r_dim * SCALE)),
                "Omega_scaled": str(int(omega * SCALE)),
                "commitment": commitment
            }
            
        except Exception as e:
            logger.error(f"Failed to convert state to ZK input: {e}")
            return {}
    
    def generate_proof_async(self, zk_input: Dict) -> Dict:
        """Generate ZK proof asynchronously"""
        try:
            # Write input file
            input_file = f"input_{int(time.time())}.json"
            with open(input_file, "w") as f:
                json.dump(zk_input, f)
            
            witness_file = f"witness_{int(time.time())}.wtns"
            proof_file = f"proof_{int(time.time())}.json"
            public_file = f"public_{int(time.time())}.json"
            
            # Generate witness
            subprocess.run([
                "node", self.config.WITNESS_GEN_JS,
                self.config.CIRCUIT_WASM,
                input_file,
                witness_file
            ], check=True, capture_output=True)
            
            # Generate proof
            subprocess.run([
                "snarkjs", "groth16", "prove",
                self.config.CIRCUIT_ZKEY,
                witness_file,
                proof_file,
                public_file
            ], check=True, capture_output=True)
            
            # Load results
            with open(proof_file) as f:
                proof = json.load(f)
            with open(public_file) as f:
                public_inputs = json.load(f)
            
            # Cleanup
            for file in [input_file, witness_file, proof_file, public_file]:
                try:
                    os.unlink(file)
                except:
                    pass
            
            return {"proof": proof, "public": public_inputs, "input": zk_input}
            
        except Exception as e:
            logger.error(f"Proof generation failed: {e}")
            return {}

class BitcoinCovenantManager:
    def __init__(self, config: HDGLConfig):
        self.config = config
        
    def generate_covenant_script(self, vector_count: int) -> str:
        """Generate Bitcoin covenant script for current state"""
        script_lines = [
            f"{self.config.CTV_PUBKEY} OP_CHECKSIGVERIFY",
            "# -- HDGL Vector Folding --"
        ]
        
        current_vectors = self.config.SEED_VECTORS
        fib_index = 0
        
        for step in range(min(self.config.FOLD_STEPS, 24)):  # Cap at 24 for production
            prev_count = current_vectors
            current_vectors *= 2
            
            # Fold operations
            script_lines.extend([
                f"OP_VEC_DUP {prev_count} OP_HDGL_VEC_ADD OP_EQUALVERIFY",
                f"{prev_count} {prev_count} OP_HDGL_PARAM_ADD OP_EQUALVERIFY"
            ])
            
            # Fibonacci uniqueness check
            param = self.config.SEED_PARAMS[step % len(self.config.SEED_PARAMS)]
            omega_val = int(param, 16)
            fib_delta = self.config.FIBS[fib_index % len(self.config.FIBS)]
            fib_index += 1
            
            omega_bound = omega_val + fib_delta
            script_lines.append(f"{omega_val} {omega_bound} {omega_bound} OP_WITHINVERIFY")
            
            if current_vectors >= vector_count:
                break
        
        # CHG Bridge operations
        script_lines.extend([
            "",
            "# -- CHG Cross-Chain Bridge --",
            "OP_CHG_MAP_INIT",
            "OP_CHG_DEPOSIT_VERIFY <ETH_TOKEN_ID> <BTC_VECTOR_ID> <AMOUNT> <PROOF>",
            "OP_CHG_MARK_CLAIMABLE <BTC_VECTOR_ID>",
            "OP_CHG_REDEEM_VERIFY <BTC_VECTOR_ID> <ETH_RECIPIENT> <PROOF>",
            "OP_CHG_LOCK_VECTOR <BTC_VECTOR_ID>",
            "OP_CHG_UPDATE_MAPPING <BTC_VECTOR_ID> <NEW_OWNER>",
            "OP_CHECKTEMPLATEVERIFY"
        ])
        
        return "\n".join(script_lines)

class HDGLBridgeDaemon:
    def __init__(self, config: HDGLConfig):
        self.config = config
        self.solana = SolanaManager(config)
        self.ethereum = EthereumManager(config)
        self.zk_gen = ZKProofGenerator(config)
        self.bitcoin = BitcoinCovenantManager(config)
        self.running = False
        self.stats = {
            "proofs_generated": 0,
            "proofs_submitted": 0,
            "errors": 0,
            "start_time": time.time()
        }
    
    def start(self):
        """Start the daemon main loop"""
        logger.info("Starting HDGL Multi-Chain Bridge Daemon")
        self.running = True
        
        try:
            while self.running:
                self.process_cycle()
                time.sleep(self.config.POLL_INTERVAL)
                
        except KeyboardInterrupt:
            logger.info("Daemon shutdown requested")
        except Exception as e:
            logger.error(f"Daemon error: {e}")
        finally:
            self.shutdown()
    
    def process_cycle(self):
        """Process one complete bridge cycle"""
        try:
            # 1. Fetch current Solana state
            state = self.solana.fetch_hdgl_state(self.config.SOLANA_STATE_ACCOUNT)
            if not state:
                logger.warning("Could not fetch Solana state, skipping cycle")
                return
            
            logger.info(f"Fetched Solana state: fold_step={state.fold_step}, strand={state.strand_id}")
            
            # 2. Check if we need to evolve
            if self.should_trigger_evolution(state):
                if self.solana.trigger_fold_evolve(self.config.SOLANA_STATE_ACCOUNT):
                    logger.info("Triggered fold evolution")
                    time.sleep(5)  # Wait for state update
                    state = self.solana.fetch_hdgl_state(self.config.SOLANA_STATE_ACCOUNT)
            
            # 3. Generate ZK input
            zk_input = self.zk_gen.solana_to_zk_input(state)
            if not zk_input:
                logger.error("Failed to generate ZK input")
                return
            
            # 4. Generate proof
            logger.info("Generating ZK proof...")
            proof_data = self.zk_gen.generate_proof_async(zk_input)
            if not proof_data:
                logger.error("Proof generation failed")
                self.stats["errors"] += 1
                return
            
            self.stats["proofs_generated"] += 1
            
            # 5. Submit to Ethereum
            tx_hash = self.ethereum.submit_zk_proof(proof_data["proof"], zk_input)
            if tx_hash:
                self.stats["proofs_submitted"] += 1
                logger.info(f"Proof submitted successfully: {tx_hash}")
            else:
                self.stats["errors"] += 1
                
            # 6. Update Bitcoin covenant if needed
            self.update_bitcoin_covenant(state)
            
            # 7. Log stats
            self.log_stats()
            
        except Exception as e:
            logger.error(f"Cycle processing error: {e}")
            self.stats["errors"] += 1
    
    def should_trigger_evolution(self, state: HDGLState) -> bool:
        """Determine if we should trigger fold evolution"""
        # Trigger every 10 cycles or if state is stale
        time_since_update = time.time() - state.last_updated
        return (state.fold_step < 24 and 
                (time_since_update > 300 or state.fold_step % 5 == 0))
    
    def update_bitcoin_covenant(self, state: HDGLState):
        """Update Bitcoin covenant script based on current state"""
        try:
            vector_count = self.config.SEED_VECTORS * (2 ** state.fold_step)
            script = self.bitcoin.generate_covenant_script(vector_count)
            
            filename = f"hdgl_chg_covenant_fold_{state.fold_step}.script"
            with open(filename, "w") as f:
                f.write(script)
                
            logger.info(f"Updated Bitcoin covenant: {filename}")
            
        except Exception as e:
            logger.error(f"Failed to update Bitcoin covenant: {e}")
    
    def log_stats(self):
        """Log daemon statistics"""
        uptime = time.time() - self.stats["start_time"]
        logger.info(f"Stats - Uptime: {uptime:.0f}s, "
                   f"Proofs: {self.stats['proofs_generated']}/{self.stats['proofs_submitted']}, "
                   f"Errors: {self.stats['errors']}")
    
    def shutdown(self):
        """Graceful shutdown"""
        logger.info("Shutting down HDGL Bridge Daemon")
        self.running = False
        self.zk_gen.executor.shutdown(wait=True)

def load_config(config_file: str = "hdgl_config.json") -> HDGLConfig:
    """Load configuration from JSON file"""
    try:
        with open(config_file, 'r') as f:
            config_dict = json.load(f)
        
        return HDGLConfig(**config_dict)
    except Exception as e:
        logger.error(f"Failed to load config: {e}")
        return HDGLConfig()

def main():
    """Main entry point"""
    config = load_config()
    daemon = HDGLBridgeDaemon(config)
    daemon.start()

if __name__ == "__main__":
    main()
