Compressor/backend/app/main.py

from typing import Any

from fastapi import Depends, FastAPI, Header, HTTPException, Query
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, Field

from .auth import create_token, revoke_token, validate_token, verify_credentials
from .config import settings
from .db import ensure_annotation_table
from .services.data_service import data_service
from .services.pretrain_service import pretrain_service


class TimePointInput(BaseModel):
    index: int = Field(ge=0)
    sampleTime: str
    files: dict[str, Any]
    siteValues: dict[str, float | None] = Field(default={})


class WaveWindowInput(BaseModel):
    devicePart: str = Field(min_length=1)
    devicePoints: list[str]
    points: list[TimePointInput]
    maxPoints: int = Field(default=200000, ge=256, le=200000)
    noSampling: bool = False
    firstCycleOnly: bool = False
    sitePoints: list[str] = Field(default=[])


class AnnotationInput(BaseModel):
    wave_file_id: int = Field(gt=0)
    label: str = Field(min_length=1, max_length=8)
    period_start: int = Field(gt=0)
    period_end: int = Field(gt=0)
    sample_index_start: int = Field(ge=0)
    sample_index_end: int = Field(ge=0)


class AnnotationQueryInput(BaseModel):
    wave_file_ids: list[int] = Field(default=[])


class PretrainRecordInput(BaseModel):
    id: int = Field(gt=0)
    point_name: str = Field(min_length=1)
    measurement_type: str = Field(min_length=1)
    rpm: float | None = None
    sample_time: str = Field(min_length=1)
    file_name: str = Field(default="")


class LoginInput(BaseModel):
    username: str = Field(min_length=1)
    password: str = Field(min_length=1)


