api/schedule.py

"""Schedule visit routes — find matches and book via LangGraph agent."""
import uuid
import threading
from fastapi import APIRouter, Depends, HTTPException
from pydantic import BaseModel
from typing import Any
from core import data_manager
from core import settings_manager as settings
from agent.graph import build_scheduling_graph
from api.deps import require_admin_session

router = APIRouter(prefix="/api/schedule", tags=["schedule"])

_graph = None
_graph_lock = threading.Lock()


def _get_graph():
    global _graph
    with _graph_lock:
        if _graph is None:
            _graph = build_scheduling_graph()
        return _graph


# In-memory thread-id store for active scheduling sessions.
_threads: dict[str, dict] = {}


class FindMatchesRequest(BaseModel):
    patient_id: str
    visit_level: Any  # int or "Nurse"
    visit_address: str
    preferred_dates: list[str]
    preferred_time_window: dict  # {"start": "09:00", "end": "12:00"}


class BookRequest(BaseModel):
    thread_id: str
    selection_index: int


@router.get("/visit-levels")
def visit_levels():
    return [{"key": k, "label": v["label"], "duration_minutes": v["duration_minutes"]} for k, v in settings.get_visit_levels().items()]


@router.get("/config")
def schedule_config():
    """Public scheduling config used by the date/time picker."""
    return {
        "min_lead_time_hours": settings.get_min_lead_time_hours(),
        "timezone": settings.get_timezone(),
    }


@router.post("/find-matches")
def find_matches(req: FindMatchesRequest, session: dict = Depends(require_admin_session)):
    patient = data_manager.get_patient_by_id(req.patient_id)
    if not patient:
        raise HTTPException(status_code=404, detail="Patient not found")

    # Convert visit_level to int if numeric string
    vl: Any = req.visit_level
    if isinstance(vl, str) and vl.isdigit():
        vl = int(vl)

    initial_state = {
        "patient_id": req.patient_id,
        "patient_name": patient["name"],
        "patient_email": patient.get("email", ""),
        "visit_level": vl,
        "visit_address": req.visit_address,
        "visit_lat": patient.get("lat"),
        "visit_lng": patient.get("lng"),
        "preferred_dates": req.preferred_dates,
        "preferred_time_window": req.preferred_time_window,
        "available_providers": [],
        "travel_times": [],
        "ranked_matches": [],
        "user_selection": None,
        "booked_appointment": None,
        "notifications_sent": False,
        "error": None,
    }

    thread_id = str(uuid.uuid4())
    config = {"configurable": {"thread_id": thread_id}}
    result = _get_graph().invoke(initial_state, config)

    _threads[thread_id] = {"config": config, "ranked_matches": result.get("ranked_matches", [])}

    if result.get("error"):
        return {"thread_id": thread_id, "error": result["error"], "ranked_matches": []}
    return {
        "thread_id": thread_id,
        "ranked_matches": result.get("ranked_matches", []),
        "error": None,
    }


@router.post("/book")
def book(req: BookRequest, session: dict = Depends(require_admin_session)):
    record = _threads.get(req.thread_id)
    if not record:
        raise HTTPException(status_code=400, detail="Invalid or expired thread_id")
    matches = record["ranked_matches"]
    if req.selection_index < 0 or req.selection_index >= len(matches):
        raise HTTPException(status_code=400, detail="Invalid selection_index")
    match = matches[req.selection_index]
    user_selection = {
        "provider_id": match["provider_id"],
        "provider_name": match["provider_name"],
        "provider_email": match.get("provider_email", ""),
        "date": match["date"],
        "start_time": match["start_time"],
        "end_time": match["end_time"],
        "recommended_departure_time": match.get("recommended_departure_time", ""),
    }
    graph = _get_graph()
    config = record["config"]
    graph.update_state(config, {"user_selection": user_selection})
    result = graph.invoke(None, config)

    # Clean up thread
    _threads.pop(req.thread_id, None)

    if result.get("error"):
        return {"error": result["error"]}
    return {
        "booked_appointment": result.get("booked_appointment"),
        "notifications_sent": result.get("notifications_sent", False),
    }