# Copyright 2024 Flower Labs GmbH. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Fleet Simulation Engine API."""


import json
import threading
import time
import traceback
from concurrent.futures import ThreadPoolExecutor
from logging import DEBUG, ERROR, INFO, WARN
from pathlib import Path
from queue import Empty, Queue
from time import sleep
from typing import Callable, Optional

from flwr.client.client_app import ClientApp, ClientAppException, LoadClientAppError
from flwr.client.clientapp.utils import get_load_client_app_fn
from flwr.client.run_info_store import DeprecatedRunInfoStore
from flwr.common.constant import (
    NUM_PARTITIONS_KEY,
    PARTITION_ID_KEY,
    PING_MAX_INTERVAL,
    ErrorCode,
)
from flwr.common.logger import log
from flwr.common.message import Error
from flwr.common.serde import message_from_taskins, message_to_taskres
from flwr.common.typing import Run
from flwr.proto.task_pb2 import TaskIns, TaskRes  # pylint: disable=E0611
from flwr.server.superlink.linkstate import LinkState, LinkStateFactory

from .backend import Backend, error_messages_backends, supported_backends

NodeToPartitionMapping = dict[int, int]


def _register_nodes(
    num_nodes: int, state_factory: LinkStateFactory
) -> NodeToPartitionMapping:
    """Register nodes with the StateFactory and create node-id:partition-id mapping."""
    nodes_mapping: NodeToPartitionMapping = {}
    state = state_factory.state()
    for i in range(num_nodes):
        node_id = state.create_node(ping_interval=PING_MAX_INTERVAL)
        nodes_mapping[node_id] = i
    log(DEBUG, "Registered %i nodes", len(nodes_mapping))
    return nodes_mapping


def _register_node_info_stores(
    nodes_mapping: NodeToPartitionMapping,
    run: Run,
    app_dir: Optional[str] = None,
) -> dict[int, DeprecatedRunInfoStore]:
    """Create DeprecatedRunInfoStore objects and register the context for the run."""
    node_info_store: dict[int, DeprecatedRunInfoStore] = {}
    num_partitions = len(set(nodes_mapping.values()))
    for node_id, partition_id in nodes_mapping.items():
        node_info_store[node_id] = DeprecatedRunInfoStore(
            node_id=node_id,
            node_config={
                PARTITION_ID_KEY: partition_id,
                NUM_PARTITIONS_KEY: num_partitions,
            },
        )

        # Pre-register Context objects
        node_info_store[node_id].register_context(
            run_id=run.run_id, run=run, app_dir=app_dir
        )

    return node_info_store


# pylint: disable=too-many-arguments,too-many-locals
def worker(
    taskins_queue: "Queue[TaskIns]",
    taskres_queue: "Queue[TaskRes]",
    node_info_store: dict[int, DeprecatedRunInfoStore],
    backend: Backend,
    f_stop: threading.Event,
) -> None:
    """Get TaskIns from queue and pass it to an actor in the pool to execute it."""
    while not f_stop.is_set():
        out_mssg = None
        try:
            # Fetch from queue with timeout. We use a timeout so
            # the stopping event can be evaluated even when the queue is empty.
            task_ins: TaskIns = taskins_queue.get(timeout=1.0)
            node_id = task_ins.task.consumer.node_id

            # Retrieve context
            context = node_info_store[node_id].retrieve_context(run_id=task_ins.run_id)

            # Convert TaskIns to Message
            message = message_from_taskins(task_ins)

            # Let backend process message
            out_mssg, updated_context = backend.process_message(message, context)

            # Update Context
            node_info_store[node_id].update_context(
                task_ins.run_id, context=updated_context
            )
        except Empty:
            # An exception raised if queue.get times out
            pass
        # Exceptions aren't raised but reported as an error message
        except Exception as ex:  # pylint: disable=broad-exception-caught
            log(ERROR, ex)
            log(ERROR, traceback.format_exc())

            if isinstance(ex, ClientAppException):
                e_code = ErrorCode.CLIENT_APP_RAISED_EXCEPTION
            elif isinstance(ex, LoadClientAppError):
                e_code = ErrorCode.LOAD_CLIENT_APP_EXCEPTION
            else:
                e_code = ErrorCode.UNKNOWN

            reason = str(type(ex)) + ":<'" + str(ex) + "'>"
            out_mssg = message.create_error_reply(
                error=Error(code=e_code, reason=reason)
            )

        finally:
            if out_mssg:
                # Convert to TaskRes
                task_res = message_to_taskres(out_mssg)
                # Store TaskRes in state
                task_res.task.pushed_at = time.time()
                taskres_queue.put(task_res)


