fastapi_app.py

"""FastAPI service exposing wage-search functionality."""

import asyncio
import logging
from pathlib import Path
from typing import Optional

from fastapi import FastAPI, HTTPException, Query
from fastapi.responses import JSONResponse
from fastapi.staticfiles import StaticFiles

from api._shared import get_dataset
from wage_search.dataset import DatasetError, WageDataset

DEFAULT_ARCHIVE = Path("data/current/OFLC_Wages_2025-26.zip")
STATIC_DIR = Path(__file__).resolve().parent / "docs"
logger = logging.getLogger(__name__)


app = FastAPI(
    title="Wage Search API",
    description="API surface for OFLC prevailing wage data matching the CLI and serverless routes.",
    version="0.1.0",
)


def dataset_or_error(dataset_year: Optional[str] = None) -> WageDataset:
    try:
        return get_dataset(dataset_year)
    except DatasetError as exc:
        raise HTTPException(status_code=400, detail=str(exc)) from exc


@app.on_event("startup")
async def warm_dataset() -> None:
    loop = asyncio.get_running_loop()

    def _load() -> None:
        try:
            get_dataset(None)
        except DatasetError as exc:  # pragma: no cover - logged for visibility
            logger.error("Unable to warm dataset: %s", exc)

    await loop.run_in_executor(None, _load)


@app.get("/health", summary="Service health check")
async def health() -> dict:
    return {"status": "ok"}


@app.get("/metadata", summary="Dataset metadata")
@app.get("/api/metadata", include_in_schema=False)
async def metadata(
    dataset_year: Optional[str] = Query(None, description="Optional dataset end-year override"),
) -> dict:
    dataset = dataset_or_error(dataset_year=dataset_year)
    return {
        "datasetYear": dataset.dataset_year,
        "archive": str(DEFAULT_ARCHIVE),
        "context": "local",
        "states": dataset.state_metadata(),
    }


@app.get("/suggest", summary="SOC lookup suggestions")
@app.get("/api/suggest", include_in_schema=False)
async def suggest(
    search: str = Query(..., min_length=2, description="Search phrase for SOC titles/codes"),
    top: int = Query(10, ge=1, le=25, description="Number of results to return"),
    dataset_year: Optional[str] = Query(None),
) -> JSONResponse:
    dataset = dataset_or_error(dataset_year=dataset_year)
    try:
        payload = dataset.suggest(search, top)
    except DatasetError as exc:
        raise HTTPException(status_code=400, detail=str(exc)) from exc
    return JSONResponse(payload)


@app.get("/areas/options", summary="List BLS areas for a state")
@app.get("/api/areas/options", include_in_schema=False)
async def list_area_options(
    state: str = Query(..., min_length=2, description="State abbreviation or name"),
    year: str = Query(..., description="Dataset end-year to validate"),
    dataset_year: Optional[str] = Query(None),
) -> JSONResponse:
    dataset = dataset_or_error(dataset_year=dataset_year)
    try:
        payload = dataset.list_area_options(state, year)
    except DatasetError as exc:
        raise HTTPException(status_code=404, detail=str(exc)) from exc
    return JSONResponse(payload)

@app.get("/areas/counties", summary="List counties for a CBSA")
@app.get("/api/areas/counties", include_in_schema=False)
async def list_counties(
    cbsa: int = Query(..., description="CBSA/area identifier"),
    state: Optional[str] = Query(None, description="Optional state abbreviation or name"),
    dataset_year: Optional[str] = Query(None),
) -> JSONResponse:
    dataset = dataset_or_error(dataset_year=dataset_year)
    try:
        payload = dataset.list_counties(cbsa, state)
    except DatasetError as exc:
        raise HTTPException(status_code=404, detail=str(exc)) from exc
    return JSONResponse(payload)


@app.get("/wage-search", summary="Retrieve wages for a SOC and CBSA")
@app.get("/api/wage-search", include_in_schema=False)
async def wage_search(
    collection_type: str = Query(..., description="Collection type (alc or acwia)"),
    year: str = Query(..., description="Dataset end-year to validate"),
    soc_code: str = Query(..., description="SOC code (e.g. 15-1252.00)"),
    area: int = Query(..., description="Area identifier / CBSA"),
    area_type: str = Query(..., description="Area type (bls_area or county_town)"),
    rd_flag: str = Query("BOTH", description="R&D flag (BOTH, RD, NON-RD)"),
    dataset_year: Optional[str] = Query(None),
) -> JSONResponse:
    dataset = dataset_or_error(dataset_year=dataset_year)
    try:
        payload = dataset.wage_search(
            collection_type=collection_type,
            year=year,
            soc_code=soc_code,
            area=area,
            area_type=area_type,
            rd_flag=rd_flag,
        )
    except DatasetError as exc:
        raise HTTPException(status_code=404, detail=str(exc)) from exc
    return JSONResponse(payload)


def _parse_csv(value: Optional[str]) -> list[str]:
    if not value:
        return []
    return [item.strip() for item in value.split(",") if item.strip()]


def _parse_ints(items: list[str]) -> list[int]:
    parsed: list[int] = []
    for item in items:
        try:
            parsed.append(int(item))
        except ValueError as exc:
            raise HTTPException(status_code=400, detail=f"Invalid area id '{item}'") from exc
    return parsed


@app.get("/wage-table", summary="Retrieve wages for multiple areas")
@app.get("/api/wage-table", include_in_schema=False)
async def wage_table(
    collection_type: str = Query("alc", description="Collection type (alc or acwia)"),
    area_type: str = Query("bls_area", description="Area type (bls_area or county_town)"),
    rd_flag: str = Query("BOTH", description="R&D flag (BOTH, RD, NON-RD)"),
    soc_code: str = Query(..., description="SOC code (e.g. 15-1252.00)"),
    year: Optional[str] = Query(None, description="Dataset end-year to validate"),
    areas: Optional[str] = Query(None, description="Comma-separated CBSA identifiers"),
    states: Optional[str] = Query(
        None,
        description="Comma-separated state abbreviations or names. Use '*' or 'ALL' for every state.",
    ),
    limit: Optional[int] = Query(
        None, ge=1, description="Optional limit applied after expansion of states/areas."
    ),
    dataset_year: Optional[str] = Query(None),
) -> JSONResponse:
    dataset = dataset_or_error(dataset_year=dataset_year)
    dataset_year_value = dataset.dataset_year
    target_year = year or dataset_year_value

    state_tokens = [token.upper() for token in _parse_csv(states)]
    area_tokens = _parse_csv(areas)

    requested_areas: list[int] = []
    if state_tokens:
        if any(token in {"*", "ALL"} for token in state_tokens):
            requested_areas.extend(dataset.area_codes_for_states(dataset.list_states()))
        else:
            requested_areas.extend(dataset.area_codes_for_states(state_tokens))

    if area_tokens:
        requested_areas.extend(_parse_ints(area_tokens))

    if not requested_areas:
        raise HTTPException(
            status_code=400, detail="Provide at least one state (states=) or area id (areas=)."
        )

    if limit is not None:
        requested_areas = requested_areas[:limit]

    payload = dataset.bulk_wage_table(
        collection_type=collection_type,
        year=target_year,
        soc_code=soc_code,
        areas=requested_areas,
        area_type=area_type,
        rd_flag=rd_flag,
    )
    return JSONResponse(payload)

if STATIC_DIR.exists():
    app.mount("/", StaticFiles(directory=STATIC_DIR, html=True), name="site")