"""Transaction building and signing for BotVegas escrow operations.

This module handles:
- Building Solana transactions from instruction data
- Signing with local keypairs
- Submitting to the network
- ATA (Associated Token Account) management

Designed to work with the instruction data returned by the BotVegas backend's
deposit-instruction and withdraw-instruction endpoints.
"""

import logging
from dataclasses import dataclass
from typing import Any
from uuid import UUID

# Import wallet utilities
from .wallet import (
    CONFIG as WALLET_CONFIG,
    retrieve_keypair,
    get_address_from_keypair,
    get_wallet_balances,
)

logger = logging.getLogger(__name__)

# =============================================================================
# Optional Imports (graceful degradation)
# =============================================================================

try:
    from solders.keypair import Keypair
    from solders.pubkey import Pubkey
    from solders.instruction import Instruction, AccountMeta
    from solders.transaction import Transaction
    from solders.message import Message
    from solders.hash import Hash
    from solders.system_program import ID as SYSTEM_PROGRAM_ID
    from solders.signature import Signature
    
    SOLDERS_AVAILABLE = True
except ImportError:
    SOLDERS_AVAILABLE = False
    logger.warning("solders not available - transaction signing disabled")

try:
    from solana.rpc.api import Client as SolanaClient
    from solana.rpc.commitment import Confirmed, Finalized
    from solana.rpc.types import TxOpts
    
    SOLANA_RPC_AVAILABLE = True
except ImportError:
    SOLANA_RPC_AVAILABLE = False
    logger.warning("solana-py not available - RPC operations disabled")

try:
    from spl.token.constants import TOKEN_PROGRAM_ID, ASSOCIATED_TOKEN_PROGRAM_ID
    from spl.token.instructions import get_associated_token_address
    
    SPL_TOKEN_AVAILABLE = True
except ImportError:
    try:
        # Alternative import path in some solana-py versions
        from solana.spl.token.constants import TOKEN_PROGRAM_ID, ASSOCIATED_TOKEN_PROGRAM_ID
        from solana.spl.token.instructions import get_associated_token_address
        
        SPL_TOKEN_AVAILABLE = True
    except ImportError:
        SPL_TOKEN_AVAILABLE = False
        logger.warning("spl-token not available - will use manual ATA derivation")


# =============================================================================
# Constants
# =============================================================================

# SPL Token Program ID (if spl-token not available)
TOKEN_PROGRAM_ID_STR = "TokenkegQfeZyiNwAJbNbGKPFXCWuBvf9Ss623VQ5DA"
ASSOCIATED_TOKEN_PROGRAM_ID_STR = "ATokenGPvbdGVxr1b2hvZbsiqW5xWH25efTNsLJA8knL"


# =============================================================================
# ATA (Associated Token Account) Utilities
# =============================================================================


def get_ata_address(wallet_address: str, mint_address: str) -> str:
    """Get the Associated Token Account address for a wallet and mint.
    
    Args:
        wallet_address: Owner's wallet address (base58)
        mint_address: Token mint address (base58)
        
    Returns:
        ATA address (base58)
    """
    if SPL_TOKEN_AVAILABLE and SOLDERS_AVAILABLE:
        try:
            owner = Pubkey.from_string(wallet_address)
            mint = Pubkey.from_string(mint_address)
            ata = get_associated_token_address(owner, mint)
            return str(ata)
        except Exception as e:
            logger.warning(f"spl-token ATA derivation failed: {e}")
    
    # Manual derivation using PDA
    if SOLDERS_AVAILABLE:
        try:
            owner = Pubkey.from_string(wallet_address)
            mint = Pubkey.from_string(mint_address)
            token_program = Pubkey.from_string(TOKEN_PROGRAM_ID_STR)
            ata_program = Pubkey.from_string(ASSOCIATED_TOKEN_PROGRAM_ID_STR)
            
            # PDA seeds: [owner, token_program, mint]
            seeds = [bytes(owner), bytes(token_program), bytes(mint)]
            ata, _ = Pubkey.find_program_address(seeds, ata_program)
            return str(ata)
        except Exception as e:
            logger.error(f"Manual ATA derivation failed: {e}")
            raise ValueError(f"Cannot derive ATA: {e}")
    
    raise ValueError("Cannot derive ATA: solders not available")


