State Management

pub struct StateManager {
    accounts: BTreeMap<Address, CoreLaneAccount>,          // Account balances and nonces
    stored_blobs: BTreeMap<B256, Vec<u8>>,                // Large data blobs
    kv_storage: BTreeMap<String, Vec<u8>>,                // Key-value storage
    intents: BTreeMap<B256, Intent>,                      // Intent marketplace
    transactions: Vec<StoredTransaction>,                  // Transaction history
    transaction_receipts: BTreeMap<String, TransactionReceipt>, // Receipts
}

pub struct BundleStateManager {
    pub accounts: BTreeMap<Address, CoreLaneAccount>,      // Modified accounts
    pub stored_blobs: BTreeMap<B256, Vec<u8>>,            // New blobs
    pub kv_storage: BTreeMap<String, Vec<u8>>,            // Modified KV pairs
    pub removed_keys: Vec<String>,                        // Deleted KV keys
    pub intents: BTreeMap<B256, Intent>,                  // Modified intents
    pub transactions: Vec<StoredTransaction>,              // New transactions
    pub transaction_receipts: BTreeMap<String, TransactionReceipt>, // New receipts
}

pub struct CoreLaneAccount {
    pub balance: U256,  // Account balance in wei
    pub nonce: U256,    // Transaction nonce for replay protection
}

pub fn get_balance(&self, original: &StateManager, address: Address) -> U256 {
    // Check bundle state first
    if let Some(account) = self.accounts.get(&address) {
        return account.balance;
    }
    // Fall back to original state
    original.get_balance(address)
}

pub fn get_nonce(&self, original: &StateManager, address: Address) -> U256 {
    if let Some(account) = self.accounts.get(&address) {
        return account.nonce;
    }
    original.get_nonce(address)
}

pub fn get_account_mut(
    &mut self,
    original: &StateManager,
    address: Address,
) -> Option<&mut CoreLaneAccount> {
    // Ensure account exists in bundle before getting mutable reference
    self.accounts.entry(address).or_insert_with(|| {
        original
            .get_account(address)
            .cloned()
            .unwrap_or_else(CoreLaneAccount::new)  // Create if doesn't exist
    });
    
    // Now get the mutable reference (account definitely exists)
    self.accounts.get_mut(&address)
}

pub fn add_balance(
    &mut self,
    original: &StateManager,
    address: Address,
    amount: U256,
) -> Result<()> {
    let account = self.get_account_mut(original, address);
    if let Some(account) = account {
        account.add_balance(amount)?;  // Uses checked_add internally
    } else {
        let mut account = CoreLaneAccount::new();
        account.add_balance(amount)?;
        self.accounts.insert(address, account);
    }
    Ok(())
}

pub fn sub_balance(
    &mut self,
    original: &StateManager,
    address: Address,
    amount: U256,
) -> Result<()> {
    let account = self.get_account_mut(original, address);
    if let Some(account) = account {
        account.sub_balance(amount)?;  // Uses checked_sub, fails on underflow
    } else {
        return Err(anyhow::anyhow!("Account not found"));
    }
    Ok(())
}

pub fn increment_nonce(&mut self, original: &StateManager, address: Address) -> Result<()> {
    let account = self.get_account_mut(original, address);
    if let Some(account) = account {
        account.increment_nonce()?;
    } else {
        let mut account = CoreLaneAccount::new();
        account.increment_nonce()?;
        self.accounts.insert(address, account);
    }
    Ok(())
}

pub fn insert_blob(&mut self, blob_hash: B256, data: Vec<u8>) {
    self.stored_blobs.insert(blob_hash, data);
}

pub fn contains_blob(&self, original: &StateManager, blob_hash: &B256) -> bool {
    self.stored_blobs.contains_key(blob_hash) || original.contains_blob(blob_hash)
}

