"""LangChain tools for PowerSun.vip TRON Energy marketplace."""

from __future__ import annotations

import json
from typing import Any, Type

from langchain_core.tools import BaseTool
from pydantic import BaseModel, Field

from langchain_powersun.client import PowerSunClient


# ─── Market Tools (public) ──────────────────────────────────────────────


class GetPricesTool(BaseTool):
    """Get current TRON Energy & Bandwidth prices for all duration tiers (5min to 30d)."""

    name: str = "powersun_get_prices"
    description: str = (
        "Get current TRON Energy & Bandwidth prices from PowerSun.vip marketplace. "
        "Returns price tiers for durations from 5 minutes to 30 days. "
        "No authentication required."
    )
    client: Any = None

    def _run(self, **kwargs: Any) -> str:
        result = self.client.post("/api/v2/agent/estimate", {"txCount": 1})
        return json.dumps(result, indent=2)


class EstimateCostInput(BaseModel):
    tx_count: int = Field(description="Number of TRON transactions to estimate cost for", ge=1)
    tx_type: str = Field(
        default="trc20_transfer",
        description="Transaction type: trc20_transfer, trc20_transfer_new, trc10_transfer, or contract_call",
    )
    duration_minutes: int = Field(
        default=60,
        description="Rental duration in minutes. Allowed: 5, 10, 15, 30, 60, 120, 1440, 10080",
    )


class EstimateCostTool(BaseTool):
    """Estimate the cost of renting TRON Energy for a given number of transactions."""

    name: str = "powersun_estimate_cost"
    description: str = (
        "Calculate the exact cost of renting TRON Energy from PowerSun.vip. "
        "Specify transaction count, type, and duration. Returns cost in TRX/SUN, "
        "savings vs burning TRX, and minimum order info. No authentication required."
    )
    args_schema: Type[BaseModel] = EstimateCostInput
    client: Any = None

    def _run(self, tx_count: int = 1, tx_type: str = "trc20_transfer", duration_minutes: int = 60) -> str:
        result = self.client.post("/api/v2/agent/estimate", {
            "txCount": tx_count,
            "txType": tx_type,
            "durationMinutes": duration_minutes,
        })
        return json.dumps(result, indent=2)


class GetAvailableResourcesTool(BaseTool):
    """Check available TRON Energy and Bandwidth in the PowerSun marketplace."""

    name: str = "powersun_get_available_resources"
    description: str = (
        "Check how much TRON Energy and Bandwidth is currently available "
        "for rent on PowerSun.vip marketplace. No authentication required."
    )
    client: Any = None

    def _run(self, **kwargs: Any) -> str:
        result = self.client.get("/api/v2/agent/resources")
        return json.dumps(result, indent=2)


class GetMarketOverviewTool(BaseTool):
    """Get a full market snapshot from PowerSun.vip."""

    name: str = "powersun_get_market_overview"
    description: str = (
        "Get a comprehensive market snapshot from PowerSun.vip: current prices, "
        "available resources, 24h order stats, and active pool count. "
        "No authentication required."
    )
    client: Any = None

    def _run(self, **kwargs: Any) -> str:
        result = self.client.get("/api/v2/agent/market-overview")
        return json.dumps(result, indent=2)


# ─── Buyer Tools (auth recommended) ─────────────────────────────────────


class BuyEnergyInput(BaseModel):
    target_address: str = Field(description="TRON address to delegate Energy to (T-address, 34 chars)")
    tx_count: int = Field(description="Number of transactions worth of energy to buy", ge=1)
    tx_type: str = Field(
        default="trc20_transfer",
        description="Transaction type: trc20_transfer, trc20_transfer_new, trc10_transfer, or contract_call",
    )
    duration_minutes: int = Field(
        default=60,
        description="Rental duration in minutes. Allowed: 5, 10, 15, 30, 60, 120, 1440, 10080",
    )
    max_price_trx: float | None = Field(
        default=None,
        description="Maximum acceptable total price in TRX. Order rejected if cost exceeds this.",
    )