def check_ata_exists(ata_address: str, rpc_url: str) -> bool:
    """Check if an ATA exists on-chain.
    
    Args:
        ata_address: The ATA address to check
        rpc_url: Solana RPC endpoint
        
    Returns:
        True if account exists and is initialized
    """
    if not SOLANA_RPC_AVAILABLE:
        logger.warning("Cannot check ATA existence: solana-py not available")
        return False
    
    try:
        client = SolanaClient(rpc_url)
        pubkey = Pubkey.from_string(ata_address) if SOLDERS_AVAILABLE else ata_address
        response = client.get_account_info(pubkey, commitment=Confirmed)
        return response.value is not None
    except Exception as e:
        logger.warning(f"ATA existence check failed: {e}")
        return False


# =============================================================================
# Transaction Building
# =============================================================================


@dataclass
class TransactionResult:
    """Result of a transaction submission."""
    success: bool
    signature: str | None
    error: str | None
    solscan_url: str | None
    
    def to_dict(self) -> dict[str, Any]:
        return {
            "success": self.success,
            "signature": self.signature,
            "error": self.error,
            "solscan_url": self.solscan_url,
        }


def build_anchor_instruction(
    program_id: str,
    instruction_name: str,
    accounts: dict[str, str],
    data: dict[str, Any],
    signer_pubkey: str,
) -> "Instruction":
    """Build an Anchor instruction from the backend's instruction data.
    
    Args:
        program_id: The Anchor program address
        instruction_name: "deposit" or "withdraw"
        accounts: Account addresses from the backend
        data: Instruction data from the backend
        signer_pubkey: The signing wallet's pubkey
        
    Returns:
        A solders Instruction object
    """
    if not SOLDERS_AVAILABLE:
        raise ValueError("Cannot build instruction: solders not available")
    
    program = Pubkey.from_string(program_id)
    
    # Build account metas based on instruction type
    # Order matters! Must match the Anchor IDL account order
    
    if instruction_name == "deposit":
        # Deposit accounts (order from Anchor program):
        # 1. game_escrow (writable)
        # 2. token_vault (writable) 
        # 3. player_deposit (writable, init if needed)
        # 4. player (signer)
        # 5. player_token_account (writable)
        # 6. token_program
        # 7. system_program
        account_metas = [
            AccountMeta(Pubkey.from_string(accounts["game_escrow"]), is_signer=False, is_writable=True),
            AccountMeta(Pubkey.from_string(accounts["token_vault"]), is_signer=False, is_writable=True),
            AccountMeta(Pubkey.from_string(accounts["player_deposit"]), is_signer=False, is_writable=True),
            AccountMeta(Pubkey.from_string(signer_pubkey), is_signer=True, is_writable=True),
            AccountMeta(Pubkey.from_string(accounts["player_token_account"]), is_signer=False, is_writable=True),
            AccountMeta(Pubkey.from_string(TOKEN_PROGRAM_ID_STR), is_signer=False, is_writable=False),
            AccountMeta(SYSTEM_PROGRAM_ID, is_signer=False, is_writable=False),
        ]
        
    elif instruction_name == "withdraw":
        # Withdraw accounts (order from Anchor program):
        # 1. config
        # 2. game_escrow (writable)
        # 3. token_vault (writable)
        # 4. player_deposit (writable)
        # 5. player (signer)
        # 6. player_token_account (writable)
        # 7. token_program
        account_metas = [
            AccountMeta(Pubkey.from_string(accounts["config"]), is_signer=False, is_writable=False),
            AccountMeta(Pubkey.from_string(accounts["game_escrow"]), is_signer=False, is_writable=True),
            AccountMeta(Pubkey.from_string(accounts["token_vault"]), is_signer=False, is_writable=True),
            AccountMeta(Pubkey.from_string(accounts["player_deposit"]), is_signer=False, is_writable=True),
            AccountMeta(Pubkey.from_string(signer_pubkey), is_signer=True, is_writable=True),
            AccountMeta(Pubkey.from_string(accounts["player_token_account"]), is_signer=False, is_writable=True),
            AccountMeta(Pubkey.from_string(TOKEN_PROGRAM_ID_STR), is_signer=False, is_writable=False),
        ]
    else:
        raise ValueError(f"Unknown instruction: {instruction_name}")
    
    # Build instruction data
    # Anchor uses an 8-byte discriminator based on instruction name
    import hashlib
    discriminator = hashlib.sha256(f"global:{instruction_name}".encode()).digest()[:8]
    
    # For deposit/withdraw, we don't need additional data beyond the discriminator
    # The accounts contain all the necessary information
    instruction_data = discriminator
    
    return Instruction(
        program_id=program,
        accounts=account_metas,
        data=instruction_data,
    )


