drivers_bot/app/api/entries.py

from io import BytesIO
from datetime import date

import matplotlib.pyplot as plt
from fastapi import APIRouter, Depends, HTTPException, Response, status
from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession

from app.db.session import get_session
from app.models.car import Car
from app.models.expense import FuelEntry, ServiceEntry
from app.schemas.expense import (
    FuelEntryCreate,
    FuelEntryRead,
    OdometerPrediction,
    OwnershipStats,
    ServiceEntryCreate,
    ServiceEntryRead,
)
from app.services.calculations import dataframe_from_query, get_ownership_stats, predict_odometer

router = APIRouter(tags=["entries"])


async def ensure_car(session: AsyncSession, car_id: int) -> None:
    if await session.get(Car, car_id) is None:
        raise HTTPException(status_code=404, detail="Car not found")


@router.post("/fuel", response_model=FuelEntryRead, status_code=status.HTTP_201_CREATED)
async def create_fuel_entry(
    payload: FuelEntryCreate, session: AsyncSession = Depends(get_session)
) -> FuelEntry:
    await ensure_car(session, payload.car_id)
    entry = FuelEntry(**payload.model_dump())
    session.add(entry)
    car = await session.get(Car, payload.car_id)
    if car and (car.current_odometer is None or payload.odometer > car.current_odometer):
        car.current_odometer = payload.odometer
    await session.commit()
    await session.refresh(entry)
    return entry


@router.get("/cars/{car_id}/fuel", response_model=list[FuelEntryRead])
async def list_fuel_entries(
    car_id: int,
    date_from: date | None = None,
    date_to: date | None = None,
    session: AsyncSession = Depends(get_session),
) -> list[FuelEntry]:
    stmt = select(FuelEntry).where(FuelEntry.car_id == car_id)
    if date_from:
        stmt = stmt.where(FuelEntry.entry_date >= date_from)
    if date_to:
        stmt = stmt.where(FuelEntry.entry_date <= date_to)
    result = await session.execute(
        stmt.order_by(FuelEntry.entry_date.desc())
    )
    return list(result.scalars())


@router.post("/service", response_model=ServiceEntryRead, status_code=status.HTTP_201_CREATED)
async def create_service_entry(
    payload: ServiceEntryCreate, session: AsyncSession = Depends(get_session)
) -> ServiceEntry:
    await ensure_car(session, payload.car_id)
    entry = ServiceEntry(**payload.model_dump())
    session.add(entry)
    car = await session.get(Car, payload.car_id)
    if car and payload.odometer and (
        car.current_odometer is None or payload.odometer > car.current_odometer
    ):
        car.current_odometer = payload.odometer
    await session.commit()
    await session.refresh(entry)
    return entry


@router.get("/cars/{car_id}/service", response_model=list[ServiceEntryRead])
async def list_service_entries(
    car_id: int,
    date_from: date | None = None,
    date_to: date | None = None,
    session: AsyncSession = Depends(get_session),
) -> list[ServiceEntry]:
    stmt = select(ServiceEntry).where(ServiceEntry.car_id == car_id)
    if date_from:
        stmt = stmt.where(ServiceEntry.entry_date >= date_from)
    if date_to:
        stmt = stmt.where(ServiceEntry.entry_date <= date_to)
    result = await session.execute(
        stmt.order_by(ServiceEntry.entry_date.desc())
    )
    return list(result.scalars())


@router.get("/cars/{car_id}/stats", response_model=OwnershipStats)
async def car_stats(
    car_id: int,
    date_from: date | None = None,
    date_to: date | None = None,
    session: AsyncSession = Depends(get_session),
) -> OwnershipStats:
    await ensure_car(session, car_id)
    today = date.today()
    period_from = date_from or today.replace(day=1)
    period_to = date_to or today
    return await get_ownership_stats(session, car_id, period_from, period_to)


@router.get("/cars/{car_id}/analytics", response_model=OdometerPrediction)
async def car_analytics(car_id: int, session: AsyncSession = Depends(get_session)) -> OdometerPrediction:
    await ensure_car(session, car_id)
    return await predict_odometer(session, car_id)


@router.get("/cars/{car_id}/charts/expenses.png")
async def expenses_chart(car_id: int, session: AsyncSession = Depends(get_session)) -> Response:
    await ensure_car(session, car_id)
    fuel_df = await dataframe_from_query(
        session,
        select(FuelEntry.entry_date.label("date"), FuelEntry.total_cost.label("cost")).where(
            FuelEntry.car_id == car_id
        ),
    )
    service_df = await dataframe_from_query(
        session,
        select(ServiceEntry.entry_date.label("date"), ServiceEntry.total_cost.label("cost")).where(
            ServiceEntry.car_id == car_id
        ),
    )
    if fuel_df.empty and service_df.empty:
        raise HTTPException(status_code=404, detail="No data for chart")

    frames = []
    if not fuel_df.empty:
        fuel_df["type"] = "fuel"
        frames.append(fuel_df)
    if not service_df.empty:
        service_df["type"] = "service"
        frames.append(service_df)

    import pandas as pd

    df = pd.concat(frames)
    df["date"] = pd.to_datetime(df["date"])
    pivot = df.pivot_table(index="date", columns="type", values="cost", aggfunc="sum").sort_index()

    fig, ax = plt.subplots(figsize=(8, 4.5))
    pivot.plot(kind="bar", stacked=True, ax=ax)
    ax.set_title("Car expenses")
    ax.set_xlabel("Date")
    ax.set_ylabel("Cost")
    fig.tight_layout()

    buffer = BytesIO()
    fig.savefig(buffer, format="png")
    plt.close(fig)
    return Response(buffer.getvalue(), media_type="image/png")