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

from __future__ import annotations

import json
import os
from typing import Any, Type

import httpx
from crewai.tools import BaseTool
from pydantic import BaseModel, Field

BASE_URL = "https://powersun.vip"
TIMEOUT = 30.0


def _get_client(api_key: str | None = None) -> httpx.Client:
    key = api_key or os.environ.get("POWERSUN_API_KEY")
    headers: dict[str, str] = {"User-Agent": "crewai-powersun/1.2.0"}
    if key:
        headers["X-API-Key"] = key
    return httpx.Client(base_url=BASE_URL, headers=headers, timeout=TIMEOUT)


# ─── Registration Tools ──────────────────────────────────────────────────


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


class PowerSunRegisterTool(BaseTool):
    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
    api_key: str | None = None

    def _run(self, address: str = "") -> str:
        with _get_client(self.api_key) as c:
            r = c.post("/api/v2/agent/register", json={"address": address})
            r.raise_for_status()
            return json.dumps(r.json(), 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 PowerSunVerifyRegistrationTool(BaseTool):
    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
    api_key: str | None = None

    def _run(self, challenge_id: str = "", address: str = "", signature: str = "") -> str:
        with _get_client(self.api_key) as c:
            r = c.post("/api/v2/agent/verify", json={
                "challengeId": challenge_id,
                "address": address,
                "signature": signature,
            })
            r.raise_for_status()
            return json.dumps(r.json(), indent=2)


# ─── Market Tools ────────────────────────────────────────────────────────


class PowerSunMarketOverviewTool(BaseTool):
    name: str = "PowerSun Market Overview"
    description: str = (
        "Get a comprehensive market snapshot from PowerSun.vip TRON Energy marketplace: "
        "current prices, available resources, 24h order stats. No authentication required."
    )
    api_key: str | None = None

    def _run(self, **kwargs: Any) -> str:
        with _get_client(self.api_key) as c:
            r = c.get("/api/v2/agent/market-overview")
            r.raise_for_status()
            return json.dumps(r.json(), indent=2)


class EstimateCostInput(BaseModel):
    tx_count: int = Field(description="Number of TRON transactions", ge=1)
    tx_type: str = Field(default="trc20_transfer", description="Transaction type")
    duration_minutes: int = Field(default=60, description="Duration in minutes: 5, 10, 15, 30, 60, 120, 1440, 10080")


class PowerSunEstimateCostTool(BaseTool):
    name: str = "PowerSun Estimate Cost"
    description: str = (
        "Estimate the cost of renting TRON Energy from PowerSun.vip. "
        "Returns cost in TRX, savings vs burning, and minimum order info."
    )
    args_schema: Type[BaseModel] = EstimateCostInput
    api_key: str | None = None

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


# ─── Buyer Tools ─────────────────────────────────────────────────────────


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)
    duration_minutes: int = Field(default=60, description="Duration in minutes: 5, 10, 15, 30, 60, 120, 1440, 10080")


class PowerSunBuyEnergyTool(BaseTool):
    name: str = "PowerSun Buy Energy"
    description: str = (
        "Purchase TRON Energy delegation from PowerSun.vip marketplace. "
        "Energy is delegated to the target address within seconds. "
        "Save 20-50% vs burning TRX. Only 10% commission. "
        "Requires POWERSUN_API_KEY environment variable or api_key parameter."
    )
    args_schema: Type[BaseModel] = BuyEnergyInput
    api_key: str | None = None

    def _run(self, target_address: str, tx_count: int = 1, duration_minutes: int = 60) -> str:
        with _get_client(self.api_key) as c:
            r = c.post("/api/v2/agent/buy-energy", json={
                "targetAddress": target_address,
                "txCount": tx_count,
                "durationMinutes": duration_minutes,
            })
            return json.dumps(r.json(), indent=2)


class PowerSunGetBalanceTool(BaseTool):
    name: str = "PowerSun Get Balance"
    description: str = (
        "Get your PowerSun.vip account balance in TRX and deposit address. "
        "Requires POWERSUN_API_KEY environment variable or api_key parameter."
    )
    api_key: str | None = None

    def _run(self, **kwargs: Any) -> str:
        with _get_client(self.api_key) as c:
            r = c.get("/api/v2/agent/balance")
            r.raise_for_status()
            return json.dumps(r.json(), indent=2)


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


class PowerSunGetOrderStatusTool(BaseTool):
    name: str = "PowerSun Get Order Status"
    description: str = (
        "Get detailed status of an energy purchase order from PowerSun.vip "
        "including delegation progress and transaction hashes. "
        "Requires POWERSUN_API_KEY environment variable or api_key parameter."
    )
    args_schema: Type[BaseModel] = OrderStatusInput
    api_key: str | None = None

    def _run(self, order_id: str = "") -> str:
        with _get_client(self.api_key) as c:
            r = c.get(f"/api/v2/agent/order/{order_id}")
            r.raise_for_status()
            return json.dumps(r.json(), indent=2)


# ─── Swap Tools ─────────────────────────────────────────────────────────


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 (50 = 0.5%)")


class PowerSunGetSwapQuoteTool(BaseTool):
    name: str = "PowerSun Get Swap Quote"
    description: str = (
        "Get a swap quote for SunSwap DEX via PowerSun.vip. Returns expected output, "
        "minimum output with slippage, energy cost, and unsigned TX to sign. "
        "Supports: TRX, USDT, USDC, USDD, SUN, BTT, WIN, JST, or any TRC-20 address. "
        "Requires POWERSUN_API_KEY environment variable or api_key parameter."
    )
    args_schema: Type[BaseModel] = GetSwapQuoteInput
    api_key: str | None = None

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


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


class PowerSunExecuteSwapTool(BaseTool):
    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, "
        "then submit here. Requires POWERSUN_API_KEY environment variable or api_key parameter."
    )
    args_schema: Type[BaseModel] = ExecuteSwapInput
    api_key: str | None = None

    def _run(self, tx_data: str = "{}") -> str:
        parsed = json.loads(tx_data)
        with _get_client(self.api_key) as c:
            r = c.post("/api/v2/agent/broadcast", json={"txData": parsed})
            r.raise_for_status()
            return json.dumps(r.json(), indent=2)
