import hashlib
import os
import subprocess
import sys
import time
from typing import Any, Dict, List, Optional, Tuple, Union

import colorama  # type: ignore
import ray
import yaml

from anyscale.client.openapi_client.models.session import Session  # type: ignore
from anyscale.credentials import load_credentials
from anyscale.fingerprint import fingerprint
import anyscale.project
from anyscale.sdk.anyscale_client.sdk import AnyscaleSDK  # type: ignore

# The default project directory to use, if no project is specified.
DEFAULT_SCRATCH_DIR = "~/.anyscale/scratch_project"

# The key used to identify the session's latest file hash.
ANYSCALE_FILES_HASH = "anyscale.com/local_files_hash"

# Max number of auto created sessions.
MAX_SESSIONS = 20


class _ConsoleLog:
    def __init__(self) -> None:
        self.t0 = time.time()

    def zero_time(self) -> None:
        self.t0 = time.time()

    def info(self, *msg: Any) -> None:
        print(
            "{}{}(anyscale +{}){} ".format(
                colorama.Style.BRIGHT,
                colorama.Fore.CYAN,
                self._time_string(),
                colorama.Style.RESET_ALL,
            ),
            end="",
        )
        print(*msg)

    def debug(self, *msg: str) -> None:
        if os.environ.get("DEBUG") == "1":
            print(
                "{}{}(anyscale +{}){} ".format(
                    colorama.Style.DIM,
                    colorama.Fore.CYAN,
                    self._time_string(),
                    colorama.Style.RESET_ALL,
                ),
                end="",
            )
            print(*msg)

    def _time_string(self) -> str:
        delta = time.time() - self.t0
        hours = 0
        minutes = 0
        while delta > 3600:
            hours += 1
            delta -= 3600
        while delta > 60:
            minutes += 1
            delta -= 60
        output = ""
        if hours:
            output += "{}h".format(hours)
        if minutes:
            output += "{}m".format(minutes)
        output += "{}s".format(round(delta, 1))
        return output


def _write_session_yaml(project_dir: str, requirements: Optional[str]) -> None:
    """Write session-default.yaml into project directory.

    Args:
        project_dir (str): Path to the project directory.
        requirements (Optional[str]): List of requirements in
            requirements.txt format.
    """
    # TODO(ekl): generate a serverless compute configuraton here.
    cluster_yaml = yaml.safe_load(anyscale.project.CLUSTER_YAML_TEMPLATE)
    cluster_yaml["setup_commands"] = [
        "pip install -U https://s3-us-west-2.amazonaws.com/ray-wheels/"
        "latest/ray-1.2.0.dev0-cp37-cp37m-manylinux2014_x86_64.whl"
    ]
    # Install requirements:
    if requirements:
        cluster_yaml["setup_commands"].append(
            f'echo "{requirements}" | pip install -r /dev/stdin'
        )
    # TODO(ekl) ray client server should be started by default
    cluster_yaml["head_start_ray_commands"] = [
        "ray stop",
        "ulimit -n 65536; cd ~/"
        + os.path.basename(project_dir)
        + "; ray start --head --port=6379 "
        "--object-manager-port=8076 "
        "--ray-client-server-port=10001 "
        "--autoscaling-config=~/ray_bootstrap_config.yaml",
    ]
    with open(os.path.join(project_dir, "session-default.yaml"), "w+") as f:
        f.write(yaml.dump(cluster_yaml))