def build_create_ata_instruction(
    payer: str,
    owner: str,
    mint: str,
) -> "Instruction":
    """Build an instruction to create an Associated Token Account.
    
    Args:
        payer: Address paying for account creation
        owner: Owner of the new ATA
        mint: Token mint address
        
    Returns:
        A solders Instruction object
    """
    if not SOLDERS_AVAILABLE:
        raise ValueError("Cannot build ATA instruction: solders not available")
    
    payer_pubkey = Pubkey.from_string(payer)
    owner_pubkey = Pubkey.from_string(owner)
    mint_pubkey = Pubkey.from_string(mint)
    token_program = Pubkey.from_string(TOKEN_PROGRAM_ID_STR)
    ata_program = Pubkey.from_string(ASSOCIATED_TOKEN_PROGRAM_ID_STR)
    
    # Derive ATA address
    ata_address = get_ata_address(owner, mint)
    ata_pubkey = Pubkey.from_string(ata_address)
    
    # Create ATA instruction accounts:
    # 1. payer (signer, writable)
    # 2. ata (writable)
    # 3. owner
    # 4. mint
    # 5. system_program
    # 6. token_program
    account_metas = [
        AccountMeta(payer_pubkey, is_signer=True, is_writable=True),
        AccountMeta(ata_pubkey, is_signer=False, is_writable=True),
        AccountMeta(owner_pubkey, is_signer=False, is_writable=False),
        AccountMeta(mint_pubkey, is_signer=False, is_writable=False),
        AccountMeta(SYSTEM_PROGRAM_ID, is_signer=False, is_writable=False),
        AccountMeta(token_program, is_signer=False, is_writable=False),
    ]
    
    # ATA program create instruction has no data
    return Instruction(
        program_id=ata_program,
        accounts=account_metas,
        data=bytes(),
    )


# =============================================================================
# Transaction Signing and Submission
# =============================================================================


async def sign_and_submit_transaction(
    instructions: list["Instruction"],
    keypair_bytes: list[int],
    rpc_url: str,
) -> TransactionResult:
    """Sign and submit a transaction to the Solana network.
    
    Args:
        instructions: List of instructions to include
        keypair_bytes: The 64-byte keypair for signing
        rpc_url: Solana RPC endpoint
        
    Returns:
        TransactionResult with signature or error
    """
    if not SOLDERS_AVAILABLE:
        return TransactionResult(
            success=False,
            signature=None,
            error="solders not available - cannot sign transactions",
            solscan_url=None,
        )
    
    if not SOLANA_RPC_AVAILABLE:
        return TransactionResult(
            success=False,
            signature=None,
            error="solana-py not available - cannot submit transactions",
            solscan_url=None,
        )
    
    try:
        # Create keypair from bytes
        keypair = Keypair.from_bytes(bytes(keypair_bytes))
        
        # Get recent blockhash
        client = SolanaClient(rpc_url)
        blockhash_response = client.get_latest_blockhash(commitment=Confirmed)
        recent_blockhash = blockhash_response.value.blockhash
        
        # Build message
        message = Message.new_with_blockhash(
            instructions,
            keypair.pubkey(),
            recent_blockhash,
        )
        
        # Sign transaction
        transaction = Transaction.new_unsigned(message)
        transaction.sign([keypair], recent_blockhash)
        
        # Submit transaction
        opts = TxOpts(skip_preflight=False, preflight_commitment=Confirmed)
        response = client.send_transaction(transaction, opts=opts)
        
        signature = str(response.value)
        
        # Wait for confirmation
        client.confirm_transaction(signature, commitment=Confirmed)
        
        return TransactionResult(
            success=True,
            signature=signature,
            error=None,
            solscan_url=f"https://solscan.io/tx/{signature}",
        )
        
    except Exception as e:
        logger.error(f"Transaction failed: {e}")
        return TransactionResult(
            success=False,
            signature=None,
            error=str(e),
            solscan_url=None,
        )


