simple-agent/app/main.py

from __future__ import annotations

import asyncio
import logging
from datetime import datetime
from typing import Annotated

from apscheduler.schedulers.asyncio import AsyncIOScheduler
from fastapi import Depends, FastAPI, Header, HTTPException, Query, status
from pydantic import BaseModel

from app.calendar_agent import CalendarAvailabilityAgent
from app.config import get_settings
from app.gmail_agent import GmailTriageAgent
from app.google_clients import build_calendar_service, build_gmail_service

settings = get_settings()
logging.basicConfig(level=getattr(logging, settings.log_level.upper(), logging.INFO))
logger = logging.getLogger("personal-agent")

app = FastAPI(title="Personal Agent", version="0.1.0")
scheduler: AsyncIOScheduler | None = None
scan_lock: asyncio.Lock | None = None


class ScanResponse(BaseModel):
    scanned: int
    linkedin: int
    advertising: int
    skipped: int
    failed: int


class AvailabilityRequest(BaseModel):
    start: str
    end: str
    calendar_ids: list[str] | None = None


class BusySlot(BaseModel):
    calendar_id: str
    start: str
    end: str


class AvailabilityResponse(BaseModel):
    start: str
    end: str
    available: bool
    busy_slots: list[BusySlot]
    checked_calendars: list[str]


def verify_api_key(
    x_api_key: Annotated[str | None, Header(alias="X-API-Key")] = None,
    authorization: Annotated[str | None, Header()] = None,
) -> None:
    expected = settings.agent_api_key
    if not expected:
        return

    provided = x_api_key
    if not provided and authorization:
        parts = authorization.split(" ", 1)
        if len(parts) == 2 and parts[0].lower() == "bearer":
            provided = parts[1]

    if provided != expected:
        raise HTTPException(
            status_code=status.HTTP_401_UNAUTHORIZED, detail="Invalid API key."
        )


def _run_scan_once(max_results: int) -> ScanResponse:
    gmail_service = build_gmail_service(settings)
    gmail_agent = GmailTriageAgent(gmail_service=gmail_service, query=settings.gmail_query)
    result = gmail_agent.scan_and_route_messages(max_results=max_results)
    return ScanResponse(
        scanned=result.scanned,
        linkedin=result.linkedin,
        advertising=result.advertising,
        skipped=result.skipped,
        failed=result.failed,
    )


def _get_scan_lock() -> asyncio.Lock:
    global scan_lock
    if scan_lock is None:
        scan_lock = asyncio.Lock()
    return scan_lock


async def _scheduled_scan() -> None:
    lock = _get_scan_lock()
    if lock.locked():
        logger.info("Previous scan still running, skipping this tick.")
        return

    async with lock:
        try:
            result = await asyncio.to_thread(_run_scan_once, 100)
            logger.info("Scheduled scan complete: %s", result.model_dump())
        except Exception:
            logger.exception("Scheduled scan failed")


@app.on_event("startup")
async def startup_event() -> None:
    global scheduler
    _get_scan_lock()
    scheduler = AsyncIOScheduler()
    scheduler.add_job(
        _scheduled_scan,
        "interval",
        minutes=settings.gmail_scan_interval_minutes,
        next_run_time=datetime.now(),
    )
    scheduler.start()
    logger.info(
        "Scheduler started (interval=%s min)", settings.gmail_scan_interval_minutes
    )


@app.on_event("shutdown")
async def shutdown_event() -> None:
    if scheduler:
        scheduler.shutdown(wait=False)


@app.get("/health")
def health() -> dict[str, object]:
    return {
        "status": "ok",
        "scan_interval_minutes": settings.gmail_scan_interval_minutes,
    }


@app.post(
    "/scan",
    response_model=ScanResponse,
    dependencies=[Depends(verify_api_key)],
)
async def scan_now(max_results: int = Query(100, ge=1, le=500)) -> ScanResponse:
    async with _get_scan_lock():
        try:
            return await asyncio.to_thread(_run_scan_once, max_results)
        except FileNotFoundError as exc:
            raise HTTPException(
                status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
                detail=str(exc),
            ) from exc
        except Exception as exc:
            raise HTTPException(
                status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
                detail=f"Gmail scan failed: {exc}",
            ) from exc


@app.post(
    "/availability",
    response_model=AvailabilityResponse,
    dependencies=[Depends(verify_api_key)],
)
async def availability(request: AvailabilityRequest) -> AvailabilityResponse:
    try:
        calendar_service = build_calendar_service(settings)
        availability_agent = CalendarAvailabilityAgent(calendar_service=calendar_service)
        result = await asyncio.to_thread(
            availability_agent.get_availability,
            request.start,
            request.end,
            request.calendar_ids,
        )
        return AvailabilityResponse(
            start=result.start,
            end=result.end,
            available=result.available,
            busy_slots=result.busy_slots,
            checked_calendars=result.checked_calendars,
        )
    except ValueError as exc:
        raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=str(exc)) from exc
    except FileNotFoundError as exc:
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            detail=str(exc),
        ) from exc
    except Exception as exc:
        raise HTTPException(
            status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
            detail=f"Availability lookup failed: {exc}",
        ) from exc