"""Tau 2 Benchmark - Retail Domain Tools.

All tools for the retail e-commerce domain. These tools execute REAL
business logic that modifies database state - they are NOT LLM-simulated.

Original benchmark: https://github.com/sierra-research/tau2-bench
Version: v0.2.0 (commit f8de30c, 2025-10-06)
Copyright (c) 2025 Sierra Research (MIT License)

Adapted from: src/tau2/domains/retail/tools.py
"""

import json
from typing import List

from maseval.benchmark.tau2.domains.base import ToolKitBase, ToolType, is_tool
from maseval.benchmark.tau2.domains.retail.db import RetailDB
from maseval.benchmark.tau2.domains.retail.models import (
    GiftCard,
    Order,
    OrderPayment,
    PaymentMethod,
    Product,
    User,
    UserAddress,
    Variant,
)


class RetailTools(ToolKitBase[RetailDB]):
    """All tools for the retail domain.

    These tools execute real business logic that modifies the database state.
    Each tool that modifies state is decorated with @is_tool(ToolType.WRITE).

    Adapted from: tau2-bench src/tau2/domains/retail/tools.py:RetailTools
    """

    def __init__(self, db: RetailDB) -> None:
        """Initialize retail tools.

        Args:
            db: RetailDB instance to operate on
        """
        super().__init__(db)

    # =========================================================================
    # Internal Helper Methods (not exposed as tools)
    # =========================================================================

    def _get_order(self, order_id: str) -> Order:
        """Get the order from the database.

        Args:
            order_id: The order id, such as '#W0000000'. Be careful there is a '#' symbol at the beginning of the order id.

        Returns:
            The order.

        Raises:
            ValueError: If the order is not found.
        """
        if self.db is None:
            raise ValueError("Database not initialized")
        if order_id not in self.db.orders:
            raise ValueError("Order not found")
        return self.db.orders[order_id]

    def _get_user(self, user_id: str) -> User:
        """Get the user from the database.

        Args:
            user_id: The user id, such as 'sara_doe_496'.

        Returns:
            The user.

        Raises:
            ValueError: If the user is not found.
        """
        if self.db is None:
            raise ValueError("Database not initialized")
        if user_id not in self.db.users:
            raise ValueError("User not found")
        return self.db.users[user_id]

    def _get_product(self, product_id: str) -> Product:
        """Get the product from the database.

        Args:
            product_id: The product id, such as '6086499569'. Be careful the product id is different from the item id.

        Returns:
            The product.

        Raises:
            ValueError: If the product is not found.
        """
        if self.db is None:
            raise ValueError("Database not initialized")
        if product_id not in self.db.products:
            raise ValueError("Product not found")
        return self.db.products[product_id]

    def _get_variant(self, product_id: str, variant_id: str) -> Variant:
        """Get the variant from the database.

        Args:
            product_id: The product id, such as '6086499569'.
            variant_id: The variant id, such as '1008292230'.

        Returns:
            The variant.

        Raises:
            ValueError: If the variant is not found.
        """
        product = self._get_product(product_id)
        if variant_id not in product.variants:
            raise ValueError("Variant not found")
        return product.variants[variant_id]

    def _get_payment_method(self, user_id: str, payment_method_id: str) -> PaymentMethod:
        """Get the payment method from the database.

        Args:
            user_id: The user id.
            payment_method_id: The payment method id.

        Returns:
            The payment method.

        Raises:
            ValueError: If the payment method is not found.
        """
        user = self._get_user(user_id)
        if payment_method_id not in user.payment_methods:
            raise ValueError("Payment method not found")
        return user.payment_methods[payment_method_id]

    def _is_pending_order(self, order: Order) -> bool:
        """Check if the order is pending.

        This is not a strict check, and not meant to be used for modify_items in pending orders.

        Args:
            order: The order.

        Returns:
            True if order status contains 'pending'.
        """
        return "pending" in order.status

    # =========================================================================
    # Generic Tools
    # =========================================================================

    @is_tool(ToolType.GENERIC)
    def calculate(self, expression: str) -> str:
        """Calculate the result of a mathematical expression.

        Args:
            expression: The mathematical expression to calculate, such as '2 + 2'.
                       The expression can contain numbers, operators (+, -, *, /),
                       parentheses, and spaces.

        Returns:
            The result of the mathematical expression.

        Raises:
            ValueError: If the expression is invalid.
        """
        if not all(char in "0123456789+-*/(). " for char in expression):
            raise ValueError("Invalid characters in expression")
        return str(round(float(eval(expression, {"__builtins__": None}, {})), 2))

    @is_tool(ToolType.GENERIC)
    def transfer_to_human_agents(self, summary: str) -> str:
        """Transfer the user to a human agent, with a summary of the user's issue.

        Only transfer if:
        - the user explicitly asks for a human agent
        - given the policy and the available tools, you cannot solve the user's issue.

        Args:
            summary: A summary of the user's issue.

        Returns:
            A message indicating the user has been transferred to a human agent.
        """
        return "Transfer successful"

    # =========================================================================
    # Read Tools
    # =========================================================================

    @is_tool(ToolType.READ)
    def find_user_id_by_email(self, email: str) -> str:
        """Find user id by email.

        If the user is not found, the function will return an error message.

        Args:
            email: The email of the user, such as 'something@example.com'.

        Returns:
            The user id if found.

        Raises:
            ValueError: If the user is not found.
        """
        if self.db is None:
            raise ValueError("Database not initialized")
        for user_id, user in self.db.users.items():
            if user.email.lower() == email.lower():
                return user_id
        raise ValueError("User not found")

    @is_tool(ToolType.READ)
    def find_user_id_by_name_zip(self, first_name: str, last_name: str, zip: str) -> str:
        """Find user id by first name, last name, and zip code.

        If the user is not found, the function will return an error message.
        By default, find user id by email, and only call this function
        if the user is not found by email or cannot remember email.

        Args:
            first_name: The first name of the customer, such as 'John'.
            last_name: The last name of the customer, such as 'Doe'.
            zip: The zip code of the customer, such as '12345'.

        Returns:
            The user id if found.

        Raises:
            ValueError: If the user is not found.
        """
        if self.db is None:
            raise ValueError("Database not initialized")
        for user_id, user in self.db.users.items():
            if (
                user.name.first_name.lower() == first_name.lower()
                and user.name.last_name.lower() == last_name.lower()
                and user.address.zip == zip
            ):
                return user_id
        raise ValueError("User not found")

    @is_tool(ToolType.READ)
    def get_order_details(self, order_id: str) -> Order:
        """Get the status and details of an order.

        Args:
            order_id: The order id, such as '#W0000000'. Be careful there is a '#' symbol
                     at the beginning of the order id.

        Returns:
            The order details.

        Raises:
            ValueError: If the order is not found.
        """
        return self._get_order(order_id)

    @is_tool(ToolType.READ)
    def get_product_details(self, product_id: str) -> Product:
        """Get the inventory details of a product.

        Args:
            product_id: The product id, such as '6086499569'. Be careful the product id
                       is different from the item id.

        Returns:
            The product details.

        Raises:
            ValueError: If the product is not found.
        """
        return self._get_product(product_id)

    @is_tool(ToolType.READ)
    def get_user_details(self, user_id: str) -> User:
        """Get the details of a user, including their orders.

        Args:
            user_id: The user id, such as 'sara_doe_496'.

        Returns:
            The user details.

        Raises:
            ValueError: If the user is not found.
        """
        return self._get_user(user_id)

    @is_tool(ToolType.READ)
    def list_all_product_types(self) -> str:
        """List the name and product id of all product types.

        Each product type has a variety of different items with unique item ids and options.
        There are only 50 product types in the store.

        Returns:
            A JSON string mapping product names to their product IDs, sorted alphabetically by name.
        """
        if self.db is None:
            raise ValueError("Database not initialized")
        product_dict = {product.name: product.product_id for product in self.db.products.values()}
        return json.dumps(product_dict, sort_keys=True)

    # =========================================================================
    # Write Tools
    # =========================================================================

    @is_tool(ToolType.WRITE)
    def cancel_pending_order(self, order_id: str, reason: str) -> Order:
        """Cancel a pending order.

        If the order is already processed or delivered, it cannot be cancelled.
        The agent needs to explain the cancellation detail and ask for explicit
        user confirmation (yes/no) to proceed.

        If the user confirms, the order status will be changed to 'cancelled'
        and the payment will be refunded. The refund will be added to the user's
        gift card balance immediately if the payment was made using a gift card,
        otherwise the refund would take 5-7 business days to process.

        Args:
            order_id: The order id, such as '#W0000000'. Be careful there is a '#' symbol
                     at the beginning of the order id.
            reason: The reason for cancellation, which should be either 'no longer needed'
                   or 'ordered by mistake'.

        Returns:
            The order details after the cancellation.

        Raises:
            ValueError: If order is not pending or reason is invalid.
        """
        # Check order exists and is pending
        order = self._get_order(order_id)
        if order.status != "pending":
            raise ValueError("Non-pending order cannot be cancelled")

        # Check reason
        if reason not in {"no longer needed", "ordered by mistake"}:
            raise ValueError("Invalid reason")

        # Handle refund
        refunds = []
        for payment in order.payment_history:
            payment_id = payment.payment_method_id
            refund = OrderPayment(
                transaction_type="refund",
                amount=payment.amount,
                payment_method_id=payment_id,
            )
            refunds.append(refund)
            user = self._get_user(order.user_id)
            payment_method = self._get_payment_method(user.user_id, payment_id)
            if isinstance(payment_method, GiftCard):
                # Refund to gift card immediately
                payment_method.balance += payment.amount
                payment_method.balance = round(payment_method.balance, 2)

        # Update order status
        order.status = "cancelled"
        order.cancel_reason = reason  # type: ignore[assignment]
        order.payment_history.extend(refunds)

        return order

    @is_tool(ToolType.WRITE)
    def exchange_delivered_order_items(
        self,
        order_id: str,
        item_ids: List[str],
        new_item_ids: List[str],
        payment_method_id: str,
    ) -> Order:
        """Exchange items in a delivered order to new items of the same product type.

        For a delivered order, return or exchange can be only done once by the agent.
        The agent needs to explain the exchange detail and ask for explicit user
        confirmation (yes/no) to proceed.

        Args:
            order_id: The order id, such as '#W0000000'.
            item_ids: The item ids to be exchanged, each such as '1008292230'.
                     There could be duplicate items in the list.
            new_item_ids: The item ids to be exchanged for, each such as '1008292230'.
                         There could be duplicate items in the list. Each new item id
                         should match the item id in the same position and be of the same product.
            payment_method_id: The payment method id to pay or receive refund for the
                              item price difference.

        Returns:
            The order details after the exchange.

        Raises:
            ValueError: If order is not delivered, items don't exist, or items don't match.
        """
        # Check order exists and is delivered
        order = self._get_order(order_id)
        if order.status != "delivered":
            raise ValueError("Non-delivered order cannot be exchanged")

        # Check the items to be exchanged exist
        all_item_ids = [item.item_id for item in order.items]
        for item_id in item_ids:
            if item_ids.count(item_id) > all_item_ids.count(item_id):
                raise ValueError(f"Number of {item_id} not found.")

        # Check new items exist and match old items and are available
        if len(item_ids) != len(new_item_ids):
            raise ValueError("The number of items to be exchanged should match.")

        diff_price = 0.0
        for item_id, new_item_id in zip(item_ids, new_item_ids):
            item = next((item for item in order.items if item.item_id == item_id), None)
            if item is None:
                raise ValueError(f"Item {item_id} not found")
            product_id = item.product_id
            variant = self._get_variant(product_id, new_item_id)
            if not variant.available:
                raise ValueError(f"New item {new_item_id} not found or available")

            old_price = item.price
            new_price = variant.price
            diff_price += new_price - old_price

        diff_price = round(diff_price, 2)

        # Check payment method exists and can cover the price difference if gift card
        payment_method = self._get_payment_method(order.user_id, payment_method_id)

        if isinstance(payment_method, GiftCard) and payment_method.balance < diff_price:
            raise ValueError("Insufficient gift card balance to pay for the price difference")

        # Modify the order
        order.status = "exchange requested"
        order.exchange_items = sorted(item_ids)
        order.exchange_new_items = sorted(new_item_ids)
        order.exchange_payment_method_id = payment_method_id
        order.exchange_price_difference = diff_price

        return order

    @is_tool(ToolType.WRITE)
    def modify_pending_order_address(
        self,
        order_id: str,
        address1: str,
        address2: str,
        city: str,
        state: str,
        country: str,
        zip: str,
    ) -> Order:
        """Modify the shipping address of a pending order.

        The agent needs to explain the modification detail and ask for explicit
        user confirmation (yes/no) to proceed.

        Args:
            order_id: The order id, such as '#W0000000'.
            address1: The first line of the address, such as '123 Main St'.
            address2: The second line of the address, such as 'Apt 1' or ''.
            city: The city, such as 'San Francisco'.
            state: The state, such as 'CA'.
            country: The country, such as 'USA'.
            zip: The zip code, such as '12345'.

        Returns:
            The order details after the modification.

        Raises:
            ValueError: If the order is not pending.
        """
        # Check if the order exists and is pending
        order = self._get_order(order_id)
        if not self._is_pending_order(order):
            raise ValueError("Non-pending order cannot be modified")

        # Modify the address
        order.address = UserAddress(
            address1=address1,
            address2=address2,
            city=city,
            state=state,
            country=country,
            zip=zip,
        )
        return order

    @is_tool(ToolType.WRITE)
    def modify_pending_order_items(
        self,
        order_id: str,
        item_ids: List[str],
        new_item_ids: List[str],
        payment_method_id: str,
    ) -> Order:
        """Modify items in a pending order to new items of the same product type.

        For a pending order, this function can only be called once.
        The agent needs to explain the exchange detail and ask for explicit
        user confirmation (yes/no) to proceed.

        Args:
            order_id: The order id, such as '#W0000000'.
            item_ids: The item ids to be modified, each such as '1008292230'.
                     There could be duplicate items in the list.
            new_item_ids: The item ids to be modified for. Each new item id should
                         match the item id in the same position and be of the same product.
            payment_method_id: The payment method id to pay or receive refund for
                              the item price difference.

        Returns:
            The order details after the modification.

        Raises:
            ValueError: If order is not pending, items don't exist, or items don't match.
        """
        # Check if the order exists and is pending
        order = self._get_order(order_id)
        if order.status != "pending":
            raise ValueError("Non-pending order cannot be modified")

        # Check if the items to be modified exist
        all_item_ids = [item.item_id for item in order.items]
        for item_id in item_ids:
            if item_ids.count(item_id) > all_item_ids.count(item_id):
                raise ValueError(f"{item_id} not found")

        # Check new items exist, match old items, and are available
        if len(item_ids) != len(new_item_ids):
            raise ValueError("The number of items to be exchanged should match")

        diff_price = 0
        for item_id, new_item_id in zip(item_ids, new_item_ids):
            if item_id == new_item_id:
                raise ValueError("The new item id should be different from the old item id")
            item = next((item for item in order.items if item.item_id == item_id), None)
            if item is None:
                raise ValueError(f"Item {item_id} not found")
            product_id = item.product_id
            variant = self._get_variant(product_id, new_item_id)
            if not variant.available:
                raise ValueError(f"New item {new_item_id} not found or available")

            old_price = item.price
            new_price = variant.price
            diff_price += new_price - old_price

        # Check if the payment method exists
        payment_method = self._get_payment_method(order.user_id, payment_method_id)

        # If the new item is more expensive, check if the gift card has enough balance
        if isinstance(payment_method, GiftCard) and payment_method.balance < diff_price:
            raise ValueError("Insufficient gift card balance to pay for the new item")

        # Handle the payment or refund
        order.payment_history.append(
            OrderPayment(
                transaction_type="payment" if diff_price > 0 else "refund",
                amount=abs(diff_price),
                payment_method_id=payment_method_id,
            )
        )
        if isinstance(payment_method, GiftCard):
            payment_method.balance -= diff_price
            payment_method.balance = round(payment_method.balance, 2)

        # Modify the order
        for item_id, new_item_id in zip(item_ids, new_item_ids):
            item = next((item for item in order.items if item.item_id == item_id), None)
            if item is None:
                raise ValueError(f"Item {item_id} not found")
            item.item_id = new_item_id
            item.price = variant.price
            item.options = variant.options
        order.status = "pending (item modified)"

        return order

    @is_tool(ToolType.WRITE)
    def modify_pending_order_payment(
        self,
        order_id: str,
        payment_method_id: str,
    ) -> Order:
        """Modify the payment method of a pending order.

        The agent needs to explain the modification detail and ask for explicit
        user confirmation (yes/no) to proceed.

        Args:
            order_id: The order id, such as '#W0000000'.
            payment_method_id: The new payment method id.

        Returns:
            The order details after the modification.

        Raises:
            ValueError: If order is not pending, payment method doesn't exist,
                       or there are multiple payments.
        """
        order = self._get_order(order_id)

        # Check if the order exists and is pending
        if not self._is_pending_order(order):
            raise ValueError("Non-pending order cannot be modified")

        # Check if the payment method exists
        payment_method = self._get_payment_method(order.user_id, payment_method_id)

        # Check that the payment history should only have one payment
        if len(order.payment_history) != 1 or order.payment_history[0].transaction_type != "payment":
            raise ValueError("There should be exactly one payment for a pending order")

        # Check that the payment method is different
        if order.payment_history[0].payment_method_id == payment_method_id:
            raise ValueError("The new payment method should be different from the current one")

        amount = order.payment_history[0].amount

        # Check if the new payment method has enough balance if it is a gift card
        if isinstance(payment_method, GiftCard) and payment_method.balance < amount:
            raise ValueError("Insufficient gift card balance to pay for the order")

        # Modify the payment method
        order.payment_history.extend(
            [
                OrderPayment(
                    transaction_type="payment",
                    amount=amount,
                    payment_method_id=payment_method_id,
                ),
                OrderPayment(
                    transaction_type="refund",
                    amount=amount,
                    payment_method_id=order.payment_history[0].payment_method_id,
                ),
            ]
        )

        # If payment is made by gift card, update the balance
        if isinstance(payment_method, GiftCard):
            payment_method.balance -= amount
            payment_method.balance = round(payment_method.balance, 2)

        # If refund is made to a gift card, update the balance
        old_payment_method = self._get_payment_method(order.user_id, order.payment_history[0].payment_method_id)
        if isinstance(old_payment_method, GiftCard):
            old_payment_method.balance += amount
            old_payment_method.balance = round(old_payment_method.balance, 2)

        return order

    @is_tool(ToolType.WRITE)
    def modify_user_address(
        self,
        user_id: str,
        address1: str,
        address2: str,
        city: str,
        state: str,
        country: str,
        zip: str,
    ) -> User:
        """Modify the default address of a user.

        The agent needs to explain the modification detail and ask for explicit
        user confirmation (yes/no) to proceed.

        Args:
            user_id: The user id, such as 'sara_doe_496'.
            address1: The first line of the address, such as '123 Main St'.
            address2: The second line of the address, such as 'Apt 1' or ''.
            city: The city, such as 'San Francisco'.
            state: The state, such as 'CA'.
            country: The country, such as 'USA'.
            zip: The zip code, such as '12345'.

        Returns:
            The user details after the modification.

        Raises:
            ValueError: If the user is not found.
        """
        user = self._get_user(user_id)
        user.address = UserAddress(
            address1=address1,
            address2=address2,
            city=city,
            state=state,
            country=country,
            zip=zip,
        )
        return user

    @is_tool(ToolType.WRITE)
    def return_delivered_order_items(
        self,
        order_id: str,
        item_ids: List[str],
        payment_method_id: str,
    ) -> Order:
        """Return some items of a delivered order.

        The order status will be changed to 'return requested'.
        The agent needs to explain the return detail and ask for explicit
        user confirmation (yes/no) to proceed.
        The user will receive follow-up email for how and where to return the item.

        Args:
            order_id: The order id, such as '#W0000000'.
            item_ids: The item ids to be returned, each such as '1008292230'.
                     There could be duplicate items in the list.
            payment_method_id: The payment method id to receive the refund.
                              Should be the original payment method or a gift card.

        Returns:
            The order details after requesting the return.

        Raises:
            ValueError: If order is not delivered, payment method is invalid,
                       or items don't exist.
        """
        order = self._get_order(order_id)
        if order.status != "delivered":
            raise ValueError("Non-delivered order cannot be returned")

        # Check if the payment method exists and is either the original payment method or a gift card
        user = self._get_user(order.user_id)
        payment_method = self._get_payment_method(user.user_id, payment_method_id)

        if not isinstance(payment_method, GiftCard) and payment_method_id != order.payment_history[0].payment_method_id:
            raise ValueError("Payment method should be the original payment method")

        # Check if the items to be returned exist
        all_item_ids = [item.item_id for item in order.items]
        for item_id in item_ids:
            if item_ids.count(item_id) > all_item_ids.count(item_id):
                raise ValueError("Some item not found")

        # Update the order status
        order.status = "return requested"
        order.return_items = sorted(item_ids)
        order.return_payment_method_id = payment_method_id

        return order
