agent/nodes.py

"""
LangGraph Node Functions — each function represents one step in the scheduling graph.

Each node takes the current SchedulingState and returns a partial state update (dict).
"""
from agent.state import SchedulingState
from core import data_manager, scheduling_engine, google_services
from core import settings_manager as settings
from datetime import datetime, timedelta
from zoneinfo import ZoneInfo


# Lazy-initialize the LLM to avoid import-time failures if the API key isn't set yet.
_llm = None

def _get_llm():
    global _llm
    if _llm is None:
        from langchain_openai import ChatOpenAI
        _llm = ChatOpenAI(model=settings.get_openai_model())
    return _llm


def validate_inputs(state: SchedulingState) -> dict:
    """Validate all inputs from the dashboard before proceeding.

    Checks:
        - Patient exists
        - Visit level is valid
        - Address is not empty
        - Preferred dates are at least 24 hours in the future
        - Time window is valid (start < end, within reason)

    Returns:
        {"error": "message"} if validation fails, otherwise empty dict {}
    """
    # Patient must exist
    patient = data_manager.get_patient_by_id(state["patient_id"])
    if not patient:
        return {"error": f"Patient '{state['patient_id']}' not found."}

    # Visit level must be in the configured set
    if state["visit_level"] not in settings.get_visit_levels():
        return {"error": f"Invalid visit level: {state['visit_level']}"}

    # Visit address must not be blank
    if not state.get("visit_address", "").strip():
        return {"error": "Visit address is required."}

    # At least one preferred date must be provided
    if not state.get("preferred_dates"):
        return {"error": "At least one preferred date is required."}

    # Every preferred date must be at least the configured lead time from now.
    # Lead time always applies, even for on-call providers.
    # Compare in the configured timezone so a UTC server doesn't mis-evaluate local times.
    lead_hours = settings.get_min_lead_time_hours()
    now_local = datetime.now(ZoneInfo(settings.get_timezone())).replace(tzinfo=None)
    min_allowed_dt = now_local + timedelta(hours=lead_hours)
    window_start = state.get("preferred_time_window", {}).get("start", "08:00")
    for date_str in state["preferred_dates"]:
        try:
            proposed_dt = datetime.strptime(f"{date_str} {window_start}", "%Y-%m-%d %H:%M")
        except ValueError:
            return {"error": f"Invalid date format: '{date_str}'. Expected YYYY-MM-DD."}
        if proposed_dt < min_allowed_dt:
            return {"error": f"Date {date_str} at {window_start} does not meet the {lead_hours}-hour lead time requirement."}

    # Time window must have start before end
    window = state.get("preferred_time_window", {})
    w_start = window.get("start", "")
    w_end = window.get("end", "")
    if w_start and w_end:
        if scheduling_engine.parse_time(w_start) >= scheduling_engine.parse_time(w_end):
            return {"error": "Time window start must be before end."}

    return {}


def check_availability(state: SchedulingState) -> dict:
    """Check which providers are available for the requested visit.

    For each preferred date/time combo:
        1. Get all active providers
        2. Get each provider's existing appointments for that date
        3. Run scheduling_engine.is_provider_available for each

    Returns:
        {"available_providers": [...]} — list of provider dicts with an added 'available_slots' field
        If none available: {"available_providers": [], "error": "No providers available..."}
    """
    # Skip if a previous node already set an error
    if state.get("error"):
        return {"available_providers": []}

    all_providers = data_manager.get_all_providers(active_only=True)
    visit_level = state["visit_level"]
    visit_duration = scheduling_engine.get_visit_duration(visit_level)
    window_start = state["preferred_time_window"]["start"]
    window_end = state["preferred_time_window"]["end"]

    available = []
    for date in state["preferred_dates"]:
        # Build a per-provider appointment index for this date (one read per provider)
        appointments_by_provider = {
            p["id"]: data_manager.get_appointments_for_provider(p["id"], date)
            for p in all_providers
        }

        for provider in all_providers:
            existing = appointments_by_provider[provider["id"]]
            # Iterate through the time window in 30-minute increments
            current_time = window_start
            while scheduling_engine.parse_time(current_time) < scheduling_engine.parse_time(window_end):
                if scheduling_engine.is_provider_available(
                    provider, date, current_time, visit_level, existing
                ):
                    end_time = scheduling_engine.add_minutes(current_time, visit_duration)
                    enriched = dict(provider)
                    enriched["available_date"] = date
                    enriched["available_start_time"] = current_time
                    enriched["available_end_time"] = end_time
                    enriched["is_oncall"] = scheduling_engine.is_oncall_for_slot(
                        provider, date, current_time, end_time
                    )
                    available.append(enriched)
                current_time = scheduling_engine.add_minutes(current_time, 30)

    if not available:
        return {
            "available_providers": [],
            "error": (
                "No providers are available for the requested dates and time. "
                "Try different dates or a different time window."
            ),
        }

    return {"available_providers": available}