pub fn get_intent_mut(
    &mut self,
    original: &StateManager,
    intent_id: &B256,
) -> Option<&mut Intent> {
    // Ensure the intent exists in our bundle before getting a mutable reference
    if !self.intents.contains_key(intent_id) {
        if let Some(orig_intent) = original.get_intent(intent_id) {
            self.intents.insert(*intent_id, orig_intent.clone());
        } else {
            return None;
        }
    }
    self.intents.get_mut(intent_id)
}

pub fn insert_intent(&mut self, intent_id: B256, intent: Intent) {
    self.intents.insert(intent_id, intent);
}

// Bundle state
pub fn get_kv(&self, key: &str) -> Option<&Vec<u8>> {
    self.kv_storage.get(key)
}

pub fn insert_kv(&mut self, key: String, value: Vec<u8>) {
    self.kv_storage.insert(key.clone(), value);
    self.removed_keys.retain(|k| k != &key);  // Un-delete if previously removed
}

pub fn remove_kv(&mut self, key: &str) -> Option<Vec<u8>> {
    let removed = self.kv_storage.remove(key);
    let key_string = key.to_string();
    if !self.removed_keys.iter().any(|k| k == &key_string) {
        self.removed_keys.push(key_string);  // Track deletion
    }
    removed
}

pub struct StoredTransaction {
    pub envelope: TxEnvelope,    // Full transaction envelope
    pub raw_data: Vec<u8>,       // Raw transaction bytes
    pub block_number: u64,       // Block height when included
}

impl BorshSerialize for StoredTransaction {
    fn serialize<W: std::io::Write>(&self, writer: &mut W) -> std::io::Result<()> {
        // Serialize raw_data and block_number
        // Envelope reconstructed on deserialization from raw_data
        BorshSerialize::serialize(&self.raw_data, writer)?;
        BorshSerialize::serialize(&self.block_number, writer)?;
        Ok(())
    }
}

impl BorshDeserialize for StoredTransaction {
    fn deserialize_reader<R: std::io::Read>(reader: &mut R) -> std::io::Result<Self> {
        let raw_data: Vec<u8> = borsh::BorshDeserialize::deserialize_reader(reader)?;
        let block_number: u64 = borsh::BorshDeserialize::deserialize_reader(reader)?;
        
        // Reconstruct envelope from raw_data
        let envelope = TxEnvelope::decode(&mut raw_data.as_slice())
            .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))?;
        
        Ok(StoredTransaction { envelope, raw_data, block_number })
    }
}

pub struct TransactionReceipt {
    pub transaction_hash: String,
    pub block_number: u64,
    pub transaction_index: u64,
    pub from: String,
    pub to: Option<String>,
    pub cumulative_gas_used: String,
    pub gas_used: String,
    pub contract_address: Option<String>,
    pub logs: Vec<Log>,
    pub status: String,  // "0x1" for success, "0x0" for failure
    pub effective_gas_price: String,
    pub tx_type: String,
    pub logs_bloom: String,
}

pub fn apply_changes(&mut self, bundle_state_manager: BundleStateManager) {
    // Apply account changes
    for (address, account) in bundle_state_manager.accounts.into_iter() {
        tracing::info!("Applying changes for account {}", address);
        self.set_account(address, account);
    }
    
    // Apply blob storage changes
    for (blob_hash, data) in bundle_state_manager.stored_blobs.into_iter() {
        self.insert_blob(blob_hash, data);
    }
    
    // Apply intent changes
    for (intent_id, intent) in bundle_state_manager.intents.into_iter() {
        self.insert_intent(intent_id, intent);
    }
    
    // Remove deleted keys
    for key in bundle_state_manager.removed_keys.into_iter() {
        self.remove_kv(&key);
    }
    
    // Apply KV storage changes
    for (key, value) in bundle_state_manager.kv_storage.into_iter() {
        self.insert_kv(key, value);
    }
    
    // Apply transaction storage
    for transaction in bundle_state_manager.transactions.into_iter() {
        self.add_transaction(transaction);
    }
    
    // Apply transaction receipts
    for (tx_hash, receipt) in bundle_state_manager.transaction_receipts.into_iter() {
        self.add_receipt(tx_hash, receipt);
    }
}

