from fastapi import APIRouter, Request, Query, Path, Body
from typing import Any, List, Dict, Union
from app.services.tjnetwork import (
    Any,
    ChangeSet,
    calculate_demand_to_network,
    calculate_demand_to_nodes,
    calculate_demand_to_region,
    get_demand,
    get_demand_schema,
    set_demand,
)

router = APIRouter()

############################################################
# demand 9.[DEMANDS]
############################################################

@router.get(
    "/getdemandschema",
    summary="获取需水量属性架构",
    description="获取指定水网中需水量(Demand)的属性架构定义"
)
async def fastapi_get_demand_schema(network: str = Query(..., description="管网名称（或数据库名称）")) -> dict[str, dict[str, Any]]:
    """
    获取需水量属性架构。
    
    返回指定水网的需水量属性架构，包括所有可配置的属性及其类型定义。
    """
    return get_demand_schema(network)


@router.get(
    "/getdemandproperties/",
    summary="获取需水量属性",
    description="获取指定水网中节点的需水量属性信息"
)
async def fastapi_get_demand_properties(
    network: str = Query(..., description="管网名称（或数据库名称）"),
    junction: str = Query(..., description="节点ID")
) -> dict[str, Any]:
    """
    获取节点的需水量属性。
    
    返回指定节点的所有需水量信息，包括需水量值、水压等级等。
    """
    return get_demand(network, junction)


# example: set_demand(p, ChangeSet({'junction': 'j1', 'demands': [{'demand': 10.0, 'pattern': None, 'category': 'x'}, {'demand': 20.0, 'pattern': None, 'category': None}]}))
@router.post(
    "/setdemandproperties/",
    response_model=None,
    summary="设置需水量属性",
    description="设置指定水网中节点的需水量属性信息"
)
async def fastapi_set_demand_properties(
    network: str = Query(..., description="管网名称（或数据库名称）"),
    junction: str = Query(..., description="节点ID"),
    req: Request = None
) -> ChangeSet:
    """
    设置节点的需水量属性。
    
    修改指定节点的需水量信息。请求体应包含需水量值、水压等级等属性。
    """
    props = await req.json()
    ps = {"junction": junction} | props
    return set_demand(network, ChangeSet(ps))

############################################################
# water distribution 36.[Water Distribution]
############################################################
@router.get(
    "/calculatedemandtonodes/",
    summary="计算需水量到节点分配",
    description="将总需水量按指定方式分配到多个节点"
)
async def fastapi_calculate_demand_to_nodes(
    network: str = Query(..., description="管网名称（或数据库名称）"),
    req: Request = None
) -> dict[str, float]:
    """
    计算需水量到节点分配。
    
    将指定的总需水量均匀或按比例分配到指定的节点列表中。
    
    请求体格式:
    {
        "demand": 需水量值(float),
        "nodes": 节点ID列表(list[str])
    }
    """
    props = await req.json()
    demand = props["demand"]
    nodes = props["nodes"]
    return calculate_demand_to_nodes(network, demand, nodes)

@router.get(
    "/calculatedemandtoregion/",
    summary="计算需水量到区域分配",
    description="将总需水量按区域特征分配到该区域内的节点"
)
async def fastapi_calculate_demand_to_region(
    network: str = Query(..., description="管网名称（或数据库名称）"),
    req: Request = None
) -> dict[str, float]:
    """
    计算需水量到区域分配。
    
    根据区域内节点的特征(如面积、人口等)将总需水量分配到该区域的各个节点。
    
    请求体格式:
    {
        "demand": 需水量值(float),
        "region": 区域ID(str)
    }
    """
    props = await req.json()
    demand = props["demand"]
    region = props["region"]
    return calculate_demand_to_region(network, demand, region)

@router.get(
    "/calculatedemandtonetwork/",
    summary="计算需水量到整网分配",
    description="将需水量均匀分配到整个水网的所有需水节点"
)
async def fastapi_calculate_demand_to_network(
    network: str = Query(..., description="管网名称（或数据库名称）"),
    demand: float = Query(..., description="总需水量(m³/h)", gt=0)
) -> dict[str, float]:
    """
    计算需水量到整网分配。
    
    将指定的需水量均匀分配到整个水网的所有需水节点。
    """
    return calculate_demand_to_network(network, demand)