import json
from datetime import datetime
from functools import wraps
from math import pi

import pytz

from app.algorithms.simulation.runner import run_simulation_ex
from app.native.wndb.core.projects import temporary_project_database
from app.services.tjnetwork import (
    get_pipe,
    get_tank,
)


def _isolated_operation(purpose: str):
    def decorator(func):
        @wraps(func)
        def wrapper(prj_name: str, *args, **kwargs):
            with temporary_project_database(prj_name, purpose) as temporary:
                kwargs["_temporary_project"] = temporary
                return func(prj_name, *args, **kwargs)

        return wrapper

    return decorator


############################################################
# project management 07  ***暂时不使用，与业务需求无关***
############################################################


@_isolated_operation("project_management")
def project_management(
    prj_name,
    start_datetime,
    pump_control,
    tank_initial_level_control=None,
    region_demand_control=None,
    _temporary_project=None,
) -> str:
    print(
        datetime.now(pytz.timezone("Asia/Shanghai")).strftime("%Y-%m-%d %H:%M:%S")
        + " -- Start Analysis."
    )
    if _temporary_project is None:
        raise RuntimeError("Project-management isolation was not prepared")
    new_name = _temporary_project
    print(
        datetime.now(pytz.timezone("Asia/Shanghai")).strftime("%Y-%m-%d %H:%M:%S")
        + " -- Start Copying Database."
    )
    print(
        datetime.now(pytz.timezone("Asia/Shanghai")).strftime("%Y-%m-%d %H:%M:%S")
        + " -- Start Opening Database."
    )
    print(
        datetime.now(pytz.timezone("Asia/Shanghai")).strftime("%Y-%m-%d %H:%M:%S")
        + " -- Database Loading OK."
    )
    result = run_simulation_ex(
        name=new_name,
        simulation_type="realtime",
        start_datetime=start_datetime,
        duration=86400,
        pump_control=pump_control,
        tank_initial_level_control=tank_initial_level_control,
        region_demand_control=region_demand_control,
        downloading_prohibition=True,
    )
    return result


############################################################
# scheduling analysis 08   ***暂时不使用，与业务需求无关***
############################################################


@_isolated_operation("scheduling")
def scheduling_simulation(
    prj_name,
    start_time,
    pump_control,
    tank_id,
    water_plant_output_id,
    time_delta=300,
    _temporary_project=None,
) -> str:
    print(
        datetime.now(pytz.timezone("Asia/Shanghai")).strftime("%Y-%m-%d %H:%M:%S")
        + " -- Start Analysis."
    )
    if _temporary_project is None:
        raise RuntimeError("Scheduling isolation was not prepared")
    new_name = _temporary_project
    print(
        datetime.now(pytz.timezone("Asia/Shanghai")).strftime("%Y-%m-%d %H:%M:%S")
        + " -- Start Copying Database."
    )
    print(
        datetime.now(pytz.timezone("Asia/Shanghai")).strftime("%Y-%m-%d %H:%M:%S")
        + " -- Start Opening Database."
    )
    print(
        datetime.now(pytz.timezone("Asia/Shanghai")).strftime("%Y-%m-%d %H:%M:%S")
        + " -- Database Loading OK."
    )

    simulation_result = json.loads(
        run_simulation_ex(
            new_name, "realtime", start_time, duration=0, pump_control=pump_control
        )
    )

    output_data = simulation_result.get("output")
    if not isinstance(output_data, dict):
        raise RuntimeError("run_simulation_ex did not return JSON output content")

    tank = get_tank(new_name, tank_id)  # 水塔信息
    tank_floor_space = pi * pow(tank["diameter"] / 2, 2)  # 水塔底面积(m^2)
    tank_init_level = tank["init_level"]  # 水塔初始水位(m)
    tank_pipes_id = tank["links"]  # pipes list

    tank_pipe_flow_direction = (
        {}
    )  # 管道流向修正系数, 水塔为下游节点时为1, 水塔为上游节点时为-1
    for pipe_id in tank_pipes_id:
        if get_pipe(new_name, pipe_id)["node2"] == tank_id:  # 水塔为下游节点
            tank_pipe_flow_direction[pipe_id] = 1
        else:
            tank_pipe_flow_direction[pipe_id] = -1

    node_results = output_data.get("node_results") or []  # [{'node': str, 'result': [{'pressure': float}]}]
    water_plant_output_pressure = 0
    for node_result in node_results:
        if node_result["node"] == water_plant_output_id:  # 水厂出水压力(m)
            water_plant_output_pressure = node_result["result"][-1]["pressure"]
    water_plant_output_pressure /= 100  # 预计水厂出水压力(Mpa)

    pipe_results = output_data.get("link_results") or []  # [{'link': str, 'result': [{'flow': float}]}]
    tank_inflow = 0
    for pipe_result in pipe_results:
        for pipe_id in tank_pipes_id:  # 遍历与水塔相连的管道
            if pipe_result["link"] == pipe_id:  # 水塔入流流量(L/s)
                tank_inflow += (
                    pipe_result["result"][-1]["flow"]
                    * tank_pipe_flow_direction[pipe_id]
                )
    tank_inflow /= 1000  # 水塔入流流量(m^3/s)
    tank_level_delta = tank_inflow * time_delta / tank_floor_space  # 水塔水位改变值(m)
    tank_level = tank_init_level + tank_level_delta  # 预计水塔水位(m)

    simulation_results = {
        "water_plant_output_pressure": water_plant_output_pressure,
        "tank_init_level": tank_init_level,
        "tank_level": tank_level,
    }

    return json.dumps(simulation_results)