class BuyEnergyTool(BaseTool):
    """Purchase TRON Energy delegation from PowerSun.vip."""

    name: str = "powersun_buy_energy"
    description: str = (
        "Purchase TRON Energy delegation from PowerSun.vip. Energy is delegated to the "
        "target TRON address within seconds. With API key: deducts from balance. "
        "Without API key: returns HTTP 402 payment instructions (TRX deposit or x402 USDC). "
        "Save 20-50% vs burning TRX. Only 10% commission."
    )
    args_schema: Type[BaseModel] = BuyEnergyInput
    client: Any = None

    def _run(
        self,
        target_address: str,
        tx_count: int = 1,
        tx_type: str = "trc20_transfer",
        duration_minutes: int = 60,
        max_price_trx: float | None = None,
    ) -> str:
        body: dict[str, Any] = {
            "targetAddress": target_address,
            "txCount": tx_count,
            "txType": tx_type,
            "durationMinutes": duration_minutes,
        }
        if max_price_trx is not None:
            body["maxPriceTrx"] = max_price_trx
        result = self.client.post("/api/v2/agent/buy-energy", body)
        return json.dumps(result, indent=2)


class GetBalanceTool(BaseTool):
    """Get PowerSun.vip account balance and deposit info."""

    name: str = "powersun_get_balance"
    description: str = (
        "Get your PowerSun.vip account balance in TRX, deposit address, "
        "and estimated transaction capacity. Requires API key authentication."
    )
    client: Any = None

    def _run(self, **kwargs: Any) -> str:
        result = self.client.get("/api/v2/agent/balance")
        return json.dumps(result, indent=2)


class GetOrderStatusInput(BaseModel):
    order_id: str = Field(description="Order ID (UUID) to check status for")


class GetOrderStatusTool(BaseTool):
    """Get detailed order status from PowerSun.vip."""

    name: str = "powersun_get_order_status"
    description: str = (
        "Get detailed status of an energy purchase order including delegation progress, "
        "fill details, and transaction hashes. Requires API key authentication."
    )
    args_schema: Type[BaseModel] = GetOrderStatusInput
    client: Any = None

    def _run(self, order_id: str = "") -> str:
        result = self.client.get(f"/api/v2/agent/order/{order_id}")
        return json.dumps(result, indent=2)


class BroadcastInput(BaseModel):
    tx_data: str = Field(
        description="JSON string of the pre-signed TRON transaction (with txID, raw_data, signature)"
    )


class BroadcastTransactionTool(BaseTool):
    """Broadcast a pre-signed TRON transaction with automatic Energy delegation."""

    name: str = "powersun_broadcast_transaction"
    description: str = (
        "Broadcast a pre-signed TRON transaction through PowerSun.vip. "
        "The platform automatically estimates required energy, delegates it to the sender, "
        "then broadcasts the transaction. Requires API key authentication."
    )
    args_schema: Type[BaseModel] = BroadcastInput
    client: Any = None

    def _run(self, tx_data: str = "{}") -> str:
        parsed = json.loads(tx_data)
        result = self.client.post("/api/v2/agent/broadcast", {"txData": parsed})
        return json.dumps(result, indent=2)


# ─── Swap Tools (auth required) ────────────────────────────────────────


class GetSwapQuoteInput(BaseModel):
    from_token: str = Field(description="Token to sell — symbol (TRX, USDT, SUN, ...) or TRC-20 contract address")
    to_token: str = Field(description="Token to buy — symbol or TRC-20 contract address")
    amount_in: str = Field(description="Amount in base units (SUN for TRX, 6 decimals for USDT, etc.)")
    slippage_bps: int = Field(
        default=50,
        description="Slippage tolerance in basis points (default: 50 = 0.5%). Range: 1-5000",
    )