# =============================================================================
# High-Level Operations
# =============================================================================


@dataclass 
class DepositResult:
    """Result of a deposit operation."""
    success: bool
    signature: str | None
    amount_deposited: float
    wallet_address: str
    game_id: str
    error: str | None
    solscan_url: str | None
    new_sol_balance: float | None
    new_botvegas_balance: float | None
    
    def to_dict(self) -> dict[str, Any]:
        return {
            "success": self.success,
            "signature": self.signature,
            "amount_deposited": self.amount_deposited,
            "wallet_address": self.wallet_address,
            "game_id": self.game_id,
            "error": self.error,
            "solscan_url": self.solscan_url,
            "new_sol_balance": self.new_sol_balance,
            "new_botvegas_balance": self.new_botvegas_balance,
        }


@dataclass
class WithdrawResult:
    """Result of a withdraw operation."""
    success: bool
    signature: str | None
    amount_withdrawn: float
    wallet_address: str
    game_id: str
    error: str | None
    solscan_url: str | None
    new_sol_balance: float | None
    new_botvegas_balance: float | None
    
    def to_dict(self) -> dict[str, Any]:
        return {
            "success": self.success,
            "signature": self.signature,
            "amount_withdrawn": self.amount_withdrawn,
            "wallet_address": self.wallet_address,
            "game_id": self.game_id,
            "error": self.error,
            "solscan_url": self.solscan_url,
            "new_sol_balance": self.new_sol_balance,
            "new_botvegas_balance": self.new_botvegas_balance,
        }