def add_taskins_to_queue(
    state: LinkState,
    queue: "Queue[TaskIns]",
    nodes_mapping: NodeToPartitionMapping,
    f_stop: threading.Event,
) -> None:
    """Put TaskIns in a queue from State."""
    while not f_stop.is_set():
        for node_id in nodes_mapping.keys():
            task_ins_list = state.get_task_ins(node_id=node_id, limit=1)
            for task_ins in task_ins_list:
                queue.put(task_ins)
        sleep(0.1)


def put_taskres_into_state(
    state: LinkState, queue: "Queue[TaskRes]", f_stop: threading.Event
) -> None:
    """Put TaskRes into State from a queue."""
    while not f_stop.is_set():
        try:
            taskres = queue.get(timeout=1.0)
            state.store_task_res(taskres)
        except Empty:
            # queue is empty when timeout was triggered
            pass


# pylint: disable=too-many-positional-arguments
def run_api(
    app_fn: Callable[[], ClientApp],
    backend_fn: Callable[[], Backend],
    nodes_mapping: NodeToPartitionMapping,
    state_factory: LinkStateFactory,
    node_info_stores: dict[int, DeprecatedRunInfoStore],
    f_stop: threading.Event,
) -> None:
    """Run the VCE."""
    taskins_queue: "Queue[TaskIns]" = Queue()
    taskres_queue: "Queue[TaskRes]" = Queue()

    try:

        # Instantiate backend
        backend = backend_fn()

        # Build backend
        backend.build(app_fn)

        # Add workers (they submit Messages to Backend)
        state = state_factory.state()

        extractor_th = threading.Thread(
            target=add_taskins_to_queue,
            args=(
                state,
                taskins_queue,
                nodes_mapping,
                f_stop,
            ),
        )
        extractor_th.start()

        injector_th = threading.Thread(
            target=put_taskres_into_state,
            args=(
                state,
                taskres_queue,
                f_stop,
            ),
        )
        injector_th.start()

        with ThreadPoolExecutor() as executor:
            _ = [
                executor.submit(
                    worker,
                    taskins_queue,
                    taskres_queue,
                    node_info_stores,
                    backend,
                    f_stop,
                )
                for _ in range(backend.num_workers)
            ]

        extractor_th.join()
        injector_th.join()

    except Exception as ex:

        log(ERROR, "An exception occured!! %s", ex)
        log(ERROR, traceback.format_exc())
        log(WARN, "Stopping Simulation Engine.")

        # Manually trigger stopping event
        f_stop.set()

        # Raise exception
        raise RuntimeError("Simulation Engine crashed.") from ex

    finally:

        # Terminate backend
        backend.terminate()