class GetSwapQuoteTool(BaseTool):
    """Get a swap quote for SunSwap DEX with energy cost estimate."""

    name: str = "powersun_get_swap_quote"
    description: str = (
        "Get a price quote and unsigned transaction for swapping tokens on SunSwap DEX "
        "via PowerSun.vip. Returns expected output, minimum output with slippage, "
        "energy cost, and an unsigned TX to sign. Supported: TRX, USDT, USDC, USDD, "
        "SUN, BTT, WIN, JST, or any TRC-20 address. Requires API key."
    )
    args_schema: Type[BaseModel] = GetSwapQuoteInput
    client: Any = None

    def _run(
        self,
        from_token: str = "TRX",
        to_token: str = "USDT",
        amount_in: str = "1000000",
        slippage_bps: int = 50,
    ) -> str:
        result = self.client.post("/api/v2/agent/swap", {
            "fromToken": from_token,
            "toToken": to_token,
            "amountIn": amount_in,
            "slippageBps": slippage_bps,
        })
        return json.dumps(result, indent=2)


class ExecuteSwapInput(BaseModel):
    tx_data: str = Field(
        description="JSON string of the pre-signed swap transaction (from get_swap_quote unsignedTx, after signing)"
    )


class ExecuteSwapTool(BaseTool):
    """Execute a pre-signed swap transaction with automatic energy delegation."""

    name: str = "powersun_execute_swap"
    description: str = (
        "Execute a pre-signed swap transaction through PowerSun.vip with automatic "
        "energy delegation. Get the unsigned TX from powersun_get_swap_quote, sign it "
        "with your private key, then submit here. Requires API key."
    )
    args_schema: Type[BaseModel] = ExecuteSwapInput
    client: Any = None

    def _run(self, tx_data: str = "{}") -> str:
        parsed = json.loads(tx_data)
        result = self.client.post("/api/v2/agent/broadcast", {"txData": parsed})
        return json.dumps(result, indent=2)


# ─── Registration Tools (public) ───────────────────────────────────────


class RegisterInput(BaseModel):
    address: str = Field(description="TRON wallet address (T-address, 34 chars) to register")


class RegisterTool(BaseTool):
    """Start PowerSun.vip registration by requesting a signature challenge."""

    name: str = "powersun_register"
    description: str = (
        "Start registration on PowerSun.vip by providing a TRON wallet address. "
        "Returns a challenge text that must be signed with the wallet's private key "
        "using tronWeb.trx.signMessageV2(). No authentication required."
    )
    args_schema: Type[BaseModel] = RegisterInput
    client: Any = None

    def _run(self, address: str = "") -> str:
        result = self.client.post("/api/v2/agent/register", {"address": address})
        return json.dumps(result, indent=2)


class VerifyRegistrationInput(BaseModel):
    challenge_id: str = Field(description="Challenge ID returned by the register tool")
    address: str = Field(description="TRON wallet address (same as used in register)")
    signature: str = Field(description="Hex signature from tronWeb.trx.signMessageV2(challenge, privateKey)")


class VerifyRegistrationTool(BaseTool):
    """Verify wallet signature and receive an API key."""

    name: str = "powersun_verify_registration"
    description: str = (
        "Complete PowerSun.vip registration by submitting the signed challenge. "
        "Returns an API key (ps_*) for authenticated endpoints. "
        "Returning users receive their existing API key. No authentication required."
    )
    args_schema: Type[BaseModel] = VerifyRegistrationInput
    client: Any = None

    def _run(self, challenge_id: str = "", address: str = "", signature: str = "") -> str:
        result = self.client.post("/api/v2/agent/verify", {
            "challengeId": challenge_id,
            "address": address,
            "signature": signature,
        })
        return json.dumps(result, indent=2)


# ─── Tool Registry ──────────────────────────────────────────────────────

REGISTRATION_TOOLS = [RegisterTool, VerifyRegistrationTool]
MARKET_TOOLS = [GetPricesTool, EstimateCostTool, GetAvailableResourcesTool, GetMarketOverviewTool]
BUYER_TOOLS = [BuyEnergyTool, GetBalanceTool, GetOrderStatusTool, BroadcastTransactionTool]
SWAP_TOOLS = [GetSwapQuoteTool, ExecuteSwapTool]
ALL_TOOLS = REGISTRATION_TOOLS + MARKET_TOOLS + BUYER_TOOLS + SWAP_TOOLS