def calculate_travel(state: SchedulingState) -> dict:
    """Calculate travel times for all available providers to the patient's address.

    For each available provider:
        1. Determine their "origin" — the address of their previous appointment that day,
           or their home address if it's their first visit
        2. Call google_services.calculate_travel_time(origin, visit_address)
        3. Also check if provider has other visits nearby (for area-affinity boost)

    Returns:
        {"travel_times": [...]} — provider dicts enriched with 'travel_time_minutes' and 'has_nearby_visits'
    """
    visit_address = state["visit_address"]
    enriched_providers = []

    for provider in state["available_providers"]:
        available_date = provider["available_date"]
        available_start = provider["available_start_time"]
        arrival_by = datetime.strptime(f"{available_date} {available_start}", "%Y-%m-%d %H:%M")

        # Fetch existing appointments for this provider on the proposed date
        existing = data_manager.get_appointments_for_provider(provider["id"], available_date)

        # Origin = home address, or the last appointment address that ends before proposed start
        origin = provider["address"]
        for appt in sorted(existing, key=lambda a: a["end_time"]):
            if appt["end_time"] <= available_start:
                origin = appt["visit_address"]

        # Travel time from origin to patient visit address (traffic-aware two-pass)
        travel_result = google_services.calculate_travel_time(origin, visit_address, arrival_by)
        travel_time = travel_result["travel_minutes"]
        if travel_time < 0:
            travel_time = 0  # Graceful fallback if the Maps API fails
        recommended_departure = travel_result["recommended_departure"]

        # Area affinity: does the provider have any visits close to the patient that day?
        has_nearby = False
        if existing:
            appt_addresses = [a["visit_address"] for a in existing]
            nearby_results = google_services.calculate_travel_times_batch(
                visit_address, appt_addresses, arrival_by
            )
            # "Nearby" = within 15 minutes of the proposed visit address
            has_nearby = any(r["travel_minutes"] >= 0 and r["travel_minutes"] <= 15 for r in nearby_results)

        enriched = dict(provider)
        enriched["travel_time_minutes"] = travel_time
        enriched["has_nearby_visits"] = has_nearby
        enriched["recommended_departure_time"] = recommended_departure.strftime("%H:%M")
        enriched_providers.append(enriched)

    return {"travel_times": enriched_providers}


def rank_and_present(state: SchedulingState) -> dict:
    """Rank available providers by travel time and area affinity, then format for display.

    Uses scheduling_engine.rank_providers_by_travel.

    Returns:
        {"ranked_matches": [...]} — sorted list of match dicts with fields:
            provider_id, provider_name, provider_email, date, start_time, end_time,
            travel_time_minutes, has_nearby_visits, summary
    """
    ranked = scheduling_engine.rank_providers_by_travel(state["travel_times"])
    level_label = settings.get_visit_levels()[state["visit_level"]]["label"]

    matches = []
    for provider in ranked:
        is_oncall = provider.get("is_oncall", False)
        nearby_note = ", nearby visits" if provider["has_nearby_visits"] else ""
        departure = provider.get("recommended_departure_time", "")
        departure_note = f", depart by {departure}" if departure else ""
        oncall_note = "[ON CALL] " if is_oncall else ""
        summary = (
            f"{oncall_note}{provider['name']} — {level_label} on {provider['available_date']} "
            f"at {provider['available_start_time']}–{provider['available_end_time']} "
            f"({provider['travel_time_minutes']} min travel{nearby_note}{departure_note})"
        )
        matches.append({
            "provider_id": provider["id"],
            "provider_name": provider["name"],
            "provider_email": provider.get("email", ""),
            "date": provider["available_date"],
            "start_time": provider["available_start_time"],
            "end_time": provider["available_end_time"],
            "travel_time_minutes": provider["travel_time_minutes"],
            "has_nearby_visits": provider["has_nearby_visits"],
            "is_oncall": is_oncall,
            "recommended_departure_time": departure,
            "summary": summary,
        })

    return {"ranked_matches": matches}