class SessionBuilder:
    """This class lets you set session options and connect to Anyscale.

    This feature is ***EXPERIMENTAL***.

    It should not be constructed directly, but instead via anyscale.* methods
    exported at the package level.

    Examples:
        >>> # Raw client, creates new session on behalf of user
        >>> anyscale.connect()

        >>> # Get or create a named session
        >>> anyscale
        ...   .session("my_named_session")
        ...   .connect()

        >>> # Specify a previously created env / app template
        >>> anyscale
        ...   .app_env("prev_created_env:v2")
        ...   .autosuspend(hours=2)
        ...   .connect()

        >>> # Create new session from local env / from scratch
        >>> anyscale
        ...   .project_dir("~/dev/my-project-folder")
        ...   .base_docker_image("anyscale/ray-ml:latest")
        ...   .require("~/dev/my-project-folder/requirements.txt")
        ...   .connect()

        >>> # Ray client connect is setup automatically
        >>> @ray.remote
        ... def my_func(value):
        ...   return value ** 2

        >>> # Remote functions are executed in the Anyscale session
        >>> print(ray.get([my_func.remote(x) for x in range(5)]))
        >>> [0, 1, 4, 9, 16]
    """

    def __init__(
        self,
        scratch_dir: str = DEFAULT_SCRATCH_DIR,
        anyscale_sdk: AnyscaleSDK = None,
        subprocess: Any = subprocess,
        ray: Any = ray,
        log: Any = _ConsoleLog(),
    ) -> None:
        # Class dependencies.
        self._anyscale_sdk: Any = None
        if anyscale_sdk:
            self._anyscale_sdk = anyscale_sdk
        else:
            self._anyscale_sdk = AnyscaleSDK(load_credentials())
        self._log = log
        self._ray: Any = ray
        self._subprocess: Any = subprocess

        # Builder args.
        self._scratch_dir: str = scratch_dir
        self._project_dir: Optional[str] = None
        self._session_name: Optional[str] = None
        self._requirements: Optional[str] = None

        # Whether to update the session when connecting to a fixed session.
        self._needs_update: bool = True

        # A temporary connection to use to lock the right session.
        self._tmp_conn: Any = None

    def connect(self) -> None:
        """Connect to Anyscale using previously specified options.

        Examples:
            >>> anyscale.connect()
        """

        # TODO(ekl) check for duplicate connections
        self._log.zero_time()

        # Autodetect or create a scratch project.
        if self._project_dir is None:
            self._project_dir = anyscale.project.find_project_root(os.getcwd())
        if self._project_dir:
            self._ensure_project_setup_at_dir(self._project_dir)
        else:
            self._project_dir = self._get_or_create_scratch_project()
        proj_def = anyscale.project.ProjectDefinition(self._project_dir)
        project_id = anyscale.project.get_project_id(proj_def.root)
        self._log.info("Using project dir", proj_def.root)

        _write_session_yaml(self._project_dir, self._requirements)

        # Run finger printing after session-default.yaml is written so it is part
        # of the finger print.
        files_hash = self._fingerprint(proj_def.root)

        # Locate a updated session and update it if needed.
        if self._session_name is None or self._needs_update:
            self._session_name = self._get_or_create_updated_session(
                files_hash.encode("utf-8"), project_id, self._session_name
            )

        # Connect Ray client.
        self._connect_to_session(project_id, self._session_name)

        # Can release the session lock now that we are connected for real.
        if self._tmp_conn:
            self._tmp_conn.disconnect()
            self._tmp_conn = None

    def project_dir(self, local_dir: str) -> "SessionBuilder":
        """Set the project directory.

        This sets the project code directory that will be synced to all nodes
        in the cluster as required by Ray. If not specified, the project
        directory will be autodetected based on the current working directory.
        If no Anyscale project is found, a "scratch" project will be used.

        Args:
            local_dir (str): path to the project directory.

        Examples:
            >>> anyscale.project_dir("~/my-proj-dir").connect()
        """
        self._project_dir = os.path.abspath(os.path.expanduser(local_dir))
        return self

    def session(self, session_name: str, update: bool = False) -> "SessionBuilder":
        """Set a fixed session name.

        By default, Anyscale connect will pick an idle session updated
        with the connect parameters, creating a new session if no updated
        idle sessions. Setting a fixed session name will force connecting to
        the given named session, creating it if it doesn't exist.

        Args:
            session_name (str): fixed name of the session.
            update (bool): whether to update session configurations when
                connecting to an existing session. Note that this may restart
                the Ray runtime.

        Examples:
            >>> anyscale.session("prod_deployment", update=True).connect()
        """
        if not update:
            self._needs_update = False
        self._session_name = session_name
        return self

    def base_docker_image(self, image_name: str) -> "SessionBuilder":
        """Set the docker image to use for the session.

        Args:
            image_name (str): docker image name.

        Examples:
            >>> anyscale.base_docker_image("anyscale/ray-ml:latest").connect()
        """
        raise NotImplementedError()

    def require(self, requirements: Union[str, List[str]]) -> "SessionBuilder":
        """Set the Python requirements for the session.

        Args:
            requirements: either be a list of pip library specifications, or
            the path to a requirements.txt file.

        Examples:
            >>> anyscale.require("~/proj/requirements.txt").connect()
            >>> anyscale.require(["gym", "torch>=1.4.0"]).connect()
        """
        if isinstance(requirements, str):
            with open(requirements, "r") as f:
                data = f.read()
            # Escape quotes
            self._requirements = data.replace('"', r"\"")
        else:
            assert isinstance(requirements, list)
            self._requirements = "\n".join(requirements)
        return self

    def app_env(self, env_name: str) -> "SessionBuilder":
        """Set the Anyscale app env to use for the session.

        Args:
            env_name (str): app env name.

        Examples:
            >>> anyscale.app_env("prev_created_env:v2").connect()
        """
        raise NotImplementedError()

    def file_mount(self, *, local_dir: str, remote_dir: str) -> "SessionBuilder":
        """Add additional directories to sync up to worker nodes.

        Args:
            local_dir (str): the local directory path to mount.
            remote_dir (str): where in the remote node to mount the local dir.

        Examples:
            >>> anyscale
            ...   .file_mount(local_dir="~/data1", remote_dir="/tmp/d1")
            ...   .file_mount(local_dir="~/data2", remote_dir="/tmp/d2")
            ...   .connect()
        """
        raise NotImplementedError()

    def download_results(
        self, *, remote_dir: str, local_dir: str, autosync: bool = False
    ) -> "SessionBuilder":
        """Specify a directory to sync down from the cluster head node.

        Args:
            remote_dir (str): the remote result dir on the head node.
            local_dir (str): the local path to sync results to.
            autosync (bool): whether to sync the files continuously. By
                default, results will only be synced on job completion.

        Examples:
            >>> anyscale
            ...   .download_results(
            ...       remote_dir="~/ray_results", remote_dir="~/proj_output",
            ...       autosync=True)
            ...   .connect()
        """
        raise NotImplementedError()

    def autoshutdown(
        self, enabled: bool = True, *, hours: int = 1, minutes: int = 0,
    ) -> "SessionBuilder":
        """Configure or disable session autoshutdown behavior.

        The session will be autoshutdown after the specified time period. By
        default, sessions auto terminate after one hour of idle.

        Args:
            enabled (bool): whether autoshutdown is enabled.
            hours (int): specify idle time in hours.
            minutes (int): specify idle time in minutes. This is added to the
                idle time in hours.

        Examples:
            >>> anyscale.autoshutdown(False).connect()
            >>> anyscale.autoshutdown(hours=10).connect()
            >>> anyscale.autoshutdown(hours=1, minutes=30).connect()
        """
        raise NotImplementedError()

    def initial_scale(
        self,
        num_cpus: Optional[int] = None,
        num_gpus: Optional[int] = None,
        resource_bundles: Optional[List[Dict[str, float]]] = None,
    ) -> "SessionBuilder":
        """Configure the initial resources to scale to.

        The session will immediately attempt to scale to accomodate the
        requested resources, bypassing normal upscaling speed constraints.
        The requested resources are pinned and exempt from downscaling.

        Args:
            num_cpus (int): number of cpus to request.
            num_gpus (int): number of gpus to request.
            resource_bundles (List[Dict[str, float]): resource bundles to
                request. Each bundle is a dict of resource_name to quantity
                that can be allocated on a single machine. Note that the
                ``num_cpus`` and ``num_gpus`` args simply desugar into
                ``[{"CPU": 1}] * num_cpus`` and ``[{"GPU": 1}] * num_gpus``
                respectively.

        Examples:
            >>> anyscale.initial_scale(num_cpus=200, num_gpus=30).connect()
            >>> anyscale.initial_scale(
            ...     num_cpus=8,
            ...     resource_bundles=[{"GPU": 8}, {"GPU": 8}, {"GPU": 1}],
            ... ).connect()
        """
        raise NotImplementedError()

    def _get_or_create_scratch_project(self) -> str:
        """Get or create a scratch project, including the directory."""
        project_dir = os.path.expanduser(self._scratch_dir)
        self._ensure_project_setup_at_dir(project_dir)
        return project_dir

    def _ensure_project_setup_at_dir(self, project_dir: str) -> None:
        """Get or create an Anyscale project rooted at the given dir."""
        os.makedirs(project_dir, exist_ok=True)
        proj_name = os.path.basename(project_dir)

        # If the project yaml exists, assume we're already setup.
        project_yaml = os.path.join(project_dir, ".anyscale.yaml")
        if os.path.exists(project_yaml):
            return

        self._log.info("Creating new project for", project_dir)
        project_response = self._anyscale_sdk.create_project(
            {
                "name": proj_name,
                "description": "Automatically created by Anyscale Connect",
            }
        )
        project_id = project_response.result.id

        if not os.path.exists(project_yaml):
            with open(project_yaml, "w+") as f:
                f.write(yaml.dump({"project_id": project_id}))

    def _get_or_create_updated_session(  # noqa: C901
        self, files_hash: bytes, project_id: str, session_name: Optional[str]
    ) -> str:
        """Get or create a session updated with the given hash.

        Args:
            files_hash (str): The files hash.
            project_id (str): The project to use.
            session_name (Optional[str]): If specified, the given session
                will be created or updated as needed. Otherwise the session
                name will be picked automatically.

        Returns:
            The name of the session to connect to.
        """

        session_hash = None

        if session_name is None:
            self._log.debug("-> Looking for any running session matching hash")
            query = self._anyscale_sdk.search_sessions(project_id=project_id)
            self._log.debug("Session states", [s.state for s in query.results])
            for session in query.results:
                if session.state != "Running":
                    continue
                self._log.debug("Trying to acquire lock on", session.name)
                self._tmp_conn = self._acquire_session_lock(
                    session, raise_connection_error=False, connection_retries=0
                )
                if self._tmp_conn is None:
                    continue
                session_hash = self._tmp_conn._internal_kv_get(ANYSCALE_FILES_HASH)
                if session_hash == files_hash:
                    self._log.debug("Acquired lock on session", session.name)
                    session_name = session.name
                    break
                else:
                    self._tmp_conn.disconnect()
                    self._tmp_conn = None

        session_name_conn_failed = None
        if session_name is None:
            self._log.debug("-> Looking for any running session (ignoring hash)")
            for session in query.results:
                if session.state != "Running":
                    continue
                self._log.debug("Trying to acquire lock on", session.name)
                try:
                    self._tmp_conn = self._acquire_session_lock(
                        session, raise_connection_error=True, connection_retries=0
                    )
                except ConnectionError:
                    self._log.debug(
                        "ConnectionError to session, will revisit", session.name
                    )
                    session_name_conn_failed = session.name
                    continue
                if self._tmp_conn:
                    self._log.debug("Acquired lock on session", session.name)
                    session_name = session.name
                    break

        if session_name is None and session_name_conn_failed:
            self._log.debug("-> Fallback to restarting an errored session")
            session_name = session_name_conn_failed

        if session_name is None:
            self._log.debug("-> Fallback to starting a new session")
            used_names = [s.name for s in query.results if s.state == "Running"]
            for i in range(MAX_SESSIONS):
                name = "session-{}".format(i)
                if name not in used_names:
                    session_name = name
                    self._log.info("Starting session", session_name)
                    break

        # Should not happen.
        if session_name is None:
            raise RuntimeError("Could not create new session to connect to.")

        self._log.debug("Session hash", session_hash)
        self._log.debug("Local files hash", files_hash)

        if session_hash != files_hash:
            self._log.info("Syncing latest project files to", session_name)
            # TODO(ekl): race condition here since "up" breaks the lock.
            if self._tmp_conn:
                self._tmp_conn.disconnect()
                self._tmp_conn = None
            self._subprocess.check_call(
                [
                    "anyscale",
                    "up",
                    "--config",
                    "session-default.yaml",
                    "--cloud-name",
                    "anyscale_default_cloud",
                    session_name,
                ],
                cwd=self._project_dir,
            )
            # Need to re-acquire the connection after the update.
            self._tmp_conn = self._acquire_session_lock(
                self._get_session_or_die(project_id, session_name),
                raise_connection_error=True,
                connection_retries=10,
            )
            # TODO(ekl) retry this, might be race condition with another client
            if not self._tmp_conn:
                raise RuntimeError("Failed to acquire session we created")
            self._tmp_conn._internal_kv_put(
                ANYSCALE_FILES_HASH, files_hash, overwrite=True
            )
            self._log.debug(
                "Updated files hash",
                files_hash,
                self._tmp_conn._internal_kv_get(ANYSCALE_FILES_HASH),
            )

        return session_name

    def _acquire_session_lock(
        self, session_meta: Any, raise_connection_error: bool, connection_retries: int
    ) -> Optional[Any]:
        """Connect to and acquire a lock on the session.

        The session lock is released by calling disconnect() on the returned
        Ray connection.

        Returns:
            Connection to the session, or None if a lock could not be acquired.
        """
        # TODO(ekl) refactor Ray client to avoid this internal API access.
        conn = self._ray.util.client.RayAPIStub()
        session_url, secure, metadata = self._get_connect_params(session_meta)
        if connection_retries > 0:
            self._log.debug("Beginning connection attempts")
        try:
            # Disable retries when acquiring session lock for fast failure.
            info = conn.connect(
                session_url, secure, metadata, connection_retries=connection_retries
            )
        except Exception as e:
            if raise_connection_error:
                self._log.debug(
                    "Connection error after {} retries".format(connection_retries)
                )
                raise
            else:
                self._log.debug("Ignoring connection error", e)
                return None
        if info["num_clients"] > 1:
            conn.disconnect()
            return None
        else:
            return conn

    def _connect_to_session(self, project_id: str, session_name: str) -> None:
        """Connect Ray client to the specified session."""
        session_found = self._get_session_or_die(project_id, session_name)
        session_url, secure, metadata = self._get_connect_params(session_found)
        self._log.debug("Connecting to Ray", session_url, metadata)
        self._ray.util.connect(
            session_url, secure=secure, metadata=metadata, connection_retries=3
        )
        self._log.debug("Checking connection...")

        def func() -> str:
            return "Connected!"

        f_remote = self._ray.remote(func)
        self._log.debug(self._ray.get(f_remote.remote()))
        self._log.info(
            "View session at: https://beta.anyscale.com/o/anyscale-internal/"
            "projects/{}/sessions/{}".format(project_id, session_found.id)
        )

    def _get_session_or_die(self, project_id: str, session_name: str) -> Session:
        """Query Anyscale for the given session's metadata."""
        query = self._anyscale_sdk.search_sessions(project_id=project_id)
        session_found: Session = None
        for session in query.results:
            if session.name == session_name:
                session_found = session
                break
        if not session_found:
            raise RuntimeError("Failed to locate session: {}".format(session_name))
        return session_found

    def _get_connect_params(self, session_meta: Session) -> Tuple[str, bool, Any]:
        """Get the params from the session needed to use Ray client."""
        full_url = session_meta.service_proxy_url
        # like "session-fqsx0p3pzfna71xxxxxxx.anyscaleuserdata.com:8081"
        session_url = full_url.split("/")[2].lower() + ":8081"
        # like "8218b528-7363-4d04-8358-57936cdxxxxx"
        auth_token = full_url.split("?token=")[1].split("&")[0]
        metadata = [("cookie", "anyscale-token=" + auth_token), ("port", "10001")]
        return session_url, False, metadata

    def _fingerprint(self, dir_path: str) -> str:
        """Calculate file hash for the given dir."""
        fingerprint_hasher = hashlib.blake2b()
        session_default_yaml = os.path.join(dir_path, "session-default.yaml")
        dir_fingerprint = fingerprint(
            dir_path,
            exclude_dirs=[".git", "__pycache__"],
            exclude_paths=[
                # Ignore the current file, it's not a necessary dependency.
                os.path.abspath(sys.argv[0]),
                # Ignore the session YAMLs since we're updating it.
                session_default_yaml,
                os.path.join(dir_path, ".anyscale.yaml"),
            ],
            mtime_hash=True,
        )
        fingerprint_hasher.update(dir_fingerprint.encode("utf-8"))
        if os.path.exists(session_default_yaml):
            with open(session_default_yaml) as f:
                fingerprint_hasher.update(f.read().encode("utf-8"))
        return fingerprint_hasher.hexdigest()