// StateManager
pub fn borsh_serialize(&self) -> Result<Vec<u8>> {
    borsh::to_vec(self)
        .map_err(|e| anyhow::anyhow!("Failed to borsh serialize StateManager: {}", e))
}

pub fn borsh_deserialize(bytes: &[u8]) -> Result<Self> {
    let mut cursor = std::io::Cursor::new(bytes);
    let result = borsh::BorshDeserialize::deserialize_reader(&mut cursor);
    let consumed = cursor.position() as usize;
    
    match result {
        Ok(value) => {
            // Ensure all bytes consumed
            if consumed != bytes.len() {
                return Err(anyhow::anyhow!(
                    "Not all bytes read: state file length {} bytes, consumed {} bytes ({} trailing)",
                    bytes.len(), consumed, bytes.len().saturating_sub(consumed)
                ));
            }
            Ok(value)
        }
        Err(e) => Err(anyhow::anyhow!(
            "Failed to borsh deserialize StateManager (file length {} bytes, consumed {} bytes before error): {}",
            bytes.len(), consumed, e
        )),
    }
}

pub fn borsh_serialize(&self) -> Result<Vec<u8>> {
    borsh::to_vec(self)
        .map_err(|e| anyhow::anyhow!("Failed to borsh serialize BundleStateManager: {}", e))
}

pub fn borsh_deserialize(bytes: &[u8]) -> Result<Self> {
    borsh::from_slice(bytes)
        .map_err(|e| anyhow::anyhow!("Failed to borsh deserialize BundleStateManager: {}", e))
}

pub trait ProcessingContext {
    fn state_manager(&self) -> &StateManager;
    fn state_manager_mut(&mut self) -> &mut StateManager;
    fn bitcoin_client_read(&self) -> Option<Arc<dyn BitcoinRpcReadClient>>;
    fn bitcoin_network(&self) -> bitcoin::Network;
    fn handle_cmio_query(
        &mut self,
        message: CmioMessage,
        current_intent_id: Option<B256>,
    ) -> Option<CmioMessage>;
}

use core_lane::{CoreLaneStateForLib, StateManager, BundleStateManager};
use core_lane::{execute_transaction, TxEnvelope};

let mut state = CoreLaneStateForLib::new(
    StateManager::new(),
    bitcoin_client_read,
    bitcoin_client_write,
    bitcoin::Network::Regtest
);

let mut bundle = BundleStateManager::new();

// Execute transaction
let result = execute_transaction(
    &tx_envelope,
    sender,
    &mut bundle,
    &mut state,
    block_timestamp,
)?;

if result.success {
    // Apply bundle to state
    state.state_manager_mut().apply_changes(bundle);
}

#[test]
fn test_bundle_state_manager() {
    let state = StateManager::new();
    let mut bundle = BundleStateManager::new();
    
    let test_addr = Address::from([1u8; 20]);
    let amount = U256::from(1000);
    
    // Add balance in bundle
    bundle.add_balance(&state, test_addr, amount).unwrap();
    
    // Check it's reflected in bundle
    assert_eq!(bundle.get_balance(&state, test_addr), amount);
    
    // Original state should be unchanged
    assert_eq!(state.get_balance(test_addr), U256::ZERO);
    
    // Apply changes
    let mut new_state = state;
    new_state.apply_changes(bundle);
    assert_eq!(new_state.get_balance(test_addr), amount);
}

#[test]
fn test_state_serialization() {
    let mut state = StateManager::new();
    let test_addr = Address::from([1u8; 20]);
    let amount = U256::from(1000);
    
    // Add balance
    let mut bundle = BundleStateManager::new();
    bundle.add_balance(&state, test_addr, amount).unwrap();
    state.apply_changes(bundle);
    
    // Serialize
    let serialized = state.borsh_serialize().unwrap();
    
    // Deserialize
    let deserialized = StateManager::borsh_deserialize(&serialized).unwrap();
    assert_eq!(deserialized.get_balance(test_addr), amount);
}