full-stack-fastapi/backend/app/api/routes/mqtt_raw.py

from __future__ import annotations

from datetime import date, datetime, time, timedelta, timezone
from pathlib import Path
from typing import Any

from fastapi import APIRouter, HTTPException, Query
from fastapi.responses import Response
from sqlmodel import col, delete, func, select

from app.api.deps import CurrentUser, SessionDep
from app.models import (
    DeleteResult,
    MqttRawDaysPublic,
    MqttRawDaySummary,
    TerminalMessageRaw,
    TerminalRawArchiveManifest,
)

router = APIRouter(prefix="/mqtt-raw", tags=["mqtt_raw"])
UTC_PLUS_8 = timezone(timedelta(hours=8))


def _normalize_day(value: Any) -> date:
    if isinstance(value, date) and not isinstance(value, datetime):
        return value
    if isinstance(value, datetime):
        return value.date()
    if isinstance(value, str):
        return date.fromisoformat(value)
    raise ValueError(f"Unsupported day value: {value!r}")


def _day_bounds(day: date) -> tuple[datetime, datetime]:
    start_local = datetime.combine(day, time.min, tzinfo=UTC_PLUS_8)
    end_local = start_local + timedelta(days=1)
    return start_local.astimezone(timezone.utc), end_local.astimezone(timezone.utc)


def _local_day_expr(session: SessionDep, field: Any) -> Any:
    bind = session.get_bind()
    dialect_name = bind.dialect.name if bind is not None else ""

    if dialect_name == "sqlite":
        return func.date(func.datetime(field, "+8 hours"))
    if dialect_name in {"postgresql", "postgres"}:
        return func.date(func.timezone("Asia/Shanghai", field))
    return func.date(field)


def _build_window(
    *, day: date, start_seconds: int | None, end_seconds: int | None
) -> tuple[datetime, datetime, int, int]:
    start_of_day_local = datetime.combine(day, time.min, tzinfo=UTC_PLUS_8)
    start_sec = 0 if start_seconds is None else start_seconds
    end_sec = 24 * 3600 if end_seconds is None else end_seconds

    if not 0 <= start_sec <= 24 * 3600:
        raise HTTPException(status_code=400, detail="start_seconds must be 0..86400")
    if not 0 <= end_sec <= 24 * 3600:
        raise HTTPException(status_code=400, detail="end_seconds must be 0..86400")
    if start_sec >= end_sec:
        raise HTTPException(
            status_code=400, detail="start_seconds must be smaller than end_seconds"
        )

    start = (start_of_day_local + timedelta(seconds=start_sec)).astimezone(timezone.utc)
    end = (start_of_day_local + timedelta(seconds=end_sec)).astimezone(timezone.utc)
    return start, end, start_sec, end_sec


def _unlink_file(path_str: str) -> None:
    path = Path(path_str)
    try:
        path.unlink(missing_ok=True)
    except OSError:
        return


@router.get(
    "/days",
    response_model=MqttRawDaysPublic,
    operation_id="mqtt_raw_list_days",
)
def list_mqtt_raw_days(
    session: SessionDep,
    _current_user: CurrentUser,
    skip: int = Query(default=0, ge=0),
    limit: int = Query(default=30, ge=1, le=200),
) -> MqttRawDaysPublic:
    days: dict[date, MqttRawDaySummary] = {}

    raw_day_expr = _local_day_expr(session, TerminalMessageRaw.received_at)
    raw_statement = (
        select(
            raw_day_expr,
            func.count(),
            func.min(TerminalMessageRaw.received_at),
            func.max(TerminalMessageRaw.received_at),
        )
        .group_by(raw_day_expr)
        .order_by(raw_day_expr.desc())
    )
    for day_value, message_count, first_at, last_at in session.exec(
        raw_statement
    ).all():
        day_key = _normalize_day(day_value)
        days[day_key] = MqttRawDaySummary(
            day=day_key,
            message_count=int(message_count or 0),
            first_received_at=first_at,
            last_received_at=last_at,
        )

    archive_day_expr = _local_day_expr(session, TerminalRawArchiveManifest.range_start)
    archive_statement = (
        select(
            archive_day_expr,
            func.sum(TerminalRawArchiveManifest.record_count),
            func.min(TerminalRawArchiveManifest.range_start),
            func.max(TerminalRawArchiveManifest.range_end),
        )
        .group_by(archive_day_expr)
        .order_by(archive_day_expr.desc())
    )
    for day_value, message_count, first_at, last_at in session.exec(
        archive_statement
    ).all():
        day_key = _normalize_day(day_value)
        existing = days.get(day_key)
        archive_count = int(message_count or 0)
        if existing is None:
            days[day_key] = MqttRawDaySummary(
                day=day_key,
                message_count=archive_count,
                first_received_at=first_at,
                last_received_at=last_at,
            )
            continue
        existing.message_count += archive_count
        if first_at and (
            existing.first_received_at is None or first_at < existing.first_received_at
        ):
            existing.first_received_at = first_at
        if last_at and (
            existing.last_received_at is None or last_at > existing.last_received_at
        ):
            existing.last_received_at = last_at

    ordered = sorted(days.values(), key=lambda item: item.day, reverse=True)
    paged = ordered[skip : skip + limit]
    return MqttRawDaysPublic(data=paged, count=len(ordered))


