from functools import partial
import os
import subprocess
import tempfile
from typing import cast, Optional, Tuple, Union

from anyscale._private.workload.workload_sdk import WorkloadSDK
from anyscale.client.openapi_client.models import (
    CreateExperimentalWorkspace,
    ExperimentalWorkspace,
)
from anyscale.client.openapi_client.models.session_state import SessionState
from anyscale.utils.runtime_env import parse_requirements_file
from anyscale.workspace.models import WorkspaceConfig, WorkspaceState


SSH_TEMPLATE = """
Host Head-Node
  HostName {head_node_ip}
  User ubuntu
  IdentityFile {key_path}
  StrictHostKeyChecking false
  IdentitiesOnly yes

Host {name}
  HostName 0.0.0.0
  ProxyJump Head-Node
  Port 5020
  User ray
  IdentityFile {key_path}
  StrictHostKeyChecking false
  IdentitiesOnly yes
"""

ANYSCALE_WORKSPACES_SSH_OPTIONS = [
    "-o",
    "StrictHostKeyChecking=no",
    "-o",
    "UserKnownHostsFile=/dev/null",
    "-o",
    "IdentitiesOnly=yes",
]


class PrivateWorkspaceSDK(WorkloadSDK):
    _POLLING_INTERVAL_SECONDS = 10.0
    _WAIT_TIMEOUT_SECONDS = 1800.0

    _BACKEND_SESSION_STATE_TO_WORKSPACE_STATE = {
        SessionState.STOPPED: WorkspaceState.TERMINATED,
        SessionState.TERMINATED: WorkspaceState.TERMINATED,
        SessionState.STARTINGUP: WorkspaceState.STARTING,
        SessionState.STARTUPERRORED: WorkspaceState.ERRORED,
        SessionState.RUNNING: WorkspaceState.RUNNING,
        SessionState.UPDATING: WorkspaceState.UPDATING,
        SessionState.UPDATINGERRORED: WorkspaceState.ERRORED,
        SessionState.STOPPING: WorkspaceState.TERMINATING,
        SessionState.TERMINATING: WorkspaceState.TERMINATING,
        SessionState.AWAITINGSTARTUP: WorkspaceState.STARTING,
        SessionState.AWAITINGFILEMOUNTS: WorkspaceState.STARTING,
        SessionState.TERMINATINGERRORED: WorkspaceState.ERRORED,
        SessionState.STOPPINGERRORED: WorkspaceState.ERRORED,
    }

    def _resolve_to_workspace_model(
        self,
        *,
        id: Optional[str] = None,  # noqa: A002
        name: Optional[str] = None,
        cloud: Optional[str] = None,
        project: Optional[str] = None,
    ) -> ExperimentalWorkspace:
        if name is None and id is None:
            raise ValueError("One of 'name' or 'id' must be provided.")

        if name is not None and id is not None:
            raise ValueError("Only one of 'name' or 'id' can be provided.")

        if id is not None and (cloud is not None or project is not None):
            raise ValueError("'cloud' and 'project' should only be used with 'name'.")

        model: Optional[ExperimentalWorkspace] = self.client.get_workspace(
            name=name, id=id, cloud=cloud, project=project
        )
        if model is None:
            if name is not None:
                raise ValueError(f"Workspace with name '{name}' was not found.")
            else:
                raise ValueError(f"Workspace with ID '{id}' was not found.")

        return model

    def create(self, config: WorkspaceConfig) -> str:
        if not config or config.name is None:
            raise ValueError("Workspace name must be configured")

        name = config.name

        compute_config_id, cloud_id = self.resolve_compute_config_and_cloud_id(
            compute_config=config.compute_config, cloud=config.cloud,
        )

        project_id = self.client.get_project_id(
            parent_cloud_id=cloud_id, name=config.project
        )

        build_id = None

        if config.containerfile is not None:
            build_id = self._image_sdk.build_image_from_containerfile(
                name=f"image-for-workspace-{name}",
                containerfile=self.get_containerfile_contents(config.containerfile),
                ray_version=config.ray_version,
            )
        elif config.image_uri is not None:
            build_id = self._image_sdk.registery_image(
                image_uri=config.image_uri,
                registry_login_secret=config.registry_login_secret,
                ray_version=config.ray_version,
            )

        dynamic_requirements = None
        if (
            config.requirements
            and self._image_sdk.enable_image_build_for_tracked_requirements
        ):
            requirements = (
                parse_requirements_file(config.requirements)
                if isinstance(config.requirements, str)
                else config.requirements
            )
            if requirements:
                build_id = self._image_sdk.build_image_from_requirements(
                    name=f"image-for-workspace-{name}",
                    base_build_id=self.client.get_default_build_id(),
                    requirements=requirements,
                )
        elif config.requirements:
            dynamic_requirements = (
                parse_requirements_file(config.requirements)
                if isinstance(config.requirements, str)
                else config.requirements
            )

        if build_id is None:
            build_id = self.client.get_default_build_id()

        workspace_id = self.client.create_workspace(
            model=CreateExperimentalWorkspace(
                name=name,
                project_id=project_id,
                compute_config_id=compute_config_id,
                cluster_environment_build_id=build_id,
                idle_timeout_minutes=config.idle_termination_minutes,
                cloud_id=cloud_id,
                skip_start=True,
            )
        )

        self._logger.info(f"Workspace created successfully id: {workspace_id}")

        if dynamic_requirements:
            self.client.update_workspace_dependencies_offline_only(
                workspace_id=workspace_id, requirements=dynamic_requirements
            )
            self._logger.info(f"Applied dynamic requirements to workspace id: {name}")
        if config.env_vars:
            self.client.update_workspace_env_vars_offline_only(
                workspace_id=workspace_id, env_vars=config.env_vars
            )
            self._logger.info(f"Applied environment variables to workspace id: {name}")

        return workspace_id

    def start(
        self,
        *,
        id: Optional[str] = None,  # noqa: A002
        name: Optional[str] = None,
        cloud: Optional[str] = None,
        project: Optional[str] = None,
    ) -> str:
        workspace_model = self._resolve_to_workspace_model(
            id=id, name=name, cloud=cloud, project=project
        )
        self.client.start_workspace(workspace_model.id)
        self.logger.info(f"Starting workspace '{workspace_model.name}'")
        return workspace_model.id

    def terminate(
        self,
        *,
        id: Optional[str] = None,  # noqa: A002
        name: Optional[str] = None,
        cloud: Optional[str] = None,
        project: Optional[str] = None,
    ) -> str:
        workspace_model = self._resolve_to_workspace_model(
            id=id, name=name, cloud=cloud, project=project
        )
        self.client.terminate_workspace(workspace_model.id)
        self.logger.info(f"Terminating workspace '{workspace_model.name}'")
        return workspace_model.id

    def status(
        self,
        *,
        id: Optional[str] = None,  # noqa: A002
        name: Optional[str] = None,
        cloud: Optional[str] = None,
        project: Optional[str] = None,
    ) -> str:
        workspace_model = self._resolve_to_workspace_model(
            id=id, name=name, cloud=cloud, project=project
        )
        status = self._get_workspace_status(workspace_model.id)
        return status

    def wait(  # noqa: PLR0912
        self,
        *,
        id: Optional[str] = None,  # noqa: A002
        name: Optional[str] = None,
        cloud: Optional[str] = None,
        project: Optional[str] = None,
        timeout_s: float = _WAIT_TIMEOUT_SECONDS,
        state: Union[str, WorkspaceState] = WorkspaceState.RUNNING,
        interval_s: float = _POLLING_INTERVAL_SECONDS,
    ):
        if not isinstance(timeout_s, (int, float)):
            raise TypeError("timeout_s must be a float")
        if timeout_s <= 0:
            raise ValueError("timeout_s must be >= 0")

        if not isinstance(interval_s, (int, float)):
            raise TypeError("interval_s must be a float")
        if interval_s <= 0:
            raise ValueError("interval_s must be >= 0")

        if isinstance(state, str):
            try:
                state = WorkspaceState.validate(state)
            except KeyError:
                raise ValueError(f"Invalid state: {state}")

        if not isinstance(state, WorkspaceState):
            raise TypeError("'state' must be a WorkspaceState.")

        workspace_model = self._resolve_to_workspace_model(
            id=id, name=name, cloud=cloud, project=project
        )

        workspace_id_or_name = workspace_model.id or workspace_model.name
        curr_state = self._get_workspace_status(workspace_model.id)

        self.logger.info(
            f"Waiting for workspace '{workspace_id_or_name}' to reach target state {state}, currently in state: {curr_state}"
        )
        for _ in self.timer.poll(timeout_s=timeout_s, interval_s=interval_s):
            new_state = self._get_workspace_status(workspace_model.id)

            if new_state != curr_state:
                self.logger.info(
                    f"Workspace '{workspace_id_or_name}' transitioned from {curr_state} to {new_state}"
                )
                curr_state = new_state

            if curr_state == state:
                self.logger.info(
                    f"Workspace '{workspace_id_or_name}' reached target state, exiting"
                )
                break
        else:
            raise TimeoutError(
                f"Workspace '{workspace_id_or_name}' did not reach target state {state} within {timeout_s}s. Last seen state: {curr_state}."
            )

    def generate_ssh_config_file(
        self,
        *,
        id: Optional[str] = None,  # noqa: A002
        name: Optional[str] = None,
        cloud: Optional[str] = None,
        project: Optional[str] = None,
        ssh_config_path: Optional[str] = None,
    ) -> Tuple[str, str]:
        workspace_model = self._resolve_to_workspace_model(
            id=id, name=name, cloud=cloud, project=project
        )

        assert (
            self.status(id=workspace_model.id) == WorkspaceState.RUNNING
        ), "Workspace must be running to generate SSH config file"

        head_node_ip = self.client.get_cluster_head_node_ip(workspace_model.cluster_id)
        ssh_key = self.client.get_cluster_ssh_key(workspace_model.cluster_id)

        def _store_ssh_key(name: str, ssh_key: str, key_dir: str) -> str:
            key_path = os.path.join(key_dir, f"{name}.pem")
            os.makedirs(os.path.dirname(key_path), exist_ok=True)

            with open(key_path, "w", opener=partial(os.open, mode=0o600)) as f:
                f.write(ssh_key)

            return key_path

        if ssh_config_path is None:
            ssh_config_path = tempfile.mkdtemp()

        key_path = _store_ssh_key(
            ssh_key.key_name, ssh_key.private_key, ssh_config_path
        )

        ssh_config = SSH_TEMPLATE.format(
            head_node_ip=head_node_ip, key_path=key_path, name=workspace_model.name
        )

        config_file_name = os.path.join(ssh_config_path, "config")
        with open(config_file_name, "w") as f:
            f.write(ssh_config)

        return workspace_model.name, config_file_name

    def run_command(
        self,
        command: str,
        *,
        id: Optional[str] = None,  # noqa: A002
        name: Optional[str] = None,
        cloud: Optional[str] = None,
        project: Optional[str] = None,
        **kwargs,
    ) -> subprocess.CompletedProcess:
        workspace_model = self._resolve_to_workspace_model(
            id=id, name=name, cloud=cloud, project=project
        )
        host_name, config_file = self.generate_ssh_config_file(id=workspace_model.id)
        return subprocess.run(
            ["ssh"]
            + ANYSCALE_WORKSPACES_SSH_OPTIONS
            + ["-F", config_file, host_name, command],
            **kwargs,
        )

    def get_default_dir_name(
        self,
        *,
        id: Optional[str] = None,  # noqa: A002
        name: Optional[str] = None,
        cloud: Optional[str] = None,
        project: Optional[str] = None,
    ) -> str:
        workspace_model = self._resolve_to_workspace_model(
            id=id, name=name, cloud=cloud, project=project
        )
        return self.client.get_workspace_default_dir_name(workspace_model.id)

    def _get_workspace_status(self, workspace_id: str) -> WorkspaceState:
        workspace_session_state = self.client.get_workspace_status(workspace_id)
        return cast(
            WorkspaceState,
            self._BACKEND_SESSION_STATE_TO_WORKSPACE_STATE.get(
                workspace_session_state, WorkspaceState.UNKNOWN
            ),
        )
