"""
DFlow Swap tool for Solana wallets.

Enables fast token swaps using DFlow's Swap API with a Solana keypair.
DFlow offers faster swaps compared to Jupiter Ultra with similar liquidity
and supports platform fees for monetization.
"""

import base64
import logging
from typing import Dict, Any, List, Optional

from solana_agent import AutoTool, ToolRegistry
from solders.keypair import Keypair  # type: ignore
from solders.transaction import VersionedTransaction  # type: ignore
from solders.message import to_bytes_versioned  # type: ignore
from solana.rpc.async_api import AsyncClient  # type: ignore
from solana.rpc.commitment import Confirmed  # type: ignore

from sakit.utils.dflow import DFlowSwap

logger = logging.getLogger(__name__)

# Default Solana RPC endpoint
DEFAULT_RPC_URL = "https://api.mainnet-beta.solana.com"


def _sign_dflow_transaction(
    transaction_base64: str,
    sign_message_func,
    payer_sign_func=None,
) -> str:
    """Sign a DFlow transaction with the user's keypair (and optionally a payer).

    Args:
        transaction_base64: Base64-encoded transaction from DFlow
        sign_message_func: Function to sign message bytes (user's keypair.sign_message)
        payer_sign_func: Optional function to sign message bytes (payer's keypair.sign_message)

    Returns:
        Base64-encoded signed transaction
    """
    tx_bytes = base64.b64decode(transaction_base64)
    transaction = VersionedTransaction.from_bytes(tx_bytes)
    message_bytes = to_bytes_versioned(transaction.message)

    if payer_sign_func:
        # Payer signs first, then user
        payer_signature = payer_sign_func(message_bytes)
        user_signature = sign_message_func(message_bytes)
        signed_tx = VersionedTransaction.populate(
            transaction.message,
            [payer_signature, user_signature],
        )
    else:
        # Just user signs
        user_signature = sign_message_func(message_bytes)
        signed_tx = VersionedTransaction.populate(
            transaction.message,
            [user_signature],
        )

    return base64.b64encode(bytes(signed_tx)).decode("utf-8")


