import sys
import threading
import uuid
from abc import abstractmethod
from contextlib import AbstractContextManager
from typing import Generic, NamedTuple, Optional, TypeVar, Union, cast

import pendulum

import dagster._check as check
from dagster.core.errors import DagsterUserCodeProcessError
from dagster.core.host_representation.origin import (
    ManagedGrpcPythonEnvRepositoryLocationOrigin,
    RepositoryLocationOrigin,
)
from dagster.core.types.loadable_target_origin import LoadableTargetOrigin
from dagster.grpc.client import DagsterGrpcClient
from dagster.grpc.server import GrpcServerProcess
from dagster.utils.error import SerializableErrorInfo, serializable_error_info_from_exc_info


class GrpcServerEndpoint(
    NamedTuple(
        "_GrpcServerEndpoint",
        [
            ("server_id", str),
            ("host", str),
            ("port", Optional[int]),
            ("socket", Optional[str]),
        ],
    )
):
    def __new__(cls, server_id, host, port, socket):
        return super(GrpcServerEndpoint, cls).__new__(
            cls,
            check.str_param(server_id, "server_id"),
            check.str_param(host, "host"),
            check.opt_int_param(port, "port"),
            check.opt_str_param(socket, "socket"),
        )

    def create_client(self) -> DagsterGrpcClient:
        return DagsterGrpcClient(port=self.port, socket=self.socket, host=self.host)


T = TypeVar("T")

# Daemons in different threads can use a shared GrpcServerRegistry to ensure that
# a single GrpcServerProcess is created for each origin
class GrpcServerRegistry(AbstractContextManager, Generic[T]):
    @abstractmethod
    def supports_origin(self, repository_location_origin: RepositoryLocationOrigin) -> bool:
        pass

    @abstractmethod
    def get_grpc_endpoint(self, repository_location_origin: T) -> GrpcServerEndpoint:
        pass

    @abstractmethod
    def reload_grpc_endpoint(self, repository_location_origin: T) -> GrpcServerEndpoint:
        pass

    @property
    @abstractmethod
    def supports_reload(self) -> bool:
        pass


class ProcessRegistryEntry(
    NamedTuple(
        "_ProcessRegistryEntry",
        [
            ("process_or_error", Union[GrpcServerProcess, SerializableErrorInfo]),
            ("loadable_target_origin", LoadableTargetOrigin),
            ("creation_timestamp", float),
            ("server_id", Optional[str]),
        ],
    )
):
    def __new__(cls, process_or_error, loadable_target_origin, creation_timestamp, server_id):
        return super(ProcessRegistryEntry, cls).__new__(
            cls,
            check.inst_param(
                process_or_error, "process_or_error", (GrpcServerProcess, SerializableErrorInfo)
            ),
            check.inst_param(
                loadable_target_origin, "loadable_target_origin", LoadableTargetOrigin
            ),
            check.float_param(creation_timestamp, "creation_timestamp"),
            check.opt_str_param(server_id, "server_id"),
        )


