from random import Random

import sampo.generator.config.gen_counts as gen_c
import sampo.generator.config.worker_req as wr
from sampo.generator.config.worker_req import scale_reqs
from sampo.schemas.graph import GraphNode, EdgeType
from sampo.schemas.interval import IntervalUniform
from sampo.schemas.utils import uuid_str
from sampo.schemas.works import WorkUnit


def _add_addition_work(probability: float, rand: Random | None = None) -> bool:
    """Return answer, if addition work will be added"""
    return IntervalUniform(0, 1).rand_float(rand) <= probability


def _get_roads(parents: list[GraphNode], cluster_name: str, dist: float,
               rand: Random | None = None) -> tuple[dict[str, GraphNode], int]:
    road_nodes = dict()
    min_r = WorkUnit(uuid_str(rand), {'granular_name': 'minimal road',
                                      'measurement': 'km'},
                     scale_reqs(wr.MIN_ROAD, dist), group=f'{cluster_name}:road', volume=dist)
    road_nodes['min'] = GraphNode(min_r, parents)
    temp_r = WorkUnit(uuid_str(rand), {'granular_name': 'temporary road',
                                       'measurement': 'km'},
                      scale_reqs(wr.TEMP_ROAD, dist), group=f'{cluster_name}:road', volume=dist)
    road_nodes['temp'] = GraphNode(temp_r, [(road_nodes['min'], wr.ATOMIC_ROAD_LEN, EdgeType.LagFinishStart)])

    final_r = WorkUnit(uuid_str(rand), {'granular_name': 'final road',
                                        'measurement': 'km'},
                       scale_reqs(wr.FINAL_ROAD, dist), group=f'{cluster_name}:road', volume=dist)
    road_nodes['final'] = GraphNode(final_r, [(road_nodes['temp'], wr.ATOMIC_ROAD_LEN, EdgeType.LagFinishStart)])
    return road_nodes, len(road_nodes)


def _get_engineering_preparation(parents: list[GraphNode], cluster_name: str, boreholes_count: int,
                                 rand: Random | None = None) -> GraphNode:
    worker_req = wr.mul_borehole_volume(wr.ENGINEERING_PREPARATION, boreholes_count, wr.ENGINEERING_PREPARATION_BASE)
    work = WorkUnit(uuid_str(rand), 'engineering preparation', worker_req,
                    group=f'{cluster_name}:engineering',
                    volume=wr.get_borehole_volume(boreholes_count, wr.ENGINEERING_PREPARATION_BASE))
    node = GraphNode(work, parents)
    return node


def _get_power_lines(parents: list[GraphNode], cluster_name: str, dist_line: float,
                     dist_high_line: float | None = None, rand: Random | None = None) -> tuple[list[GraphNode], int]:
    worker_req = wr.scale_reqs(wr.POWER_LINE, dist_line)
    power_line_1 = WorkUnit(uuid_str(rand),
                            {'granular_name': 'power line', 'measurement': 'km'},
                            worker_req,
                            group=f'{cluster_name}:electricity',
                            volume=dist_line)
    power_line_2 = WorkUnit(uuid_str(rand),
                            {'granular_name': 'power line', 'measurement': 'km'},
                            worker_req,
                            group=f'{cluster_name}:electricity',
                            volume=dist_line)

    power_lines = [
        GraphNode(power_line_1, parents),
        GraphNode(power_line_2, parents),
    ]
    if dist_high_line is not None:
        worker_req_high = wr.scale_reqs(wr.POWER_LINE, dist_high_line)
        high_power_line = WorkUnit(uuid_str(rand),
                                   {'granular_name': 'high power line', 'measurement': 'km'},
                                   worker_req_high,
                                   group=f'{cluster_name}:electricity', volume=dist_high_line)
        power_lines.append(GraphNode(high_power_line, parents))

    return power_lines, len(power_lines)


def _get_pipe_lines(parents: list[GraphNode], cluster_name: str, pipe_dists: list[float],
                    rand: Random | None = None) -> tuple[list[GraphNode], int]:
    worker_req_pipe = wr.scale_reqs(wr.PIPE_LINE, pipe_dists[0])
    first_pipe = WorkUnit(uuid_str(rand),
                          {'granular_name': 'pipe', 'measurement': 'km'},
                          worker_req_pipe, group=f'{cluster_name}:oil_gas_long_pipes',
                          volume=pipe_dists[0])

    graph_nodes = [GraphNode(first_pipe, parents)]
    for i in range(1, len(pipe_dists)):
        node_work = WorkUnit(uuid_str(rand),
                             {'granular_name': 'node'},
                             wr.PIPE_NODE,
                             group=f'{cluster_name}:oil_gas_long_pipes')
        graph_nodes.append(GraphNode(node_work, parents))
        worker_req_pipe = wr.scale_reqs(wr.PIPE_LINE, pipe_dists[i])
        pipe_work = WorkUnit(uuid_str(rand),
                             {'granular_name': 'pipe', 'measurement': 'km'},
                             worker_req_pipe,
                             group=f'{cluster_name}:oil_gas_long_pipes',
                             volume=pipe_dists[i])
        graph_nodes.append(GraphNode(pipe_work, parents))

    worker_req_loop = wr.scale_reqs(wr.PIPE_LINE, pipe_dists[0])
    looping = WorkUnit(uuid_str(rand),
                       {'granular_name': 'looping', 'measurement': 'km'},
                       worker_req_loop,
                       group=f'{cluster_name}:oil_gas_long_pipes',
                       volume=pipe_dists[0])
    graph_nodes.append(GraphNode(looping, graph_nodes[0:1]))
    return graph_nodes, len(graph_nodes)


