# cartpilot_cart/agent.py
from google.adk.agents import Agent
from google.adk.tools import ToolContext
from .catalog import find_products, get_product

CART_KEY = "cart"  # list[dict] of {sku, qty} stored in session state


def add_to_cart(sku: str, quantity: int, tool_context: ToolContext) -> dict:
    """Adds a product to the shopper's cart.

    Args:
        sku: The product SKU to add.
        quantity: How many units to add. Must be a positive integer.
    """
    product = get_product(sku)
    if product is None:
        return {"status": "not_found", "sku": sku}
    if quantity <= 0:
        return {"status": "invalid_quantity", "message": "Quantity must be positive."}
    if product["stock"] < quantity:
        return {"status": "insufficient_stock", "available": product["stock"]}

    cart = tool_context.state.get(CART_KEY, [])
    for line in cart:
        if line["sku"] == sku:
            line["quantity"] += quantity
            break
    else:
        cart.append({"sku": sku, "quantity": quantity})
    tool_context.state[CART_KEY] = cart
    return {"status": "success", "cart": cart}


def view_cart(tool_context: ToolContext) -> dict:
    """Returns the current contents of the shopper's cart with line totals."""
    cart = tool_context.state.get(CART_KEY, [])
    lines, subtotal = [], 0.0
    for entry in cart:
        product = get_product(entry["sku"])
        if not product:
            continue
        line_total = product["price"] * entry["quantity"]
        subtotal += line_total
        lines.append({"sku": entry["sku"], "name": product["name"],
                       "quantity": entry["quantity"], "line_total": round(line_total, 2)})
    return {"status": "success", "lines": lines, "subtotal": round(subtotal, 2)}


def remove_from_cart(sku: str, tool_context: ToolContext) -> dict:
    """Removes a product entirely from the shopper's cart.

    Args:
        sku: The product SKU to remove.
    """
    cart = tool_context.state.get(CART_KEY, [])
    new_cart = [line for line in cart if line["sku"] != sku]
    if len(new_cart) == len(cart):
        return {"status": "not_in_cart", "sku": sku}
    tool_context.state[CART_KEY] = new_cart
    return {"status": "success", "cart": new_cart}


root_agent = Agent(
    model="gemini-3.5-flash",
    name="cartpilot_agent",
    description="Helps shoppers search products and manage their cart.",
    instruction="""
You are the CartPilot shopping assistant.
- Use `search_products` to help shoppers find items.
- Use `add_to_cart`, `view_cart`, and `remove_from_cart` to manage the cart.
- Always confirm what changed after a cart action (e.g. "Added 2x Trail
  Runner Jacket — your subtotal is now $179.98.").
- Never guess at cart contents — call `view_cart` if you're not sure.
- Keep responses short (2-4 sentences).
""",
    tools=[search_products, add_to_cart, view_cart, remove_from_cart],
)

# main.py
import asyncio
from google.adk.runners import Runner
from google.adk.sessions import InMemorySessionService
from google.genai import types
from cartpilot_cart.agent import root_agent

APP_NAME = "cartpilot"
USER_ID = "demo-user"

async def main():
    session_service = InMemorySessionService()
    session = await session_service.create_session(app_name=APP_NAME, user_id=USER_ID)
    runner = Runner(agent=root_agent, app_name=APP_NAME, session_service=session_service)

    for turn in ["Do you have any water bottles?", "Add one to my cart.", "What's in my cart so far?"]:
        print(f"\n> {turn}")
        message = types.Content(role="user", parts=[types.Part(text=turn)])
        async for event in runner.run_async(user_id=USER_ID, session_id=session.id, new_message=message):
            if event.is_final_response() and event.content and event.content.parts:
                print(event.content.parts[0].text)

    final_session = await session_service.get_session(app_name=APP_NAME, user_id=USER_ID, session_id=session.id)
    print("\n--- raw session state ---")
    print(final_session.state)  # {'cart': [{'sku': 'sku-002', 'quantity': 1}]}

if __name__ == "__main__":
    asyncio.run(main())