import os
from typing import Type, Any, Optional
from pydantic import BaseModel, Field
from langchain_core.tools import BaseTool
from aegis.core.models import PaymentIntent
from aegis.client import AegisClient


class AegisPaymentInput(BaseModel):
    requested_amount: float = Field(..., description="The amount of money to request.")
    target_vendor: str = Field(..., description="The vendor to pay.")
    reasoning: str = Field(..., description="Reasoning for this payment.")


class AegisPaymentTool(BaseTool):
    name: str = "aegis_payment_tool"
    description: str = (
        "Use this tool to request a one-time virtual card for an automated purchase. "
        "ONLY call this tool when you are on the FINAL checkout page and can visually "
        "confirm that credit card input fields are visible. "
        "If auto-injection is enabled, the card will be securely filled into the browser "
        "automatically — you only need to click the submit/pay button afterward. "
        "DO NOT retry with different reasoning if rejected. "
        "Provide the amount (float), target vendor (str), and your full reasoning (str)."
    )
    args_schema: Type[BaseModel] = AegisPaymentInput

    client: Any = Field(description="The AegisClient instance")
    agent_id: str = Field(..., description="The ID of the Agent making the request")
    injector: Optional[Any] = Field(default=None, description="Optional AegisBrowserInjector instance")
    cdp_url: str = Field(default="http://localhost:9222", description="CDP endpoint for browser injection")

    def __init__(
        self,
        client: AegisClient,
        agent_id: str,
        injector=None,
        cdp_url: str = "http://localhost:9222",
        **kwargs,
    ):
        super().__init__(
            client=client,
            agent_id=agent_id,
            injector=injector,
            cdp_url=cdp_url,
            **kwargs,
        )

    def _run(
        self,
        requested_amount: float,
        target_vendor: str,
        reasoning: str,
        run_manager=None,
    ) -> str:
        return "Please use the async method ainvoke() for AegisPaymentTool."

    async def _arun(
        self,
        requested_amount: float,
        target_vendor: str,
        reasoning: str,
        run_manager=None,
    ) -> str:
        intent = PaymentIntent(
            agent_id=self.agent_id,
            requested_amount=requested_amount,
            target_vendor=target_vendor,
            reasoning=reasoning,
        )

        seal = await self.client.process_payment(intent)

        if seal.status.lower() == "rejected":
            return f"Payment rejected by guardrails. Reason: {seal.rejection_reason}"

        masked_card = f"****-****-****-{seal.card_number[-4:]}"

        # -------------------------------------------------------------------
        # Auto-injection path: if an injector is provided, fill the browser
        # -------------------------------------------------------------------
        if self.injector is not None:
            injection_ok = await self.injector.inject_payment_info(
                seal_id=seal.seal_id,
                cdp_url=self.cdp_url,
            )

            if not injection_ok:
                # Cancel the budget reservation — treat as if never issued
                self.client.state_tracker.mark_used(seal.seal_id)
                return (
                    "Payment rejected. Error: Aegis could not find credit card input "
                    "fields on your active browser tab. Please ensure you have navigated "
                    "to the FINAL checkout form and the card fields are visible, then retry."
                )

            return (
                f"Payment approved and securely auto-injected into the browser form. "
                f"Please proceed to click the submit/pay button. "
                f"Masked card: {masked_card}"
            )

        # -------------------------------------------------------------------
        # Standard path: return masked card details only
        # -------------------------------------------------------------------
        return (
            f"Payment approved. Card Issued: {masked_card}, "
            f"Expiry: {seal.expiration_date}, Authorized Amount: {seal.authorized_amount}"
        )