def _get_boreholes_equipment_group(parents: list[GraphNode], cluster_name: str, group_ind: int, borehole_count: int,
                                   rand: Random | None = None) -> tuple[list[GraphNode], int]:
    metering_install = WorkUnit(uuid_str(rand), 'metering installation',
                                wr.METERING_INSTALL, group=f'{cluster_name}:borehole_env_{group_ind}')
    worker_req_ktp_nep = wr.mul_borehole_volume(wr.KTP_NEP, borehole_count, wr.KTP_NEP_BASE)
    ktp_nep = WorkUnit(uuid_str(rand), 'KTP and NEP',
                       worker_req_ktp_nep, group=f'{cluster_name}:borehole_env_{group_ind}',
                       volume=wr.get_borehole_volume(borehole_count, wr.KTP_NEP_BASE))
    worker_req_tank = wr.mul_borehole_volume(wr.DRAINAGE_TANK, borehole_count, wr.DRAINAGE_TANK_BASE)
    drainage_tank = WorkUnit(uuid_str(rand), 'drainage tank',
                             worker_req_tank, group=f'{cluster_name}:borehole_env_{group_ind}',
                             volume=wr.get_borehole_volume(borehole_count, wr.DRAINAGE_TANK_BASE))
    nodes = [
        GraphNode(metering_install, parents),
        GraphNode(ktp_nep, parents),
        GraphNode(drainage_tank, parents),
    ]
    return nodes, len(nodes)


def _get_boreholes_equipment_shared(parents: list[GraphNode], cluster_name: str,
                                    rand: Random | None = None) -> tuple[list[GraphNode], int]:
    water_block = WorkUnit(uuid_str(rand), 'block water distribution', wr.WATER_BLOCK,
                           group=f'{cluster_name}:borehole_env')
    automation_block = WorkUnit(uuid_str(rand), 'block local automation', wr.AUTOMATION_BLOCK,
                                group=f'{cluster_name}:borehole_env')
    block_dosage = WorkUnit(uuid_str(rand), 'block dosage inhibitor', wr.BLOCK_DOSAGE,
                            group=f'{cluster_name}:borehole_env')
    start_filters = WorkUnit(uuid_str(rand), 'start filters system', wr.START_FILTER,
                             group=f'{cluster_name}:borehole_env')
    firewall = WorkUnit(uuid_str(rand), 'firewall tank', wr.FIREWALL,
                        group=f'{cluster_name}:borehole_env')
    nodes = [
        GraphNode(water_block, parents),
        GraphNode(automation_block, parents),
        GraphNode(block_dosage, parents),
        GraphNode(start_filters, parents),
        GraphNode(firewall, parents),
    ]
    return nodes, len(nodes)


def _get_boreholes(parents: list[GraphNode], cluster_name: str, group_ind: int, borehole_count: int,
                   rand: Random | None = None) -> tuple[list[GraphNode], int]:
    nodes = []
    for i in range(borehole_count):
        borehole_work = WorkUnit(uuid_str(rand), 'borehole',
                                 wr.BOREHOLE, group=f'{cluster_name}:borehole_groups_{group_ind}')
        nodes.append(GraphNode(borehole_work, parents))
    return nodes, len(nodes)


