backend/agents/ihg_agent.py

"""
IHG.com browser agent — finds reward-night (points) availability.
Uses google-genai (Gemini Flash free tier) with Playwright browser tools.
"""

from __future__ import annotations

import asyncio
import json
import os
import re
from datetime import date
from typing import Callable

from google import genai
from google.genai import types

from ..models import HotelOption
from ..tools.browser import BrowserSession

_MODEL = "gemini-2.0-flash-lite"
_MAX_ITER = 30


def _system_prompt(
    destination: str, check_in: date, check_out: date,
    preferred_brands: list[str], nights: int, nights_paid: int, nights_free: int,
) -> str:
    brands = ", ".join(preferred_brands)
    return f"""You are an IHG hotel search assistant. Search ihg.com for reward-night availability.

## Search
Destination: {destination} | Check-in: {check_in} | Check-out: {check_out}
Nights: {nights} total ({nights_paid} paid, {nights_free} free via 4th-night-free benefit)
Preferred brands (in order): {brands}

## Steps
1. Go to https://www.ihg.com/hotels/us/en/find-hotels/hotel/list
2. Enter "{destination}", set dates {check_in} → {check_out}.
3. CRITICAL: Toggle "Use Points" ON before searching. Cash rates don't qualify.
4. Only include brands: {brands}
5. Only include hotels with points availability on ALL nights.
6. total_points = points_per_night × {nights_paid}

## Output — JSON array only, no prose, no fences:
[{{"destination":"{destination}","property_name":"<name>","brand":"<brand>",
"check_in":"{check_in}","check_out":"{check_out}",
"points_per_night":<int>,"total_points":<int>,
"nights_paid":{nights_paid},"nights_free":{nights_free},"cash_rate":<float|null>}}]
Return [] if nothing found."""


def _parse(text: str, destination: str, check_in: date, check_out: date) -> list[HotelOption]:
    m = re.search(r"```(?:json)?\s*([\s\S]*?)```", text)
    text = m.group(1).strip() if m else text
    s, e = text.find("["), text.rfind("]")
    if s == -1 or e == -1:
        return []
    try:
        items = json.loads(text[s : e + 1])
    except json.JSONDecodeError:
        return []
    out = []
    for item in items:
        try:
            for k in ("check_in", "check_out"):
                if isinstance(item.get(k), str):
                    item[k] = date.fromisoformat(item[k])
            out.append(HotelOption(**item))
        except Exception:
            pass
    return out


async def _search_one(
    destination: str, check_in: date, check_out: date,
    browser: BrowserSession, preferred_brands: list[str],
    nights: int, progress: Callable[[str], None],
) -> list[HotelOption]:
    nights_free = nights // 4
    nights_paid = nights - nights_free

    client = genai.Client(api_key=os.environ["GEMINI_API_KEY"])
    config = types.GenerateContentConfig(
        system_instruction=_system_prompt(
            destination, check_in, check_out, preferred_brands, nights, nights_paid, nights_free
        ),
        tools=browser.gemini_tools,
    )
    contents = [types.Content(role="user", parts=[types.Part(text=(
        f"Search ihg.com for reward-night availability in {destination} "
        f"from {check_in} to {check_out}. Toggle Use Points. Return JSON array."
    ))])]

    last_text = ""
    for i in range(_MAX_ITER):
        for attempt in range(4):
            try:
                response = await client.aio.models.generate_content(
                    model=_MODEL, contents=contents, config=config
                )
                break
            except Exception as e:
                if "429" in str(e) and attempt < 3:
                    await asyncio.sleep(15 * (attempt + 1))
                else:
                    raise
        contents.append(response.candidates[0].content)
        try:
            if response.text:
                last_text = response.text
        except ValueError:
            pass

        fn_calls = response.function_calls
        if not fn_calls:
            break

        fn_parts = []
        for fc in fn_calls:
            progress(f"IHG [{destination}]: {fc.name} (step {i+1})")
            result = await browser.execute_tool(fc.name, dict(fc.args))
            fn_parts.append(types.Part(
                function_response=types.FunctionResponse(name=fc.name, response={"result": result})
            ))
        contents.append(types.Content(role="user", parts=fn_parts))

    results = _parse(last_text, destination, check_in, check_out)
    progress(f"IHG [{destination}]: {len(results)} property(ies) found.")
    return results


async def search_points_availability(
    destinations: list[str],
    date_pairs: list[tuple[date, date]],
    browser: BrowserSession,
    preferred_brands: list[str],
    nights: int,
    progress: Callable[[str], None],
) -> list[HotelOption]:
    all_results: list[HotelOption] = []
    for destination in destinations:
        for check_in, check_out in date_pairs:
            progress(f"IHG: {destination} {check_in} → {check_out}")
            all_results.extend(await _search_one(
                destination, check_in, check_out, browser, preferred_brands, nights, progress
            ))
    progress(f"IHG: complete — {len(all_results)} total option(s).")
    return all_results