backend/agents/delta_agent.py

"""
Delta.com browser agent — finds companion certificate availability.
Uses google-genai (Gemini Flash free tier) with Playwright browser tools.
"""

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

from google import genai
from google.genai import types

from ..models import Cabin, FlightOption, TripSearchRequest
from ..tools.browser import BrowserSession

_MODEL = "gemini-2.0-flash-lite"
_MAX_ITER = 30
_ECONOMY_CLASSES = ["L", "U", "T", "X", "V"]
_FIRST_CLASSES = ["Z", "D", "I"]


def _system_prompt(request: TripSearchRequest) -> str:
    cabin = "economy" if request.preferred_cabin == Cabin.economy else "first class"
    eligible = _ECONOMY_CLASSES if request.preferred_cabin == Cabin.economy else _FIRST_CLASSES
    dests = ", ".join(request.destinations)
    return f"""You are a travel research assistant navigating delta.com to find companion certificate availability.

## Goal
Search delta.com for flights from {request.origin} to: {dests}.
Date window: {request.earliest_departure} to {request.latest_return}.
Trip length: ~{request.trip_duration_nights} nights. Cabin: {cabin}.

## Companion Certificate Rules
- Companion flies for taxes only (~$80).
- Economy eligible fare classes: L, U, T, X, V
- First class eligible fare classes: Z, D, I
- Hawaii (HNL) is domestic — eligible.
- Use the Price Calendar to find eligible dates.

## Steps
1. User is already logged in — do NOT log in.
2. Go to My Profile → Certificates, eCredits & Vouchers. Confirm cert is present.
3. For each destination: search round-trip with flexible dates, use Price Calendar,
   find dates with eligible fare classes ({", ".join(eligible)}), note outbound date,
   return date, fare class, and base price.

## Output
After searching ALL destinations, output ONLY a JSON array, no prose, no fences:
[{{"destination":"<IATA>","outbound_date":"<YYYY-MM-DD>","return_date":"<YYYY-MM-DD>",
"fare_class":"<letter>","base_price":<float>,"companion_taxes":80.0,
"eligible_for_companion_cert":true,"cabin":"{cabin}"}}]
Omit destinations with no eligible flights. Output JSON only."""


def _user_prompt(request: TripSearchRequest) -> str:
    return (
        f"Search delta.com companion cert availability: {request.origin} → "
        f"{', '.join(request.destinations)}. "
        f"Dates {request.earliest_departure} – {request.latest_return}, "
        f"~{request.trip_duration_nights} nights. Return JSON array."
    )


def _parse(text: str) -> list[FlightOption]:
    text = re.sub(r"```(?:json)?\s*", "", text).strip().rstrip("`").strip()
    s, e = text.find("["), text.rfind("]")
    if s == -1 or e <= s:
        return []
    try:
        items = json.loads(text[s : e + 1])
    except json.JSONDecodeError:
        return []
    out = []
    for item in items:
        try:
            out.append(FlightOption(**item))
        except Exception:
            pass
    return out


async def search_companion_cert(
    request: TripSearchRequest,
    browser: BrowserSession,
    progress: Callable[[str], None],
) -> list[FlightOption]:
    client = genai.Client(api_key=os.environ["GEMINI_API_KEY"])
    config = types.GenerateContentConfig(
        system_instruction=_system_prompt(request),
        tools=browser.gemini_tools,
    )
    contents = [types.Content(role="user", parts=[types.Part(text=_user_prompt(request))])]

    progress(f"Delta: searching {request.origin} → {', '.join(request.destinations)}")

    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:
            progress("Delta: navigation complete.")
            break

        fn_parts = []
        for fc in fn_calls:
            progress(f"Delta [{i+1}/{_MAX_ITER}]: {fc.name}")
            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))
    else:
        progress(f"Delta: reached max iterations, parsing partial results.")

    options = _parse(last_text)
    progress(f"Delta: found {len(options)} eligible flight(s).")
    return options