def book_appointment(state: SchedulingState) -> dict:
    """Book the selected appointment.

    Steps:
        1. Final overlap check (in case something changed)
        2. Create appointment record via data_manager.add_appointment
        3. Create Google Calendar event via google_services.create_calendar_event
        4. Update appointment record with google_event_id
        5. Log action to audit log

    Returns:
        {"booked_appointment": {...}} or {"error": "Booking failed: ..."}
    """
    selection = state["user_selection"]
    provider_id = selection["provider_id"]
    date = selection["date"]
    start_time = selection["start_time"]

    # Final freshness check: re-fetch the provider and their current appointments
    provider = data_manager.get_provider_by_id(provider_id)
    if not provider:
        return {"error": f"Booking failed: provider '{provider_id}' not found."}

    existing = data_manager.get_appointments_for_provider(provider_id, date)
    if not scheduling_engine.is_provider_available(
        provider, date, start_time, state["visit_level"], existing
    ):
        return {"error": "Booking failed: the selected slot is no longer available."}

    # Persist the appointment record
    appointment_data = {
        "patient_id": state["patient_id"],
        "patient_name": state["patient_name"],
        "provider_id": provider_id,
        "provider_name": selection["provider_name"],
        "date": date,
        "start_time": start_time,
        "end_time": selection["end_time"],
        "visit_level": state["visit_level"],
        "visit_address": state["visit_address"],
        "recommended_departure_time": selection.get("recommended_departure_time", ""),
        "status": "scheduled",
        "google_event_id": None,
    }
    appointment = data_manager.add_appointment(appointment_data)

    # Create Google Calendar event (non-critical — failure doesn't block booking)
    try:
        level_label = settings.get_visit_levels()[state["visit_level"]]["label"]
        event_summary = f"[{level_label}] {state['patient_name']} — {selection['provider_name']}"
        description = (
            f"Patient: {state['patient_name']}\n"
            f"Provider: {selection['provider_name']}\n"
            f"Visit Level: {level_label}\n"
            f"Address: {state['visit_address']}"
        )
        event_id = google_services.create_calendar_event(
            summary=event_summary,
            description=description,
            location=state["visit_address"],
            start_datetime=f"{date}T{start_time}:00",
            end_datetime=f"{date}T{selection['end_time']}:00",
        )
        data_manager.update_appointment(appointment["id"], {"google_event_id": event_id})
        appointment["google_event_id"] = event_id
    except Exception:
        pass

    # Audit log
    data_manager.log_action(
        user="admin",
        action="SCHEDULE",
        details=(
            f"Scheduled {settings.get_visit_levels()[state['visit_level']]['label']} visit for "
            f"{state['patient_name']} with {selection['provider_name']} "
            f"on {date} at {start_time}"
        ),
        appointment_id=appointment["id"],
    )

    return {"booked_appointment": appointment}


def send_notifications(state: SchedulingState) -> dict:
    """Send confirmation emails to both the provider and patient.

    Uses google_services.send_scheduling_notification.

    Returns:
        {"notifications_sent": True/False}
    """
    appointment = state.get("booked_appointment")
    if not appointment:
        return {"notifications_sent": False}

    selection = state["user_selection"]
    appointment_details = {
        "patient_name": state["patient_name"],
        "provider_name": selection["provider_name"],
        "date": appointment["date"],
        "start_time": appointment["start_time"],
        "end_time": appointment["end_time"],
        "visit_level": settings.get_visit_levels()[state["visit_level"]]["label"],
        "address": state["visit_address"],
        "recommended_departure_time": appointment.get("recommended_departure_time", ""),
    }

    try:
        result = google_services.send_scheduling_notification(
            provider_email=selection.get("provider_email", ""),
            patient_email=state.get("patient_email", ""),
            appointment_details=appointment_details,
        )
        sent = result.get("provider_sent", False) or result.get("patient_sent", False)
    except Exception:
        sent = False

    return {"notifications_sent": sent}


def no_match(state: SchedulingState) -> dict:
    """Handle the case where no providers are available.

    Uses the LLM to generate a helpful message suggesting alternative dates/times.

    Returns:
        {"error": "No matching providers found. <suggestion>"}
    """
    existing_error = state.get("error", "No providers are available for the requested visit.")

    try:
        llm = _get_llm()
        preferred_dates = state.get("preferred_dates", [])
        visit_level = state.get("visit_level", "")
        level_label = settings.get_visit_levels().get(visit_level, {}).get("label", str(visit_level))
        window = state.get("preferred_time_window", {})

        prompt = (
            f"A home healthcare {level_label} visit could not be scheduled. "
            f"Requested dates: {', '.join(preferred_dates)}. "
            f"Requested time window: {window.get('start', '?')}–{window.get('end', '?')}. "
            f"Provide 2–3 brief, practical suggestions the scheduler can try next "
            f"(e.g., different dates, expanded time window, or different visit level if appropriate). "
            f"Keep the response to 2 sentences maximum."
        )

        response = llm.invoke(prompt)
        suggestion = response.content.strip()
    except Exception:
        suggestion = (
            "Try selecting different dates, expanding the time window, "
            "or contacting the scheduling team for assistance."
        )

    return {"error": f"{existing_error} {suggestion}"}