TJWaterServerBinary/app/api/v1/endpoints/timeseries/scheme.py

from fastapi import APIRouter, Depends, HTTPException
from typing import List
from datetime import datetime
from psycopg import AsyncConnection

from app.infra.db.timescaledb.repositories.scheme import SchemeRepository
from .dependencies import get_timescale_connection

router = APIRouter()


@router.post("/scheme/links/batch", status_code=201)
async def insert_scheme_links(
    data: List[dict], conn: AsyncConnection = Depends(get_timescale_connection)
):
    await SchemeRepository.insert_links_batch(conn, data)
    return {"message": f"Inserted {len(data)} records"}


@router.get("/scheme/links")
async def get_scheme_links(
    scheme_type: str,
    scheme_name: str,
    start_time: datetime,
    end_time: datetime,
    conn: AsyncConnection = Depends(get_timescale_connection),
):
    return await SchemeRepository.get_links_by_scheme_and_time_range(
        conn, scheme_type, scheme_name, start_time, end_time
    )


@router.get("/scheme/links/{link_id}/field")
async def get_scheme_link_field(
    scheme_type: str,
    scheme_name: str,
    link_id: str,
    start_time: datetime,
    end_time: datetime,
    field: str,
    conn: AsyncConnection = Depends(get_timescale_connection),
):
    try:
        return await SchemeRepository.get_link_field_by_scheme_and_time_range(
            conn, scheme_type, scheme_name, start_time, end_time, link_id, field
        )
    except ValueError as e:
        raise HTTPException(status_code=400, detail=str(e))


@router.patch("/scheme/links/{link_id}/field")
async def update_scheme_link_field(
    scheme_type: str,
    scheme_name: str,
    link_id: str,
    time: datetime,
    field: str,
    value: float,
    conn: AsyncConnection = Depends(get_timescale_connection),
):
    try:
        await SchemeRepository.update_link_field(
            conn, time, scheme_type, scheme_name, link_id, field, value
        )
        return {"message": "Updated successfully"}
    except ValueError as e:
        raise HTTPException(status_code=400, detail=str(e))


@router.delete("/scheme/links")
async def delete_scheme_links(
    scheme_type: str,
    scheme_name: str,
    start_time: datetime,
    end_time: datetime,
    conn: AsyncConnection = Depends(get_timescale_connection),
):
    await SchemeRepository.delete_links_by_scheme_and_time_range(
        conn, scheme_type, scheme_name, start_time, end_time
    )
    return {"message": "Deleted successfully"}


@router.post("/scheme/nodes/batch", status_code=201)
async def insert_scheme_nodes(
    data: List[dict], conn: AsyncConnection = Depends(get_timescale_connection)
):
    await SchemeRepository.insert_nodes_batch(conn, data)
    return {"message": f"Inserted {len(data)} records"}


@router.get("/scheme/nodes/{node_id}/field")
async def get_scheme_node_field(
    scheme_type: str,
    scheme_name: str,
    node_id: str,
    start_time: datetime,
    end_time: datetime,
    field: str,
    conn: AsyncConnection = Depends(get_timescale_connection),
):
    try:
        return await SchemeRepository.get_node_field_by_scheme_and_time_range(
            conn, scheme_type, scheme_name, start_time, end_time, node_id, field
        )
    except ValueError as e:
        raise HTTPException(status_code=400, detail=str(e))


@router.patch("/scheme/nodes/{node_id}/field")
async def update_scheme_node_field(
    scheme_type: str,
    scheme_name: str,
    node_id: str,
    time: datetime,
    field: str,
    value: float,
    conn: AsyncConnection = Depends(get_timescale_connection),
):
    try:
        await SchemeRepository.update_node_field(
            conn, time, scheme_type, scheme_name, node_id, field, value
        )
        return {"message": "Updated successfully"}
    except ValueError as e:
        raise HTTPException(status_code=400, detail=str(e))


@router.delete("/scheme/nodes")
async def delete_scheme_nodes(
    scheme_type: str,
    scheme_name: str,
    start_time: datetime,
    end_time: datetime,
    conn: AsyncConnection = Depends(get_timescale_connection),
):
    await SchemeRepository.delete_nodes_by_scheme_and_time_range(
        conn, scheme_type, scheme_name, start_time, end_time
    )
    return {"message": "Deleted successfully"}


@router.post("/scheme/simulation/store", status_code=201)
async def store_scheme_simulation_result(
    scheme_type: str,
    scheme_name: str,
    node_result_list: List[dict],
    link_result_list: List[dict],
    result_start_time: str,
    conn: AsyncConnection = Depends(get_timescale_connection),
):
    """Store scheme simulation results to TimescaleDB"""
    await SchemeRepository.store_scheme_simulation_result(
        conn,
        scheme_type,
        scheme_name,
        node_result_list,
        link_result_list,
        result_start_time,
    )
    return {"message": "Scheme simulation results stored successfully"}


@router.get("/scheme/query/by-scheme-time-property")
async def query_scheme_records_by_scheme_time_property(
    scheme_type: str,
    scheme_name: str,
    query_time: str,
    type: str,
    property: str,
    conn: AsyncConnection = Depends(get_timescale_connection),
):
    """Query all scheme records by scheme, time and property"""
    try:
        results = await SchemeRepository.query_all_record_by_scheme_time_property(
            conn, scheme_type, scheme_name, query_time, type, property
        )
        return {"results": results}
    except ValueError as e:
        raise HTTPException(status_code=400, detail=str(e))


@router.get("/scheme/query/by-id-time")
async def query_scheme_simulation_by_id_time(
    scheme_type: str,
    scheme_name: str,
    id: str,
    type: str,
    query_time: str,
    conn: AsyncConnection = Depends(get_timescale_connection),
):
    """Query scheme simulation results by id and time"""
    try:
        result = await SchemeRepository.query_scheme_simulation_result_by_id_time(
            conn, scheme_type, scheme_name, id, type, query_time
        )
        return {"result": result}
    except ValueError as e:
        raise HTTPException(status_code=400, detail=str(e))