app = FastAPI(
    title="压缩机故障预测波形服务",
    version="0.1.0",
)
app.add_middleware(
    CORSMiddleware,
    allow_origins=[settings.cors_origin, "http://127.0.0.1:5173"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


def require_auth(authorization: str | None = Header(default=None)) -> str:
    token = (authorization or "").removeprefix("Bearer ").strip()
    if not validate_token(token):
        raise HTTPException(status_code=401, detail="未登录或登录已过期")
    return token


@app.on_event("startup")
def startup() -> None:
    try:
        ensure_annotation_table()
    except Exception:
        # 数据库不可达时跳过建表，演示模式仍可正常启动。
        pass


@app.post("/api/login")
def login(payload: LoginInput) -> dict[str, Any]:
    if not verify_credentials(payload.username, payload.password):
        raise HTTPException(status_code=401, detail="用户名或密码错误")
    return {"token": create_token(), "username": payload.username}


@app.post("/api/logout")
def logout(authorization: str | None = Header(default=None)) -> dict[str, Any]:
    token = (authorization or "").removeprefix("Bearer ").strip()
    revoke_token(token)
    return {"ok": True}


@app.get("/api/health")
def health(_auth: str = Depends(require_auth)) -> dict[str, Any]:
    return data_service.health()


@app.get("/api/query-options")
def query_options(_auth: str = Depends(require_auth)) -> dict[str, Any]:
    try:
        result = data_service.query_options()
        # Counts are keyed by the full point_name so the UI can scope them to
        # the currently selected device_part.
        result["pretrainCounts"] = pretrain_service.recorded_counts()
        result["abnormalCounts"] = data_service.abnormal_counts()
        return result
    except Exception as error:
        raise HTTPException(status_code=503, detail=str(error)) from error


@app.get("/api/time-points")
def time_points(
    device_part: str = Query(min_length=1),
    device_points: list[str] = Query(default=[]),
    min_time: str | None = None,
    max_time: str | None = None,
    include_stopped: bool = False,
    min_status: int | None = Query(default=None, ge=0),
    status_filter: list[str] = Query(default=[]),
    site_points: list[str] = Query(default=[]),
    _auth: str = Depends(require_auth),
) -> dict[str, Any]:
    try:
        return data_service.time_points(
            device_part,
            device_points,
            min_time,
            max_time,
            include_stopped,
            min_status,
            status_filter,
            site_points,
        )
    except ValueError as error:
        raise HTTPException(status_code=400, detail=str(error)) from error
    except Exception as error:
        raise HTTPException(status_code=503, detail=str(error)) from error


@app.get("/api/site-points")
def site_points(
    device_part: str = Query(min_length=1),
    _auth: str = Depends(require_auth),
) -> dict[str, Any]:
    try:
        return data_service.site_points(device_part)
    except Exception as error:
        raise HTTPException(status_code=503, detail=str(error)) from error


@app.get("/api/tspluse-ruler")
def tspluse_ruler(_auth: str = Depends(require_auth)) -> dict[str, Any]:
    try:
        return data_service.tspluse_ruler()
    except Exception as error:
        raise HTTPException(status_code=503, detail=str(error)) from error


@app.post("/api/wave-window")
def wave_window(payload: WaveWindowInput, _auth: str = Depends(require_auth)) -> dict[str, Any]:
    try:
        return data_service.wave_window(
            payload.devicePart,
            payload.devicePoints,
            [point.model_dump() if hasattr(point, "model_dump") else point.dict() for point in payload.points],
            payload.maxPoints,
            payload.noSampling,
            payload.firstCycleOnly,
            payload.sitePoints,
        )
    except ValueError as error:
        raise HTTPException(status_code=400, detail=str(error)) from error
    except Exception as error:
        raise HTTPException(status_code=503, detail=str(error)) from error


@app.get("/api/wave-files/{wave_file_id}/periods/{period_number}")
def period_detail(
    wave_file_id: int,
    period_number: int,
    _auth: str = Depends(require_auth),
) -> dict[str, Any]:
    try:
        return data_service.period_detail(wave_file_id, period_number)
    except ValueError as error:
        raise HTTPException(status_code=400, detail=str(error)) from error
    except Exception as error:
        raise HTTPException(status_code=503, detail=str(error)) from error


@app.get("/api/annotation-config")
def annotation_config(_auth: str = Depends(require_auth)) -> dict[str, Any]:
    return data_service.annotation_config()


@app.get("/api/annotations")
def list_annotations(
    wave_file_ids: list[int] = Query(default=[]),
    _auth: str = Depends(require_auth),
) -> dict[str, Any]:
    try:
        return data_service.list_annotations(wave_file_ids)
    except ValueError as error:
        raise HTTPException(status_code=400, detail=str(error)) from error
    except Exception as error:
        raise HTTPException(status_code=503, detail=str(error)) from error


@app.post("/api/annotations/query")
def query_annotations(
    payload: AnnotationQueryInput,
    _auth: str = Depends(require_auth),
) -> dict[str, Any]:
    try:
        return data_service.list_annotations(payload.wave_file_ids)
    except ValueError as error:
        raise HTTPException(status_code=400, detail=str(error)) from error
    except Exception as error:
        raise HTTPException(status_code=503, detail=str(error)) from error


@app.post("/api/annotations")
def create_annotation(payload: AnnotationInput, _auth: str = Depends(require_auth)) -> dict[str, Any]:
    try:
        return data_service.create_annotation(payload.model_dump())
    except ValueError as error:
        raise HTTPException(status_code=400, detail=str(error)) from error
    except Exception as error:
        raise HTTPException(status_code=503, detail=str(error)) from error


@app.delete("/api/annotations/{annotation_id}")
def delete_annotation(annotation_id: int, _auth: str = Depends(require_auth)) -> dict[str, Any]:
    try:
        return data_service.delete_annotation(annotation_id)
    except ValueError as error:
        raise HTTPException(status_code=404, detail=str(error)) from error
    except Exception as error:
        raise HTTPException(status_code=503, detail=str(error)) from error


@app.get("/api/pretrain-records")
def pretrain_records_status(
    ids: list[int] = Query(default=[]),
    _auth: str = Depends(require_auth),
) -> dict[str, Any]:
    try:
        return {"recorded": pretrain_service.recorded_ids(ids)}
    except Exception as error:
        raise HTTPException(status_code=503, detail=str(error)) from error


@app.post("/api/pretrain-records")
def create_pretrain_record(
    payload: PretrainRecordInput,
    _auth: str = Depends(require_auth),
) -> dict[str, Any]:
    try:
        recorded, already = pretrain_service.record(payload.model_dump())
    except ValueError as error:
        raise HTTPException(status_code=400, detail=str(error)) from error
    except Exception as error:
        raise HTTPException(status_code=503, detail=str(error)) from error
    return {"recorded": recorded, "already": already}


@app.delete("/api/pretrain-records/{file_id}")
def delete_pretrain_record(file_id: int, _auth: str = Depends(require_auth)) -> dict[str, Any]:
    try:
        deleted = pretrain_service.remove(file_id)
    except ValueError as error:
        raise HTTPException(status_code=400, detail=str(error)) from error
    except Exception as error:
        raise HTTPException(status_code=503, detail=str(error)) from error
    if not deleted:
        raise HTTPException(status_code=404, detail="CSV 中不存在该记录")
    return {"deleted": deleted}