async def execute_deposit(
    identity_name: str,
    instruction_data: dict[str, Any],
    buy_in_amount: float,
    rpc_url: str | None = None,
) -> DepositResult:
    """Execute a deposit transaction.
    
    Args:
        identity_name: Identity whose wallet to use
        instruction_data: Response from get_deposit_instruction endpoint
        buy_in_amount: Amount being deposited (for logging)
        rpc_url: Solana RPC endpoint
        
    Returns:
        DepositResult with transaction details
    """
    rpc_url = rpc_url or WALLET_CONFIG.RPC_MAINNET
    
    # Get keypair
    keypair_bytes = retrieve_keypair(identity_name)
    if keypair_bytes is None:
        return DepositResult(
            success=False,
            signature=None,
            amount_deposited=0,
            wallet_address="",
            game_id=instruction_data.get("data", {}).get("game_id", ""),
            error=f"No wallet found for identity '{identity_name}'",
            solscan_url=None,
            new_sol_balance=None,
            new_botvegas_balance=None,
        )
    
    wallet_address = get_address_from_keypair(keypair_bytes)
    game_id = instruction_data.get("data", {}).get("game_id", "")
    
    # Check balances first
    sol_balance, botvegas_balance = get_wallet_balances(wallet_address, rpc_url)
    
    if sol_balance is None or sol_balance < WALLET_CONFIG.MIN_SOL_FOR_FEES:
        return DepositResult(
            success=False,
            signature=None,
            amount_deposited=0,
            wallet_address=wallet_address,
            game_id=game_id,
            error=f"Insufficient SOL for fees. Have {sol_balance or 0:.4f}, need at least {WALLET_CONFIG.MIN_SOL_FOR_FEES}",
            solscan_url=None,
            new_sol_balance=sol_balance,
            new_botvegas_balance=botvegas_balance,
        )
    
    if botvegas_balance is None or botvegas_balance < buy_in_amount:
        return DepositResult(
            success=False,
            signature=None,
            amount_deposited=0,
            wallet_address=wallet_address,
            game_id=game_id,
            error=f"Insufficient $BOTVEGAS. Have {botvegas_balance or 0:.2f}, need {buy_in_amount}",
            solscan_url=None,
            new_sol_balance=sol_balance,
            new_botvegas_balance=botvegas_balance,
        )
    
    # Get player's ATA for $BOTVEGAS
    try:
        player_ata = get_ata_address(wallet_address, WALLET_CONFIG.BOTVEGAS_MINT)
    except Exception as e:
        return DepositResult(
            success=False,
            signature=None,
            amount_deposited=0,
            wallet_address=wallet_address,
            game_id=game_id,
            error=f"Failed to derive ATA: {e}",
            solscan_url=None,
            new_sol_balance=sol_balance,
            new_botvegas_balance=botvegas_balance,
        )
    
    # Update accounts with player's ATA
    accounts = instruction_data["accounts"].copy()
    accounts["player_token_account"] = player_ata
    
    # Build instructions
    instructions = []
    
    # Check if ATA exists, create if not
    if not check_ata_exists(player_ata, rpc_url):
        logger.info(f"Creating ATA for {wallet_address}")
        try:
            create_ata_ix = build_create_ata_instruction(
                payer=wallet_address,
                owner=wallet_address,
                mint=WALLET_CONFIG.BOTVEGAS_MINT,
            )
            instructions.append(create_ata_ix)
        except Exception as e:
            return DepositResult(
                success=False,
                signature=None,
                amount_deposited=0,
                wallet_address=wallet_address,
                game_id=game_id,
                error=f"Failed to build ATA instruction: {e}",
                solscan_url=None,
                new_sol_balance=sol_balance,
                new_botvegas_balance=botvegas_balance,
            )
    
    # Build deposit instruction
    try:
        deposit_ix = build_anchor_instruction(
            program_id=instruction_data["program_id"],
            instruction_name=instruction_data["instruction"],
            accounts=accounts,
            data=instruction_data["data"],
            signer_pubkey=wallet_address,
        )
        instructions.append(deposit_ix)
    except Exception as e:
        return DepositResult(
            success=False,
            signature=None,
            amount_deposited=0,
            wallet_address=wallet_address,
            game_id=game_id,
            error=f"Failed to build deposit instruction: {e}",
            solscan_url=None,
            new_sol_balance=sol_balance,
            new_botvegas_balance=botvegas_balance,
        )
    
    # Sign and submit
    result = await sign_and_submit_transaction(
        instructions=instructions,
        keypair_bytes=keypair_bytes,
        rpc_url=rpc_url,
    )
    
    if not result.success:
        return DepositResult(
            success=False,
            signature=None,
            amount_deposited=0,
            wallet_address=wallet_address,
            game_id=game_id,
            error=result.error,
            solscan_url=None,
            new_sol_balance=sol_balance,
            new_botvegas_balance=botvegas_balance,
        )
    
    # Get updated balances
    new_sol, new_botvegas = get_wallet_balances(wallet_address, rpc_url)
    
    return DepositResult(
        success=True,
        signature=result.signature,
        amount_deposited=buy_in_amount,
        wallet_address=wallet_address,
        game_id=game_id,
        error=None,
        solscan_url=result.solscan_url,
        new_sol_balance=new_sol,
        new_botvegas_balance=new_botvegas,
    )


