import asyncio
import networkx
import inspect
import threading
import os
from typing import List, Dict, Union, Any, Tuple, Coroutine
from collections import defaultdict
from hedra.core.graphs.events import Event
from hedra.core.graphs.stages.base.import_tools import set_stage_hooks
from hedra.core.graphs.stages.base.stage import Stage
from hedra.core.graphs.stages.error import Error
from hedra.core.graphs.stages.types.stage_types import StageTypes
from hedra.core.graphs.hooks.registry.registrar import registrar
from hedra.core.graphs.hooks.registry.registry_types import EventHook
from hedra.core.graphs.hooks.registry.registry_types.hook import Hook, HookType
from hedra.core.graphs.stages.base.parallel.batch_executor import BatchExecutor
from hedra.core.graphs.transitions.exceptions.exceptions import IsolatedStageError
from hedra.logging import HedraLogger
from hedra.plugins.types.engine.engine_plugin import EnginePlugin
from hedra.plugins.types.reporter.reporter_plugin import ReporterPlugin
from hedra.plugins.types.plugin_types import PluginType
from .transition import Transition
from .common import (
    invalid_transition
)
from .idle import (
    invalid_idle_transition
)


class TransitionAssembler:

    def __init__(
        self, 
        transition_types, 
        graph_name: str=None,
        graph_path: str=None,
        graph_id: str=None,
        cpus: int=None, 
        worker_id: int=None,
        core_config: Dict[str, Any]={}
    ) -> None:
        self.transition_types: Dict[Tuple[StageTypes, StageTypes], Coroutine] = transition_types
        self.core_config = core_config
        self.graph_name = graph_name
        self.graph_path = graph_path
        self.graph_id = graph_id
        self.generated_stages = {}
        self.transitions = {}
        self.instances_by_type: Dict[str, List[Stage]] = {}
        self.cpus = cpus
        self.worker_id = worker_id
        self.loop = asyncio.get_event_loop()
        self.hooks_by_type: Dict[HookType, Dict[str, Hook]] = defaultdict(dict)

        self.logging = HedraLogger()
        self.logging.initialize()

        self._thread_id = threading.current_thread().ident
        self._process_id = os.getpid()

        self._graph_metadata_log_string = f'Graph - {self.graph_name}:{self.graph_id} - thread:{self._thread_id} - process:{self._process_id} - '

    def generate_stages(self, stages: Dict[str, Stage]):

        stages_count = len(stages)
        self.logging.hedra.sync.debug(f'{self._graph_metadata_log_string} - Generating - {stages_count} - stages')
        self.logging.filesystem.sync['hedra.core'].debug(f'{self._graph_metadata_log_string} - Generating - {stages_count} - stages')

        self.instances_by_type = {}

        for stage in stages.values():
            self.instances_by_type[stage.stage_type] = []

        stage_types_count = len(self.instances_by_type)
        self.logging.hedra.sync.debug(f'{self._graph_metadata_log_string} - Found - {stage_types_count} - unique stage types')
        self.logging.filesystem.sync['hedra.core'].debug(f'{self._graph_metadata_log_string} - Found - {stage_types_count} - unique stage types')

        self.generated_stages: Dict[str, Stage] = {
            stage_name: stage() for stage_name, stage in stages.items()
        }

        for stage in self.generated_stages.values():
            stage.core_config = self.core_config
            stage.graph_name = self.graph_name
            stage.graph_path = self.graph_path
            stage.graph_id = self.graph_id

            stage.workers = self.cpus
            stage.worker_id = self.worker_id

            for hook_shortname, hook in registrar.reserved[stage.name].items():
                    hook._call = hook._call.__get__(stage, stage.__class__)
                    setattr(stage, hook_shortname, hook._call)


            stage = set_stage_hooks(stage)

            for hook_type in stage.hooks:
                self.hooks_by_type[hook_type].update({
                    hook.name: hook for hook in stage.hooks[hook_type]
                })

            methods = inspect.getmembers(stage, predicate=inspect.ismethod) 

            for _, method in methods:
                method_name = method.__qualname__
                hook: Hook = registrar.all.get(method_name)
                
                if hook:
                    hook._call = hook._call.__get__(stage, stage.__class__)
                    setattr(stage, hook.shortname, hook._call)

                    if inspect.ismethod(hook.call) is False:
                        hook.call = hook.call.__get__(stage, stage.__class__)
                        setattr(stage, hook.shortname, hook.call)

                    hook.stage = stage.name
                    hook.stage_instance: Stage = stage
                    
                    self.hooks_by_type[hook.hook_type][hook.name] = hook
                    stage.hooks[hook.hook_type].append(hook)
            
            self.instances_by_type[stage.stage_type].append(stage)


            event_hooks = []
            stage_event_hooks: List[EventHook] = stage.hooks[HookType.EVENT]
            for event_hook in stage_event_hooks:
                if event_hook.names is None or len(event_hook.names)  == 0:
                    event_hooks.append(event_hook)

            stage.hooks[HookType.EVENT] = event_hooks

        event_hooks: List[EventHook] = list(self.hooks_by_type.get(HookType.EVENT, {}).values())
        for event_hook in event_hooks:
            for target_hook_name in event_hook.names:    
                target_hook = registrar.all.get(target_hook_name)
                
                if target_hook:
                    self.logging.filesystem.sync['hedra.core'].info(
                        f'{self._graph_metadata_log_string} - Appendng Event - {event_hook.name}:{event_hook.hook_id} - to target Stage - {target_hook.stage}:{target_hook.stage_instance.stage_id} Event Hooks'
                    )

                    event = Event(target_hook, event_hook)
                    event.target_key = event_hook.key

                    target_hook_idx = -1

                    try:
                        target_hook_idx = target_hook.stage_instance.hooks[target_hook.hook_type].index(target_hook)

                    except ValueError:
                        pass

                    if target_hook_idx >= 0 and isinstance(target_hook, Event):
                        if event_hook.pre is True:
                            target_hook.pre_sources[event_hook.name] = event_hook
                            target_hook.stage_instance.hooks[target_hook.hook_type][target_hook_idx] = target_hook

                        else:
                            target_hook.post_sources[event_hook.name] = event_hook
                            target_hook.stage_instance.hooks[target_hook.hook_type][target_hook_idx] = target_hook
                        
                        registrar.all[event.name] = target_hook

                    elif target_hook_idx >= 0:
                        target_hook.stage_instance.hooks[target_hook.hook_type][target_hook_idx] = event
                        registrar.all[event.name] = event
                    
                    target_hook.stage_instance.linked_events[(target_hook.stage, target_hook.hook_type, target_hook.name)].append(
                        (event_hook.stage,  event_hook.hook_type, event_hook.name)
                    )
               
    
        self.logging.hedra.sync.debug(f'{self._graph_metadata_log_string} - Successfully generated - {stages_count} - stages')
        self.logging.filesystem.sync['hedra.core'].debug(f'{self._graph_metadata_log_string} - Successfully generated - {stages_count} - stages')

    def build_transitions_graph(self, topological_generations: List[List[str]], graph: networkx.Graph):

        self.logging.hedra.sync.debug(f'Buiding transitions matrix')
        self.logging.filesystem.sync['hedra.core'].debug(f'Buiding transitions matrix')

        transitions: List[List[Transition]] = []
        plugins: Dict[PluginType, Dict[str, Union[EnginePlugin, ReporterPlugin]]] = {
            PluginType.ENGINE: {},
            PluginType.OPTIMIZER: {},
            PluginType.PERSONA: {},
            PluginType.REPORTER: {}
        }

        for isolate_stage_name in networkx.isolates(graph):
            raise IsolatedStageError(
                self.generated_stages.get(isolate_stage_name)
            )
        
        for generation in topological_generations:

            generation_transitions = []

            stage_pool_size = self.cpus

            stages = {
                stage_name: self.generated_stages.get(stage_name) for stage_name in generation
            }
            parallel_stages = []

            no_workers_stages = [
                StageTypes.WAIT, 
                StageTypes.IDLE
            ]

            stages_count = len(stages)
            self.logging.hedra.sync.debug(f'{self._graph_metadata_log_string} - Provisioning workers - {stages_count} - stages')
            self.logging.filesystem.sync['hedra.core'].debug(f'{self._graph_metadata_log_string} - Provisioning workers- {stages_count} - stages')

            for stage in stages.values():

                for plugin_name, plugin in stage.plugins.items():
                    plugins[plugin.type][plugin_name] = plugin

                if stage.allow_parallel is False and stage.stage_type not in no_workers_stages:
                    stage.workers = 1
                    stage_pool_size -= 1

                    self.logging.hedra.sync.debug(f'{self._graph_metadata_log_string} - Stage - {stage.name} - provisioned - {stage.workers} - workers')
                    self.logging.filesystem.sync['hedra.core'].debug(f'{self._graph_metadata_log_string} - Stage - {stage.name} - provisioned - {stage.workers} - workers')

                    stage.executor = BatchExecutor(stage.workers)

                else:
                    parallel_stages.append((
                        stage.name,
                        stage
                    ))
            
            if len(parallel_stages) > 0:
                batch_executor = BatchExecutor(max_workers=stage_pool_size)

                batched_stages: List[Tuple[str, Stage, int]] = batch_executor.partion_stage_batches(parallel_stages)
                
                for _, stage, assigned_workers_count in batched_stages:

                    self.logging.hedra.sync.debug(f'{self._graph_metadata_log_string} - Stage - {stage.name} - provisioned - {assigned_workers_count} - workers')
                    self.logging.filesystem.sync['hedra.core'].debug(f'{self._graph_metadata_log_string} - Stage - {stage.name} - provisioned - {assigned_workers_count} - workers')

                    stage.workers = assigned_workers_count
                    stage.executor = BatchExecutor(max_workers=assigned_workers_count)

                    stages[stage.name] = stage

            for stage in stages.values():

                stage.plugins_by_type = plugins

                neighbors = list(graph.neighbors(stage.name))

                neighbors_count = len(neighbors)
                self.logging.hedra.sync.debug(f'{self._graph_metadata_log_string} - Discovered - {neighbors_count} - neighboring stages for stage - {stage.name}')
                self.logging.filesystem.sync['hedra.core'].debug(f'{self._graph_metadata_log_string} - Discovered - {neighbors_count} - neighboring stages for stage - {stage.name}')
                
                for neighbor in neighbors:
                    neighbor_stage = self.generated_stages.get(neighbor)

                    transition_action = self.transition_types.get((
                        stage.stage_type,
                        neighbor_stage.stage_type
                    ))

                    if transition_action == invalid_transition or transition_action == invalid_idle_transition:
                        invalid_transition_error, _ = self.loop.run_until_complete(
                            transition_action(stage, neighbor_stage)
                        )

                        raise invalid_transition_error

                    self.logging.hedra.sync.debug(f'{self._graph_metadata_log_string} - Created transition from - {stage.name} - to - {neighbor_stage.name}')
                    self.logging.filesystem.sync['hedra.core'].debug(f'{self._graph_metadata_log_string} - Created transition from - {stage.name} - to - {neighbor_stage.name}')

                    transition = Transition(
                        transition_action,
                        stage,
                        neighbor_stage
                    )

                    generation_transitions.append(transition)

            if len(generation_transitions) > 0:
                transitions.append(generation_transitions)

        self.logging.hedra.sync.debug(f'{self._graph_metadata_log_string} - Transition matrix assemmbly complete')
        self.logging.filesystem.sync['hedra.core'].debug(f'{self._graph_metadata_log_string} - Transition matrix assemmbly complete')

        return transitions

    def map_to_setup_stages(self, graph: networkx.DiGraph):

        self.logging.hedra.sync.debug(f'{self._graph_metadata_log_string} - Mapping stages to requisite Setup stages')
        self.logging.filesystem.sync['hedra.core'].debug(f'{self._graph_metadata_log_string} - Mapping stages to requisite Setup stages')
        
        idle_stages = self.instances_by_type.get(StageTypes.IDLE)
        for idle_stage in idle_stages:
            idle_stage.context.stages = {}
            idle_stage.context.visited = []
            idle_stage.context.results = {}
            idle_stage.context.results_stages = []
            idle_stage.context.summaries = {}
            idle_stage.context.paths = {}
            idle_stage.context.path_lengths = {}
            
        idle_stage_name = idle_stage.__class__.__name__

        complete_stage = self.instances_by_type.get(StageTypes.COMPLETE)[0]

        for stage_type in StageTypes:

            idle_stage.context.stages[stage_type] = {}

            for stage in self.instances_by_type.get(stage_type, []):

                stage_name = stage.__class__.__name__

                has_path = networkx.has_path(
                    graph, 
                    idle_stage_name,
                    stage_name
                )

                if has_path:
                    idle_stage.context.stages[stage_type][stage_name] = stage
                    paths = networkx.all_shortest_paths(graph, stage_name, complete_stage.name)
                
                    stage_paths = []
                    for path in paths:
                        stage_paths.extend(path)
                    
                    idle_stage.context.paths[stage_name] = stage_paths

                    path_lengths = networkx.all_pairs_shortest_path_length(graph)

                    stage_path_lengths = {}
                    for path_stage_name, path_lengths_set in path_lengths:

                        del path_lengths_set[path_stage_name]
                        stage_path_lengths[path_stage_name] = path_lengths_set

                    idle_stage.context.path_lengths[stage_name] = stage_path_lengths.get(stage_name)

        for stage in self.generated_stages.values():
            for idle_stage in idle_stages:
                stage.context = idle_stage.context

        self.logging.hedra.sync.debug(f'{self._graph_metadata_log_string} - Mapped stages to requisite Setup stages')
        self.logging.filesystem.sync['hedra.core'].debug(f'{self._graph_metadata_log_string} - Mapped stages to requisite Setup stages')
    
    def create_error_transition(self, error: Exception):

        from_stage = error.from_stage
            
        error_transition = self.transition_types.get((
            from_stage.stage_type,
            StageTypes.ERROR
        ))

        error_stage = Error()
        error_stage.graph_name = self.graph_name
        error_stage.graph_id = self.graph_id
        error_stage.error = error

        return Transition(
            error_transition,
            from_stage,
            error_stage
        )