@_isolated_operation("daily_scheduling")
def daily_scheduling_simulation(
    prj_name,
    start_time,
    pump_control,
    reservoir_id,
    tank_id,
    water_plant_output_id,
    _temporary_project=None,
) -> str:
    print(
        datetime.now(pytz.timezone("Asia/Shanghai")).strftime("%Y-%m-%d %H:%M:%S")
        + " -- Start Analysis."
    )
    if _temporary_project is None:
        raise RuntimeError("Daily-scheduling isolation was not prepared")
    new_name = _temporary_project
    print(
        datetime.now(pytz.timezone("Asia/Shanghai")).strftime("%Y-%m-%d %H:%M:%S")
        + " -- Start Copying Database."
    )
    print(
        datetime.now(pytz.timezone("Asia/Shanghai")).strftime("%Y-%m-%d %H:%M:%S")
        + " -- Start Opening Database."
    )
    print(
        datetime.now(pytz.timezone("Asia/Shanghai")).strftime("%Y-%m-%d %H:%M:%S")
        + " -- Database Loading OK."
    )

    simulation_result = json.loads(
        run_simulation_ex(
            new_name,
            "realtime",
            start_time,
            duration=86400,
            pump_control=pump_control,
        )
    )

    output_data = simulation_result.get("output")
    if not isinstance(output_data, dict):
        raise RuntimeError("run_simulation_ex did not return JSON output content")

    node_results = output_data.get("node_results") or []  # [{'node': str, 'result': [{'pressure': float, 'head': float}]}]
    water_plant_output_pressure = []
    reservoir_level = []
    tank_level = []
    for node_result in node_results:
        if node_result["node"] == water_plant_output_id:
            for result in node_result["result"]:
                water_plant_output_pressure.append(
                    result["pressure"] / 100
                )  # 水厂出水压力(Mpa)
        elif node_result["node"] == reservoir_id:
            for result in node_result["result"]:
                reservoir_level.append(result["head"] - 250.35)  # 清水池液位(m)
        elif node_result["node"] == tank_id:
            for result in node_result["result"]:
                tank_level.append(result["pressure"])  # 调节池液位(m)

    simulation_results = {
        "water_plant_output_pressure": water_plant_output_pressure,
        "reservoir_level": reservoir_level,
        "tank_level": tank_level,
    }

    return json.dumps(simulation_results)