@router.get(
    "/download",
    operation_id="mqtt_raw_download",
    responses={
        200: {
            "content": {
                "application/octet-stream": {
                    "schema": {"type": "string", "format": "binary"}
                }
            }
        }
    },
)
def download_mqtt_raw(
    session: SessionDep,
    _current_user: CurrentUser,
    day: date,
    start_seconds: int | None = Query(default=None),
    end_seconds: int | None = Query(default=None),
) -> Response:
    range_start, range_end, start_sec, end_sec = _build_window(
        day=day,
        start_seconds=start_seconds,
        end_seconds=end_seconds,
    )

    chunks: list[tuple[datetime, bytes]] = []

    raw_statement = (
        select(TerminalMessageRaw)
        .where(TerminalMessageRaw.received_at >= range_start)
        .where(TerminalMessageRaw.received_at < range_end)
        .order_by(col(TerminalMessageRaw.received_at), col(TerminalMessageRaw.id))
    )
    raw_rows = session.exec(raw_statement).all()
    for row in raw_rows:
        chunks.append((row.received_at, row.payload_bytes))

    manifest_statement = (
        select(TerminalRawArchiveManifest)
        .where(TerminalRawArchiveManifest.range_start >= range_start)
        .where(TerminalRawArchiveManifest.range_end <= range_end)
        .order_by(
            col(TerminalRawArchiveManifest.range_start),
            col(TerminalRawArchiveManifest.range_end),
        )
    )
    manifests = session.exec(manifest_statement).all()
    for manifest in manifests:
        path = Path(manifest.file_path)
        if not path.exists() or not path.is_file():
            continue
        chunks.append((manifest.range_start, path.read_bytes()))

    chunks.sort(key=lambda item: item[0])
    payload = b"".join(chunk for _, chunk in chunks)

    name_suffix = (
        day.isoformat()
        if start_sec == 0 and end_sec == 24 * 3600
        else f"{day.isoformat()}_{start_sec}_{end_sec}"
    )
    filename = f"mqtt_raw_{name_suffix}.txt"
    return Response(
        content=payload,
        media_type="application/octet-stream",
        headers={"Content-Disposition": f'attachment; filename="{filename}"'},
    )


@router.delete(
    "/day/{day}",
    response_model=DeleteResult,
    operation_id="mqtt_raw_delete_day",
)
def delete_mqtt_raw_by_day(
    session: SessionDep,
    _current_user: CurrentUser,
    day: date,
) -> DeleteResult:
    range_start, range_end = _day_bounds(day)

    raw_statement = (
        select(TerminalMessageRaw.id)
        .where(TerminalMessageRaw.received_at >= range_start)
        .where(TerminalMessageRaw.received_at < range_end)
    )
    raw_ids = list(session.exec(raw_statement).all())
    if raw_ids:
        session.exec(
            delete(TerminalMessageRaw).where(col(TerminalMessageRaw.id).in_(raw_ids))
        )

    manifest_statement = (
        select(TerminalRawArchiveManifest)
        .where(TerminalRawArchiveManifest.range_start >= range_start)
        .where(TerminalRawArchiveManifest.range_end <= range_end)
    )
    manifests = session.exec(manifest_statement).all()
    for manifest in manifests:
        _unlink_file(manifest.file_path)
        session.delete(manifest)

    session.commit()

    deleted_count = len(raw_ids) + len(manifests)
    return DeleteResult(
        deleted_count=deleted_count,
        message=f"Deleted {deleted_count} record(s) for {day.isoformat()}",
    )


@router.delete(
    "/all",
    response_model=DeleteResult,
    operation_id="mqtt_raw_delete_all",
)
def delete_all_mqtt_raw(
    session: SessionDep,
    _current_user: CurrentUser,
    confirm: bool = Query(default=False),
) -> DeleteResult:
    if not confirm:
        raise HTTPException(
            status_code=400,
            detail="Set confirm=true to delete all saved MQTT raw data",
        )

    raw_ids = list(session.exec(select(TerminalMessageRaw.id)).all())
    if raw_ids:
        session.exec(
            delete(TerminalMessageRaw).where(col(TerminalMessageRaw.id).in_(raw_ids))
        )

    manifests = session.exec(select(TerminalRawArchiveManifest)).all()
    for manifest in manifests:
        _unlink_file(manifest.file_path)
        session.delete(manifest)

    session.commit()

    deleted_count = len(raw_ids) + len(manifests)
    return DeleteResult(
        deleted_count=deleted_count,
        message=f"Deleted {deleted_count} record(s) from MQTT raw storage",
    )