def _get_boreholes_equipment_general(parents: list[GraphNode], cluster_name: str, pipes_count: int, masts_count: int,
                                     rand: Random | None = None) -> tuple[list[GraphNode], int]:
    nodes = []
    dists_sum = 0
    for i in range(pipes_count):
        dist = gen_c.DIST_BETWEEN_BOREHOLES.rand_float(rand)
        dists_sum += dist
        worker_req_pipe = scale_reqs(wr.POWER_NETWORK, dist)
        pipe_net_work = WorkUnit(uuid_str(rand),
                                 {'granular_name': 'elem of pipe_network', 'measurement': 'km'},
                                 worker_req_pipe,
                                 group=f'{cluster_name}:oil_gas_pipe_net',
                                 volume=dist)
        nodes.append(GraphNode(pipe_net_work, parents))

    worker_req_power = scale_reqs(wr.POWER_NETWORK, dists_sum)
    power_net_work = WorkUnit(uuid_str(rand), {'granular_name': 'power network', 'measurement': 'km'},
                              worker_req_power,
                              group=f'{cluster_name}:electricity',
                              volume=dists_sum)
    nodes.append(GraphNode(power_net_work, parents))

    for i in range(masts_count):
        light_mast_work = WorkUnit(uuid_str(rand), {'granular_name': 'mast', 'measurement': 'km'},
                                   wr.LIGHT_MAST,
                                   group=f'{cluster_name}:light_masts')
        nodes.append(GraphNode(light_mast_work, parents))
    return nodes, len(nodes)


def _get_handing_stage(parents: list[GraphNode], cluster_name: str, borehole_count: int, group_ind: int,
                       rand: Random | None = None) -> GraphNode:
    worker_req = wr.mul_borehole_volume(wr.HANDING_STAGE, borehole_count, wr.HANDING_STAGE_BASE)
    work = WorkUnit(uuid_str(rand), 'cluster handing', worker_req,
                    group=f'{cluster_name}:handing_stage_{group_ind}',
                    volume=wr.get_borehole_volume(borehole_count, wr.HANDING_STAGE_BASE))
    node = GraphNode(work, parents)
    return node


def get_cluster_works(cluster_name: str, pipe_nodes_count: int,
                      pipe_net_count: int, light_masts_count: int, borehole_counts: list[int],
                      roads: dict[str, GraphNode] | None = None,
                      rand: Random | None = None) -> tuple[GraphNode, dict[str, GraphNode], int]:
    """
    Creates works on the development of the field on one object, i.e. a group of boreholes

    :param root_node: The parent node of the work graph for this object
    :param cluster_name: object name
    :param pipe_nodes_count: Number of pipeline segments from this field to other fields
    :param pipe_net_count: Number of pipes connecting boreholes
    :param light_masts_count: Number of floodlight masts
    :param borehole_counts: Number of boreholes
    :param roads: If the object is not connected to the central node, the road to which it is connected, otherwise None
    :param rand: Number generator with a fixed seed, or None for no fixed seed
    :return:
    """
    is_slave = roads is not None
    dist_to_parent = None

    # the whole count of nodes generated in this function
    count_nodes = 0

    if not is_slave:
        dist_to_parent = gen_c.DIST_TO_PARENT.rand_float(rand)
        roads, count_generated = _get_roads([], cluster_name, dist_to_parent, rand)
        count_nodes += count_generated

    engineering_preparation = _get_engineering_preparation([roads['min']], cluster_name, sum(borehole_counts), rand)
    preparation_stage = [roads['temp'], engineering_preparation]
    count_nodes += len(preparation_stage)

    dist_line = gen_c.ONE_SECTION_PIPE.rand_float(rand)
    power_lines, count_generated = _get_power_lines(preparation_stage, cluster_name, dist_line,
                                                    dist_high_line=dist_to_parent, rand=rand)
    count_nodes += count_generated

    pipe_dists = [gen_c.ONE_SECTION_PIPE.rand_float(rand) for _ in range(pipe_nodes_count + 1)]
    pipe_lines_and_nodes, count_generated = _get_pipe_lines(preparation_stage, cluster_name, pipe_dists, rand)
    count_nodes += count_generated

    boreholes_eq_shared, count_generated = _get_boreholes_equipment_shared(preparation_stage,
                                                                           cluster_name, rand)
    count_nodes += count_generated
    boreholes_eq_general, count_generated = _get_boreholes_equipment_general(preparation_stage,
                                                                             cluster_name,
                                                                             pipes_count=pipe_net_count,
                                                                             masts_count=light_masts_count,
                                                                             rand=rand)
    count_nodes += count_generated

    boreholes_all = []
    for i, borehole_count in enumerate(borehole_counts):
        boreholes_eq, count_generated = _get_boreholes_equipment_group(preparation_stage, cluster_name,
                                                                       group_ind=i,
                                                                       borehole_count=borehole_count,
                                                                       rand=rand)
        count_nodes += count_generated

        boreholes, count_generated = _get_boreholes(boreholes_eq + boreholes_eq_shared, cluster_name,
                                                    group_ind=i, borehole_count=borehole_count, rand=rand)
        count_nodes += count_generated
        boreholes_all += boreholes

    all_stages = [roads['final']] + power_lines + pipe_lines_and_nodes + boreholes_all + boreholes_eq_general
    handing_stage = _get_handing_stage(all_stages, cluster_name, sum(borehole_counts), group_ind=0, rand=rand)
    return handing_stage, roads, count_nodes
