"""
Simple Bitcoin Wallet Implementation
Handles real Bitcoin transactions with minimal setup
"""

import os
import json
import time
import hashlib
import secrets
import base58
from typing import Optional, Dict, Any, List
import ecdsa
from dataclasses import dataclass
from decimal import Decimal
from bitcoinlib.services.services import Service
from src.mining.reward_processor import BitcoinRewardProcessor

# For real Bitcoin network operations
try:
    from bitcoin_network import BitcoinNetwork
except ImportError:
    # Fallback implementation if bitcoin_network is not available
    class BitcoinNetwork:
        def __init__(self, network='bitcoin'):
            self.network = network
            
        def get_balance(self, address):
            # For testing, return the wallet file balance
            return Decimal('0')
            
        def create_transaction(self, from_address, to_address, amount, private_key, fee_rate, utxos=None):
            if not utxos:
                raise ValueError("No unspent transaction outputs found or no key available for UTXO's")
            return "mock_tx_id"
            
        def get_transaction_details(self, txid):
            return {
                "status": "pending",
                "confirmations": 0,
                "block_height": None
            }
            
        def estimate_fee(self, size_bytes):
            return {
                "high": 5000,
                "medium": 2000,
                "low": 1000
            }


class WalletManager:
    def __init__(self, network='bitcoin'):
        self.network = BitcoinNetwork(network)
        self.wallets = {}
        self.load_wallets()
    
    def load_wallets(self):
        """Load wallets from storage"""
        try:
            with open('wallets.json', 'r') as f:
                self.wallets = json.load(f)
        except FileNotFoundError:
            self.wallets = {}
    
    def save_wallets(self):
        """Save wallets to storage"""
        with open('wallets.json', 'w') as f:
            json.dump(self.wallets, f, indent=4)
    
    def create_wallet(self, label: str = None) -> dict:
        """Create a new wallet with optional label"""
        private_key = secrets.token_bytes(32)
        signing_key = ecdsa.SigningKey.from_string(private_key, curve=ecdsa.SECP256k1)
        verifying_key = signing_key.get_verifying_key()
        
        # Generate Bitcoin address
        public_key = b'\x04' + verifying_key.to_string()
        sha256_hash = hashlib.sha256(public_key).digest()
        ripemd160_hash = hashlib.new('ripemd160', sha256_hash).digest()
        version = b'\x00' if self.network.network == 'bitcoin' else b'\x6f'
        hash_with_version = version + ripemd160_hash
        double_hash = hashlib.sha256(hashlib.sha256(hash_with_version).digest()).digest()
        checksum = double_hash[:4]
        binary_address = hash_with_version + checksum
        address = base58.b58encode(binary_address).decode()
        
        wallet = {
            'address': address,
            'private_key': private_key.hex(),
            'label': label,
            'created_at': int(time.time())
        }
        
        self.wallets[address] = wallet
        self.save_wallets()
        return wallet
    
    def get_wallet(self, address: str) -> Optional[dict]:
        """Get wallet details by address"""
        return self.wallets.get(address)
    
    def list_wallets(self) -> List[dict]:
        """List all wallets"""
        return list(self.wallets.values())
    
    def get_balance(self, address: str) -> Decimal:
        """Get wallet balance"""
        if address not in self.wallets:
            raise ValueError("Wallet not found")
        return self.network.get_balance(address)
    
    def send_transaction(self, from_address: str, to_address: str, amount: Decimal, fee_rate: int = None) -> str:
        """Send a transaction"""
        wallet = self.get_wallet(from_address)
        if not wallet:
            raise ValueError("Source wallet not found")
            
        return self.network.create_transaction(
            from_address=from_address,
            to_address=to_address,
            amount=amount,
            private_key=bytes.fromhex(wallet['private_key']),
            fee_rate=fee_rate
        )

@dataclass
class WalletKeys:
    private_key: str
    public_key: str
    wif_private_key: str  # Wallet Import Format
    address: str


