"""LangChain tool wrappers for Wallgent financial operations.

Usage::

    from wallgent import Wallgent
    from wallgent.langchain import get_wallgent_tools

    client = Wallgent("wg_test_...")
    tools = get_wallgent_tools(client)

    # Pass tools to your LangChain agent
    agent = initialize_agent(tools=tools, llm=llm, ...)
"""

from __future__ import annotations

import json
from typing import Any, Dict, List, Optional, Type

from pydantic import BaseModel, Field

try:
    from langchain_core.tools import BaseTool
except ImportError:
    raise ImportError(
        "langchain-core is required for LangChain integration. "
        "Install it with: pip install wallgent[langchain] or pip install langchain-core"
    )

from wallgent.client import Wallgent


# ── Input schemas ────────────────────────────────────────────


class CheckBalanceInput(BaseModel):
    wallet_id: str = Field(description="The wallet ID to check balance for")


class SendPaymentInput(BaseModel):
    source_wallet_id: str = Field(description="Source wallet ID")
    destination_wallet_id: str = Field(description="Destination wallet ID")
    amount: float = Field(description="Amount to send")
    idempotency_key: str = Field(description="Unique key for idempotent requests")
    description: Optional[str] = Field(default=None, description="Payment description")


class CreateWalletInput(BaseModel):
    name: str = Field(description="Wallet name")
    agent_name: Optional[str] = Field(default=None, description="Agent identifier")


class ListWalletsInput(BaseModel):
    pass


class FundWalletInput(BaseModel):
    wallet_id: str = Field(description="Wallet ID to fund")
    amount: float = Field(description="Amount to fund")


class CreateInvoiceInput(BaseModel):
    wallet_id: str = Field(alias="walletId", description="Wallet to receive payment")
    customer_id: str = Field(alias="customerId", description="Customer to charge")
    line_items: str = Field(
        alias="lineItems",
        description='JSON array of line items, e.g. [{"description":"Service","unitPrice":"100.00"}]',
    )
    memo: Optional[str] = Field(default=None, description="Invoice memo")


class PayAgentInput(BaseModel):
    to_address: str = Field(description="Recipient Wallgent address (wg_...)")
    amount: str = Field(description='Amount to send (e.g. "50.00")')
    description: Optional[str] = Field(default=None, description="Payment description")


class SearchDirectoryInput(BaseModel):
    search: Optional[str] = Field(default=None, description="Search query for agent name or address")


# ── Tool implementations ─────────────────────────────────────


class WallgentCheckBalanceTool(BaseTool):
    name: str = "wallgent_check_balance"
    description: str = "Check the balance of a Wallgent wallet."
    args_schema: Type[BaseModel] = CheckBalanceInput
    client: Any = None

    def _run(self, wallet_id: str) -> str:
        result = self.client.wallets.get_balance(wallet_id)
        return json.dumps(result)


class WallgentSendPaymentTool(BaseTool):
    name: str = "wallgent_send_payment"
    description: str = "Send a payment from one wallet to another."
    args_schema: Type[BaseModel] = SendPaymentInput
    client: Any = None

    def _run(
        self,
        source_wallet_id: str,
        destination_wallet_id: str,
        amount: float,
        idempotency_key: str,
        description: Optional[str] = None,
    ) -> str:
        result = self.client.payments.send(
            source_wallet_id=source_wallet_id,
            destination_wallet_id=destination_wallet_id,
            amount=amount,
            idempotency_key=idempotency_key,
            description=description,
        )
        return json.dumps(result)


class WallgentCreateWalletTool(BaseTool):
    name: str = "wallgent_create_wallet"
    description: str = "Create a new wallet for an AI agent."
    args_schema: Type[BaseModel] = CreateWalletInput
    client: Any = None

    def _run(self, name: str, agent_name: Optional[str] = None) -> str:
        result = self.client.wallets.create(name=name, agent_name=agent_name)
        return json.dumps(result)


class WallgentListWalletsTool(BaseTool):
    name: str = "wallgent_list_wallets"
    description: str = "List all wallets."
    args_schema: Type[BaseModel] = ListWalletsInput
    client: Any = None

    def _run(self) -> str:
        result = self.client.wallets.list()
        return json.dumps(result)


class WallgentFundWalletTool(BaseTool):
    name: str = "wallgent_fund_wallet"
    description: str = "Add funds to a wallet (sandbox only)."
    args_schema: Type[BaseModel] = FundWalletInput
    client: Any = None

    def _run(self, wallet_id: str, amount: float) -> str:
        result = self.client.wallets.fund(wallet_id, amount=amount)
        return json.dumps(result)


class WallgentCreateInvoiceTool(BaseTool):
    name: str = "wallgent_create_invoice"
    description: str = "Create an invoice for a customer with line items."
    args_schema: Type[BaseModel] = CreateInvoiceInput
    client: Any = None

    def _run(
        self,
        wallet_id: str = "",
        customer_id: str = "",
        line_items: str = "[]",
        memo: Optional[str] = None,
    ) -> str:
        items = json.loads(line_items) if isinstance(line_items, str) else line_items
        params: Dict[str, Any] = {
            "walletId": wallet_id,
            "customerId": customer_id,
            "lineItems": items,
        }
        if memo:
            params["memo"] = memo
        result = self.client.invoices.create(**params)
        return json.dumps(result)


class WallgentPayAgentTool(BaseTool):
    name: str = "wallgent_pay_agent"
    description: str = (
        "Pay another agent by their Wallgent address (wg_...). "
        "Instant, zero fees via the Wallgent payment network."
    )
    args_schema: Type[BaseModel] = PayAgentInput
    client: Any = None

    def _run(
        self,
        to_address: str,
        amount: str,
        description: Optional[str] = None,
    ) -> str:
        result = self.client.network.pay_agent(
            to_address=to_address,
            amount=amount,
            description=description,
        )
        return json.dumps(result)


class WallgentSearchDirectoryTool(BaseTool):
    name: str = "wallgent_search_directory"
    description: str = "Search the Wallgent agent directory to find other agents you can pay."
    args_schema: Type[BaseModel] = SearchDirectoryInput
    client: Any = None

    def _run(self, search: Optional[str] = None) -> str:
        result = self.client.network.search_directory(search)
        return json.dumps(result)


# ── Factory ──────────────────────────────────────────────────


def get_wallgent_tools(client: Wallgent) -> List[BaseTool]:
    """Create LangChain tools from an authenticated Wallgent client.

    Args:
        client: An authenticated ``Wallgent`` instance.

    Returns:
        A list of LangChain ``BaseTool`` instances ready for agent use.
    """
    return [
        WallgentCheckBalanceTool(client=client),
        WallgentSendPaymentTool(client=client),
        WallgentCreateWalletTool(client=client),
        WallgentListWalletsTool(client=client),
        WallgentFundWalletTool(client=client),
        WallgentCreateInvoiceTool(client=client),
        WallgentPayAgentTool(client=client),
        WallgentSearchDirectoryTool(client=client),
    ]