# pylint: disable=too-many-arguments,unused-argument,too-many-locals,too-many-branches
# pylint: disable=too-many-statements,too-many-positional-arguments
def start_vce(
    backend_name: str,
    backend_config_json_stream: str,
    app_dir: str,
    is_app: bool,
    f_stop: threading.Event,
    run: Run,
    flwr_dir: Optional[str] = None,
    client_app: Optional[ClientApp] = None,
    client_app_attr: Optional[str] = None,
    num_supernodes: Optional[int] = None,
    state_factory: Optional[LinkStateFactory] = None,
    existing_nodes_mapping: Optional[NodeToPartitionMapping] = None,
) -> None:
    """Start Fleet API with the Simulation Engine."""
    nodes_mapping = {}

    if client_app_attr is not None and client_app is not None:
        raise ValueError(
            "Both `client_app_attr` and `client_app` are provided, "
            "but only one is allowed."
        )

    if num_supernodes is not None and existing_nodes_mapping is not None:
        raise ValueError(
            "Both `num_supernodes` and `existing_nodes_mapping` are provided, "
            "but only one is allowed."
        )
    if num_supernodes is None:
        if state_factory is None or existing_nodes_mapping is None:
            raise ValueError(
                "If not passing an existing `state_factory` and associated "
                "`existing_nodes_mapping` you must supply `num_supernodes` to indicate "
                "how many nodes to insert into a new StateFactory that will be created."
            )
    if existing_nodes_mapping:
        if state_factory is None:
            raise ValueError(
                "`existing_nodes_mapping` was passed, but no `state_factory` was "
                "passed."
            )
        log(INFO, "Using exiting NodeToPartitionMapping and StateFactory.")
        # Use mapping constructed externally. This also means nodes
        # have previously being registered.
        nodes_mapping = existing_nodes_mapping
    app_dir = str(Path(app_dir).absolute())

    if not state_factory:
        log(INFO, "A StateFactory was not supplied to the SimulationEngine.")
        # Create an empty in-memory state factory
        state_factory = LinkStateFactory(":flwr-in-memory-state:")
        log(INFO, "Created new %s.", state_factory.__class__.__name__)

    if num_supernodes:
        # Register SuperNodes
        nodes_mapping = _register_nodes(
            num_nodes=num_supernodes, state_factory=state_factory
        )

    # Construct mapping of DeprecatedRunInfoStore
    node_info_stores = _register_node_info_stores(
        nodes_mapping=nodes_mapping, run=run, app_dir=app_dir if is_app else None
    )

    # Load backend config
    log(DEBUG, "Supported backends: %s", list(supported_backends.keys()))
    backend_config = json.loads(backend_config_json_stream)

    try:
        backend_type = supported_backends[backend_name]
    except KeyError as ex:
        log(
            ERROR,
            "Backend `%s`, is not supported. Use any of %s or add support "
            "for a new backend.",
            backend_name,
            list(supported_backends.keys()),
        )
        if backend_name in error_messages_backends:
            log(ERROR, error_messages_backends[backend_name])

        raise ex

    def backend_fn() -> Backend:
        """Instantiate a Backend."""
        return backend_type(backend_config)

    # Load ClientApp if needed
    def _load() -> ClientApp:

        if client_app:
            return client_app
        if client_app_attr:
            return get_load_client_app_fn(
                default_app_ref=client_app_attr,
                app_path=app_dir,
                flwr_dir=flwr_dir,
                multi_app=False,
            )(run.fab_id, run.fab_version, run.fab_hash)

        raise ValueError("Either `client_app_attr` or `client_app` must be provided")

    app_fn = _load

    try:
        # Test if ClientApp can be loaded
        client_app = app_fn()

        # Cache `ClientApp`
        if client_app_attr:
            # Now wrap the loaded ClientApp in a dummy function
            # this prevent unnecesary low-level loading of ClientApp
            def _load_client_app() -> ClientApp:
                return client_app

            app_fn = _load_client_app

        # Run main simulation loop
        run_api(
            app_fn,
            backend_fn,
            nodes_mapping,
            state_factory,
            node_info_stores,
            f_stop,
        )
    except LoadClientAppError as loadapp_ex:
        f_stop_delay = 10
        log(
            ERROR,
            "LoadClientAppError exception encountered. Terminating simulation in %is",
            f_stop_delay,
        )
        time.sleep(f_stop_delay)
        f_stop.set()  # set termination event
        raise loadapp_ex
    except Exception as ex:
        raise ex