class BTCWallet:
    def __init__(self, wallet_file: str = 'my_wallet.json', private_key: Optional[str] = None, network: str = 'bitcoin'):
        """Initialize real Bitcoin wallet"""
        self.wallet_file = wallet_file
        
        # Initialize real Bitcoin network connection
        self.network = BitcoinNetwork(network=network)
        self.service = Service(network=network)
        
        # Initialize reward processor for mining rewards
        self.reward_processor = BitcoinRewardProcessor(network=network)
        
        # Try to load existing wallet or create new one
        if private_key:
            self.keys = self._setup_wallet(private_key)
            self.wallet_data = self._create_default_wallet_data()
            self._save_wallet()
        elif os.path.exists(wallet_file):
            self.keys = self._load_wallet()
            self.wallet_data = self._load_wallet_data()
        else:
            self.keys = self._setup_wallet()
            self.wallet_data = self._create_default_wallet_data()
            self._save_wallet()
            
        # Initialize real network balance tracking
        self._update_network_balance()

    def _update_network_balance(self):
        """Update wallet balance from network"""
        try:
            # Get balance from network
            balance = self.network.get_balance(self.keys.address)
            # Update wallet data
            self.wallet_data['balance'] = str(balance)  # Store as string to preserve decimal precision
            self._save_wallet()
        except Exception as e:
            print(f"Warning: Could not update network balance: {e}")
            
    def get_network_balance(self) -> Decimal:
        """Get current balance from network"""
        try:
            return self.network.get_balance(self.keys.address)
        except Exception as e:
            print(f"Warning: Could not fetch network balance: {e}")
            return Decimal('0')
            
    def _create_default_wallet_data(self) -> Dict[str, Any]:
        """Create default wallet data structure"""
        return {
            'private_key': self.keys.private_key,
            'public_key': self.keys.public_key,
            'wif_private_key': self.keys.wif_private_key,
            'address': self.keys.address,
            'balance': 0,
            'total_mined': 0,
            'mining_stats': {
                'last_mining_session': None,
                'total_blocks_mined': 0,
                'last_reward': None,
                'last_hash_rate': None
            },
            'transactions': [],
            'metadata': {
                'created_at': time.time(),
                'last_updated': time.time()
            }
        }

    def _load_wallet_data(self) -> Dict[str, Any]:
        """Load wallet data from file"""
        try:
            with open(self.wallet_file, 'r') as f:
                return json.load(f)
        except Exception as e:
            print(f"Warning: Could not load wallet data: {e}")
            return self._create_default_wallet_data()

    def _save_wallet(self):
        """Save complete wallet info to file"""
        # Update last_updated timestamp
        self.wallet_data['metadata']['last_updated'] = time.time()
        
        # Save to file
        with open(self.wallet_file, 'w') as f:
            json.dump(self.wallet_data, f, indent=2)
            
    def _load_wallet(self) -> WalletKeys:
        """Load complete wallet from file"""
        with open(self.wallet_file, 'r') as f:
            data = json.load(f)
            return WalletKeys(
                private_key=data['private_key'],
                public_key=data['public_key'],
                wif_private_key=data['wif_private_key'],
                address=data['address']
            )
            
    def _setup_wallet(self, private_key: Optional[str] = None) -> WalletKeys:
        """Set up a new wallet or import from private key"""
        if private_key is None:
            # Generate new private key
            private_key = secrets.token_hex(32)
        
        # Validate private key
        if not self._validate_private_key(private_key):
            raise ValueError("Invalid private key format")
            
        # Create signing key from private key
        signing_key = ecdsa.SigningKey.from_string(
            bytes.fromhex(private_key), 
            curve=ecdsa.SECP256k1
        )
        
        # Get verifying key (public key)
        verifying_key = signing_key.get_verifying_key()
        public_key_bytes = verifying_key.to_string()
        
        # Generate compressed public key
        prefix = b'\x02' if public_key_bytes[-1] % 2 == 0 else b'\x03'
        compressed_public_key = prefix + public_key_bytes[:32]
        
        # Create Bitcoin address
        address = self._generate_address(compressed_public_key)
        
        # Create WIF private key
        wif_private_key = self._create_wif(bytes.fromhex(private_key))
        
        return WalletKeys(
            private_key=private_key,
            public_key=compressed_public_key.hex(),
            wif_private_key=wif_private_key,
            address=address
        )

    def _validate_private_key(self, key: str) -> bool:
        """Validate private key format and value"""
        try:
            # Check if it's a valid hex string of correct length
            if len(key) != 64:
                return False
                
            # Check if within valid range (SECP256k1 curve order)
            value = int(key, 16)
            return 0 < value < 0xFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFEBAAEDCE6AF48A03BBFD25E8CD0364141
        except (ValueError, TypeError):
            return False

    def _generate_address(self, public_key_bytes: bytes) -> str:
        """Generate Bitcoin address from public key"""
        # SHA-256 of the public key
        sha256_hash = hashlib.sha256(public_key_bytes).digest()
        
        # RIPEMD-160 of the SHA-256
        ripemd160_hash = hashlib.new('ripemd160', sha256_hash).digest()
        
        # Add version byte (0x00 for mainnet)
        version_ripemd160_hash = b'\x00' + ripemd160_hash
        
        # Double SHA-256 for checksum
        double_sha256 = hashlib.sha256(hashlib.sha256(version_ripemd160_hash).digest()).digest()
        
        # First 4 bytes are the checksum
        checksum = double_sha256[:4]
        
        # Add checksum to versioned RIPEMD-160 hash
        binary_address = version_ripemd160_hash + checksum
        
        # Encode in base58
        address = base58.b58encode(binary_address).decode('utf-8')
        
        return address

    def _create_wif(self, private_key_bytes: bytes) -> str:
        """Create Wallet Import Format (WIF) private key"""
        # Add version byte (0x80 for mainnet)
        version = b'\x80'
        extended = version + private_key_bytes
        
        # Add compression byte
        extended += b'\x01'  # compressed public key marker
        
        # Double SHA-256 for checksum
        double_sha256 = hashlib.sha256(hashlib.sha256(extended).digest()).digest()
        
        # First 4 bytes are the checksum
        checksum = double_sha256[:4]
        
        # Combine everything and encode in base58
        final_key = extended + checksum
        wif = base58.b58encode(final_key).decode('utf-8')
        
        return wif

    def process_mining_reward(self, amount: float, block_hash: str):
        """Process a mining reward"""
        # Update wallet data
        self.wallet_data['total_mined'] += amount
        self.wallet_data['balance'] += amount
        self.wallet_data['mining_stats']['last_mining_session'] = time.time()
        self.wallet_data['mining_stats']['last_reward'] = amount
        self.wallet_data['mining_stats']['total_blocks_mined'] += 1
        
        # Record transaction
        self.wallet_data['transactions'].append({
            'tx_hash': block_hash,
            'type': 'mining_reward',
            'amount': amount,
            'block_hash': block_hash,
            'timestamp': time.time()
        })
        
        # Save changes
        self._save_wallet()

    def get_balance(self) -> Decimal:
        """Get wallet balance (mined BTC)"""
        return Decimal(str(self.wallet_data.get('balance', 0)))
        
    def send_bitcoin(self, to_address: str, amount: Decimal, fee_rate: Optional[int] = None) -> Optional[str]:
        """
        Send Bitcoin to another address
        Args:
            to_address: Destination Bitcoin address
            amount: Amount in BTC
            fee_rate: Optional fee rate in satoshis/byte
        Returns:
            Transaction ID if successful
        """
        try:
            # Create and broadcast transaction using bitcoinlib
            txid = self.network.create_transaction(
                from_address=self.keys.address,
                to_address=to_address,
                amount=amount,
                private_key=self.keys.private_key,
                fee_rate=fee_rate
            )
            
            if txid:
                # Update local balance
                self.wallet_data['balance'] -= float(amount)
                
                # Record transaction
                self.wallet_data['transactions'].append({
                    'tx_hash': txid,
                    'type': 'sent',
                    'amount': float(amount),
                    'to_address': to_address,
                    'timestamp': time.time()
                })
                
                # Save changes
                self._save_wallet()
                
            return txid
            
        except Exception as e:
            raise ValueError(f"Error creating transaction: {str(e)}")
        
    def get_transaction_details(self, txid: str) -> Optional[Dict]:
        """Get details of a transaction"""
        return self.network.get_transaction_details(txid)
        
    def estimate_transfer_fees(self, size_bytes: int = 250) -> Dict[str, int]:
        """Get estimated transfer fees"""
        return self.network.estimate_fee(size_bytes)
        
    def get_wallet_info(self) -> Dict[str, Any]:
        """Get complete wallet information"""
        return {
            'address': self.keys.address,
            'balance': self.get_balance(),
            'total_mined': self.wallet_data['total_mined'],
            'mining_stats': self.wallet_data['mining_stats'],
            'transaction_count': len(self.wallet_data['transactions'])
        }


def load_or_create_wallet(wallet_file: str = 'my_wallet.json') -> BTCWallet:
    """Create a new wallet or load existing one"""
    try:
        return BTCWallet(wallet_file=wallet_file)
    except Exception as e:
        print(f"Error loading wallet: {str(e)}")
        return None