class SolanaDFlowSwapTool(AutoTool):
    """Fast token swaps using DFlow API with a Solana keypair."""

    def __init__(self, registry: Optional[ToolRegistry] = None):
        super().__init__(
            name="solana_dflow_swap",
            description=(
                "Fast token swap on Solana using DFlow API. "
                "Swaps tokens instantly with competitive rates and low slippage. "
                "Supports SOL and all SPL tokens. Faster alternative to Jupiter Ultra."
            ),
            registry=registry,
        )
        self._private_key: Optional[str] = None
        self._platform_fee_bps: Optional[int] = None
        self._fee_account: Optional[str] = None
        self._referral_account: Optional[str] = None
        self._payer_private_key: Optional[str] = None
        self._rpc_url: Optional[str] = None

    def get_schema(self) -> Dict[str, Any]:
        return {
            "type": "object",
            "properties": {
                "input_mint": {
                    "type": "string",
                    "description": "Token mint address to sell/swap from. Use 'So11111111111111111111111111111111111111112' for native SOL.",
                },
                "output_mint": {
                    "type": "string",
                    "description": "Token mint address to buy/swap to. Use 'So11111111111111111111111111111111111111112' for native SOL.",
                },
                "amount": {
                    "type": "integer",
                    "description": "Amount to swap in the smallest unit (lamports for SOL, base units for tokens). Example: 1000000000 for 1 SOL.",
                },
                "slippage_bps": {
                    "type": "integer",
                    "description": "Maximum slippage tolerance in basis points (100 = 1%). Default is auto.",
                    "default": 0,
                },
            },
            "required": ["input_mint", "output_mint", "amount"],
            "additionalProperties": False,
        }

    def configure(self, config: Dict[str, Any]) -> None:
        super().configure(config)
        tool_cfg = config.get("tools", {}).get("solana_dflow_swap", {})
        self._private_key = tool_cfg.get("private_key")
        self._platform_fee_bps = tool_cfg.get("platform_fee_bps")
        self._fee_account = tool_cfg.get("fee_account")
        self._referral_account = tool_cfg.get("referral_account")
        self._payer_private_key = tool_cfg.get("payer_private_key")
        self._rpc_url = tool_cfg.get("rpc_url") or DEFAULT_RPC_URL

    async def execute(
        self,
        input_mint: str,
        output_mint: str,
        amount: int,
        slippage_bps: int = 0,
    ) -> Dict[str, Any]:
        if not self._private_key:
            return {"status": "error", "message": "Private key not configured."}

        try:
            keypair = Keypair.from_base58_string(self._private_key)
            user_pubkey = str(keypair.pubkey())

            # Check if payer is configured for gasless transactions
            payer_keypair = None
            sponsor = None
            if self._payer_private_key:
                payer_keypair = Keypair.from_base58_string(self._payer_private_key)
                sponsor = str(payer_keypair.pubkey())

            # Initialize DFlow client
            dflow = DFlowSwap()

            # Get order from DFlow
            order_result = await dflow.get_order(
                input_mint=input_mint,
                output_mint=output_mint,
                amount=amount,
                user_public_key=user_pubkey,
                slippage_bps=slippage_bps if slippage_bps > 0 else None,
                platform_fee_bps=self._platform_fee_bps,
                platform_fee_mode="outputMint",
                fee_account=self._fee_account,
                referral_account=self._referral_account,
                sponsor=sponsor,
            )

            if not order_result.success:
                return {"status": "error", "message": order_result.error}

            if not order_result.transaction:
                return {
                    "status": "error",
                    "message": "No transaction returned from DFlow.",
                }

            # Sign the transaction
            signed_tx = _sign_dflow_transaction(
                transaction_base64=order_result.transaction,
                sign_message_func=keypair.sign_message,
                payer_sign_func=payer_keypair.sign_message if payer_keypair else None,
            )

            # Send the transaction to Solana
            signature = await self._send_transaction(signed_tx)

            if not signature:
                return {
                    "status": "error",
                    "message": "Failed to send transaction to Solana.",
                }

            return {
                "status": "success",
                "signature": signature,
                "input_amount": order_result.in_amount,
                "output_amount": order_result.out_amount,
                "min_output_amount": order_result.min_out_amount,
                "input_mint": order_result.input_mint,
                "output_mint": order_result.output_mint,
                "price_impact": order_result.price_impact_pct,
                "platform_fee": order_result.platform_fee,
                "execution_mode": order_result.execution_mode,
                "message": f"Swap successful! Signature: {signature}",
            }

        except Exception as e:
            logger.exception(f"DFlow swap failed: {str(e)}")
            return {"status": "error", "message": str(e)}

    async def _send_transaction(self, signed_tx_base64: str) -> Optional[str]:
        """Send a signed transaction to Solana."""
        try:
            tx_bytes = base64.b64decode(signed_tx_base64)

            async with AsyncClient(self._rpc_url) as client:
                result = await client.send_raw_transaction(
                    tx_bytes,
                    opts={
                        "skip_preflight": True,
                        "preflight_commitment": Confirmed,
                        "max_retries": 3,
                    },
                )

                if result.value:
                    return str(result.value)

                return None
        except Exception as e:
            logger.error(f"Failed to send transaction: {e}")
            return None


class SolanaDFlowSwapPlugin:
    """Plugin for fast token swaps using DFlow API."""

    def __init__(self):
        self.name = "solana_dflow_swap"
        self.config = None
        self.tool_registry = None
        self._tool = None

    @property
    def description(self):
        return "Plugin for fast token swaps using DFlow API with a Solana keypair."

    def initialize(self, tool_registry: ToolRegistry) -> None:
        self.tool_registry = tool_registry
        self._tool = SolanaDFlowSwapTool(registry=tool_registry)

    def configure(self, config: Dict[str, Any]) -> None:
        self.config = config
        if self._tool:
            self._tool.configure(self.config)

    def get_tools(self) -> List[AutoTool]:
        return [self._tool] if self._tool else []


def get_plugin():
    return SolanaDFlowSwapPlugin()