async def execute_withdraw(
    identity_name: str,
    instruction_data: dict[str, Any],
    rpc_url: str | None = None,
) -> WithdrawResult:
    """Execute a withdraw transaction.
    
    Args:
        identity_name: Identity whose wallet to use
        instruction_data: Response from get_withdraw_instruction endpoint
        rpc_url: Solana RPC endpoint
        
    Returns:
        WithdrawResult with transaction details
    """
    rpc_url = rpc_url or WALLET_CONFIG.RPC_MAINNET
    
    # Get keypair
    keypair_bytes = retrieve_keypair(identity_name)
    if keypair_bytes is None:
        return WithdrawResult(
            success=False,
            signature=None,
            amount_withdrawn=0,
            wallet_address="",
            game_id=instruction_data.get("data", {}).get("game_id", ""),
            error=f"No wallet found for identity '{identity_name}'",
            solscan_url=None,
            new_sol_balance=None,
            new_botvegas_balance=None,
        )
    
    wallet_address = get_address_from_keypair(keypair_bytes)
    game_id = instruction_data.get("data", {}).get("game_id", "")
    
    # Check SOL balance for fees
    sol_balance, botvegas_balance = get_wallet_balances(wallet_address, rpc_url)
    
    if sol_balance is None or sol_balance < WALLET_CONFIG.MIN_SOL_FOR_FEES:
        return WithdrawResult(
            success=False,
            signature=None,
            amount_withdrawn=0,
            wallet_address=wallet_address,
            game_id=game_id,
            error=f"Insufficient SOL for fees. Have {sol_balance or 0:.4f}, need at least {WALLET_CONFIG.MIN_SOL_FOR_FEES}",
            solscan_url=None,
            new_sol_balance=sol_balance,
            new_botvegas_balance=botvegas_balance,
        )
    
    # Get player's ATA for $BOTVEGAS
    try:
        player_ata = get_ata_address(wallet_address, WALLET_CONFIG.BOTVEGAS_MINT)
    except Exception as e:
        return WithdrawResult(
            success=False,
            signature=None,
            amount_withdrawn=0,
            wallet_address=wallet_address,
            game_id=game_id,
            error=f"Failed to derive ATA: {e}",
            solscan_url=None,
            new_sol_balance=sol_balance,
            new_botvegas_balance=botvegas_balance,
        )
    
    # Update accounts with player's ATA
    accounts = instruction_data["accounts"].copy()
    accounts["player_token_account"] = player_ata
    
    # Build instructions
    instructions = []
    
    # Check if ATA exists, create if not (shouldn't happen for withdraw but be safe)
    if not check_ata_exists(player_ata, rpc_url):
        logger.info(f"Creating ATA for withdraw to {wallet_address}")
        try:
            create_ata_ix = build_create_ata_instruction(
                payer=wallet_address,
                owner=wallet_address,
                mint=WALLET_CONFIG.BOTVEGAS_MINT,
            )
            instructions.append(create_ata_ix)
        except Exception as e:
            return WithdrawResult(
                success=False,
                signature=None,
                amount_withdrawn=0,
                wallet_address=wallet_address,
                game_id=game_id,
                error=f"Failed to build ATA instruction: {e}",
                solscan_url=None,
                new_sol_balance=sol_balance,
                new_botvegas_balance=botvegas_balance,
            )
    
    # Build withdraw instruction
    try:
        withdraw_ix = build_anchor_instruction(
            program_id=instruction_data["program_id"],
            instruction_name=instruction_data["instruction"],
            accounts=accounts,
            data=instruction_data["data"],
            signer_pubkey=wallet_address,
        )
        instructions.append(withdraw_ix)
    except Exception as e:
        return WithdrawResult(
            success=False,
            signature=None,
            amount_withdrawn=0,
            wallet_address=wallet_address,
            game_id=game_id,
            error=f"Failed to build withdraw instruction: {e}",
            solscan_url=None,
            new_sol_balance=sol_balance,
            new_botvegas_balance=botvegas_balance,
        )
    
    # Sign and submit
    result = await sign_and_submit_transaction(
        instructions=instructions,
        keypair_bytes=keypair_bytes,
        rpc_url=rpc_url,
    )
    
    if not result.success:
        return WithdrawResult(
            success=False,
            signature=None,
            amount_withdrawn=0,
            wallet_address=wallet_address,
            game_id=game_id,
            error=result.error,
            solscan_url=None,
            new_sol_balance=sol_balance,
            new_botvegas_balance=botvegas_balance,
        )
    
    # Get updated balances
    new_sol, new_botvegas = get_wallet_balances(wallet_address, rpc_url)
    
    # Calculate withdrawn amount (difference in balance)
    amount_withdrawn = (new_botvegas or 0) - (botvegas_balance or 0)
    
    return WithdrawResult(
        success=True,
        signature=result.signature,
        amount_withdrawn=amount_withdrawn,
        wallet_address=wallet_address,
        game_id=game_id,
        error=None,
        solscan_url=result.solscan_url,
        new_sol_balance=new_sol,
        new_botvegas_balance=new_botvegas,
    )