# GrpcServerRegistry that creates local gRPC python processes from
# ManagedGrpcPythonEnvRepositoryLocationOrigins and shares them between threads.
class ProcessGrpcServerRegistry(GrpcServerRegistry):
    def __init__(
        self,
        # How long each process should run before a new process should be created the next
        # time a given origin is requested (which will pick up any changes that have been
        # made to the code)
        reload_interval,
        # How long the process can live without a heartbeat before it dies. You should ensure
        # that either heartbeat_ttl is greater than reload_interval (so that the process will reload
        # before it ends due to heartbeat failure), or if reload_interval is 0, that any processes
        # returned by this registry have at least one GrpcServerRepositoryLocation hitting the
        # server with a heartbeat while you want the process to stay running.
        heartbeat_ttl,
        # How long to wait for the server to start up and receive connections before timing out
        startup_timeout,
    ):
        # ProcessRegistryEntry map of servers being currently returned, keyed by origin ID
        self._active_entries = {}

        self._waited_for_processes = False

        check.invariant(
            heartbeat_ttl > reload_interval,
            "Heartbeat TTL must be larger than reload interval, or processes could die due to TTL failure before they are reloaded",
        )

        self._reload_interval = check.int_param(reload_interval, "reload_interval")
        self._heartbeat_ttl = check.int_param(heartbeat_ttl, "heartbeat_ttl")
        self._startup_timeout = check.int_param(startup_timeout, "startup_timeout")

        self._lock = threading.Lock()

        self._all_processes = []

        self._cleanup_thread_shutdown_event = None
        self._cleanup_thread = None

        if self._reload_interval > 0:
            self._cleanup_thread_shutdown_event = threading.Event()

            self._cleanup_thread = threading.Thread(
                target=self._clear_old_processes,
                name="grpc-server-registry-cleanup",
                args=(self._cleanup_thread_shutdown_event, self._reload_interval),
            )
            self._cleanup_thread.daemon = True
            self._cleanup_thread.start()

    def supports_origin(self, repository_location_origin):
        return isinstance(repository_location_origin, ManagedGrpcPythonEnvRepositoryLocationOrigin)

    @property
    def supports_reload(self):
        return True

    def reload_grpc_endpoint(
        self, repository_location_origin: ManagedGrpcPythonEnvRepositoryLocationOrigin
    ) -> GrpcServerEndpoint:
        check.inst_param(
            repository_location_origin, "repository_location_origin", RepositoryLocationOrigin
        )
        with self._lock:
            origin_id = repository_location_origin.get_id()
            if origin_id in self._active_entries:
                # Free the map entry for this origin so that _get_grpc_endpoint will create
                # a new process
                del self._active_entries[origin_id]

            return self._get_grpc_endpoint(repository_location_origin)

    def get_grpc_endpoint(
        self, repository_location_origin: ManagedGrpcPythonEnvRepositoryLocationOrigin
    ) -> GrpcServerEndpoint:
        check.inst_param(
            repository_location_origin, "repository_location_origin", RepositoryLocationOrigin
        )

        with self._lock:
            return self._get_grpc_endpoint(repository_location_origin)

    def _get_loadable_target_origin(
        self, repository_location_origin: ManagedGrpcPythonEnvRepositoryLocationOrigin
    ):
        check.inst_param(
            repository_location_origin,
            "repository_location_origin",
            ManagedGrpcPythonEnvRepositoryLocationOrigin,
        )
        return repository_location_origin.loadable_target_origin

    def _get_grpc_endpoint(
        self, repository_location_origin: ManagedGrpcPythonEnvRepositoryLocationOrigin
    ) -> GrpcServerEndpoint:
        origin_id = repository_location_origin.get_id()
        loadable_target_origin = self._get_loadable_target_origin(repository_location_origin)
        if not loadable_target_origin:
            raise Exception(
                f"No Python file/module information available for location {repository_location_origin.location_name}"
            )

        if not origin_id in self._active_entries:
            refresh_server = True
        else:
            active_entry = self._active_entries[origin_id]
            refresh_server = loadable_target_origin != active_entry.loadable_target_origin

        server_process: Union[GrpcServerProcess, SerializableErrorInfo]
        new_server_id: Optional[str]
        if refresh_server:
            try:
                new_server_id = str(uuid.uuid4())
                server_process = GrpcServerProcess(
                    loadable_target_origin=loadable_target_origin,
                    heartbeat=True,
                    heartbeat_timeout=self._heartbeat_ttl,
                    fixed_server_id=new_server_id,
                    startup_timeout=self._startup_timeout,
                )
                self._all_processes.append(server_process)
            except Exception:
                server_process = serializable_error_info_from_exc_info(sys.exc_info())
                new_server_id = None

            self._active_entries[origin_id] = ProcessRegistryEntry(
                process_or_error=server_process,
                loadable_target_origin=loadable_target_origin,
                creation_timestamp=pendulum.now("UTC").timestamp(),
                server_id=new_server_id,
            )

        active_entry = self._active_entries[origin_id]

        if isinstance(active_entry.process_or_error, SerializableErrorInfo):
            raise DagsterUserCodeProcessError(
                active_entry.process_or_error.to_string(),
                user_code_process_error_infos=[active_entry.process_or_error],
            )

        return GrpcServerEndpoint(
            server_id=active_entry.server_id,
            host="localhost",
            port=active_entry.process_or_error.port,
            socket=active_entry.process_or_error.socket,
        )

    # Clear out processes from the map periodically so that they'll be re-created the next
    # time the origins are requested. Lack of any heartbeats will ensure that the server will
    # eventually die once they're no longer being held by any threads.
    def _clear_old_processes(self, shutdown_event, reload_interval):
        while True:
            shutdown_event.wait(5)
            if shutdown_event.is_set():
                break

            current_time = pendulum.now("UTC").timestamp()
            with self._lock:
                origin_ids_to_clear = []

                for origin_id, entry in self._active_entries.items():
                    if (
                        current_time - entry.creation_timestamp > reload_interval
                    ):  # Use a different threshold for errors so they aren't cached as long?
                        origin_ids_to_clear.append(origin_id)

                for origin_id in origin_ids_to_clear:
                    del self._active_entries[origin_id]

                # Remove any dead processes from the all_processes map
                dead_process_indexes = []
                for index in range(len(self._all_processes)):
                    process = self._all_processes[index]
                    if not process.server_process.poll() is None:
                        dead_process_indexes.append(index)

                for index in reversed(dead_process_indexes):
                    del self._all_processes[index]

    def __exit__(self, exception_type, exception_value, traceback):
        if self._cleanup_thread:
            cast(threading.Event, self._cleanup_thread_shutdown_event).set()
            self._cleanup_thread.join()

        for process in self._all_processes:
            process.create_ephemeral_client().cleanup_server()

    def wait_for_processes(self):
        # Wait for any processes created by this registry. Generally not needed outside
        # of tests, since the processes have heartbeats and will end on their own once
        # they finish any outstanding executions.
        if self._waited_for_processes:
            return
        self._waited_for_processes = True
        for process in self._all_processes:
            process.wait()
