import os
import subprocess
import urllib.parse
from datetime import datetime
from pathlib import Path
from typing import Any, Dict, List, Optional

import click
import humanfriendly
import inquirer
import paramiko

from vessl.openapi_client import CliWorkspaceBackupCreateAPIInput
from vessl.openapi_client.models import (
    InfluxdbWorkloadLog,
    OrmVolumeMountRequests,
    OrmWorkspacePort,
    ResponseWorkspaceDetail,
    ResponseWorkspaceList,
    WorkspaceCreateAPIInput,
)
from vessl import vessl_api
from vessl.util.echo import print_warning, print_info
from vessl.kernel_cluster import list_cluster_nodes, read_cluster
from vessl.kernel_resource_spec import (
    _configure_custom_kernel_resource_spec,
    read_kernel_resource_spec,
)
from vessl.organization import _get_organization_name
from vessl.util.common import parse_time_to_ago
from vessl.util.config import VesslConfigLoader
from vessl.util.constant import (
    SSH_CONFIG_FORMAT,
    SSH_CONFIG_PATH,
    TEMP_DIR,
    WORKSPACE_BACKUP_MAX_SIZE,
    WORKSPACE_BACKUP_MAX_SIZE_FORMATTED,
)
from vessl.util.exception import InvalidWorkspaceError, VesslException
from vessl.util.random import random_string
from vessl.util.ssh import ssh_command_from_endpoint
from vessl.util.tar import Tar
from vessl.util.zipper import Zipper
from vessl.volume import (
    _configure_volume_mount_request_datasets,
    _configure_volume_mount_request_local_files,
    copy_volume_file,
)


def read_workspace(workspace_id: int, **kwargs) -> ResponseWorkspaceDetail:
    """Read workspace in the default organization. If you want to override the
    default organization, then pass `organization_name` as `**kwargs`.

    Args:
        workspace_id(int): Workspace ID.

    Example:
        ```python
        vessl.read_workspace(
            workspace_id=123456,
        )
        ```
    """
    return vessl_api.workspace_read_api(
        workspace_id=workspace_id, organization_name=_get_organization_name(**kwargs)
    )


def list_workspaces(
    cluster_id: int = None, statuses: List[str] = None, mine: bool = True, limit: Optional[int] = None, **kwargs
) -> List[ResponseWorkspaceList]:
    """List workspaces in the default organization. If you want to override the
    default organization, then pass `organization_name` as `**kwargs`.

    Args:
        cluster_id(int): Defaults to None.
        statuses(List[str]): A list of status filter. Defaults to None.
        mine(bool): True if list only my workspaces, False otherwise. Defaults
            to True.

    Example:
        ```python
        vessl.list_workspaces(
            cluster_id=123456,
            statuses=["running"],
        )
        ```
    """
    statuses = (
        [",".join(statuses)] if statuses else None
    )  # since openapi-generator uses repeating params instead of commas
    return vessl_api.workspace_list_api(
        organization_name=_get_organization_name(**kwargs),
        cluster=cluster_id,
        mine=mine,
        statuses=statuses,
        limit=limit,
    ).results


def create_workspace(
    name: str,
    cluster_name: str,
    cluster_node_names: List[str] = None,
    kernel_resource_spec_name: str = None,
    processor_type: str = None,
    cpu_limit: float = None,
    memory_limit: str = None,
    gpu_type: str = None,
    gpu_limit: int = None,
    kernel_image_url: str = None,
    max_hours: int = 24,
    dataset_mounts: List[str] = None,
    local_files: List[str] = None,
    use_vesslignore: bool = True,
    root_volume_size: str = "100Gi",
    ports: List[Dict[str, Any]] = None,
    init_script: str = None,
    **kwargs,
) -> ResponseWorkspaceDetail:
    """Create workspace in the default organization. If you want to override the
    default organization, then pass `organization_name` as `**kwargs`.

    Args:
        name(str): Workspace name.
        cluster_name(str): Cluster name(must be specified before other options).
        cluster_node_names(List[str]): A list of candidate cluster node names.
            Defaults to None.
        kernel_resource_spec_name(str):  Resource type to run an experiment (for
            managed cluster only). Defaults to None.
        processor_type(str) cpu or gpu (for custom cluster only). Defaults to
            None.
        cpu_limit(float): Number of vCPUs (for custom cluster only). Defaults to
            None.
        memory_limit(str): Memory limit in GiB (for custom cluster only).
            Defaults to None.
        gpu_type(str): GPU type (for custom cluster only). Defaults to None.
        gpu_limit(int): Number of GPU cores (for custom cluster only). Defaults
            to None.
        kernel_image_url(str): Kernel docker image URL. Defaults to None.
        max_hours(int): Max hours limit to run. Defaults to 24.
        dataset_mounts(List[str]): A list of dataset mounts. Defaults to None.
        local_files(List[str]): A list of local file mounts. Defaults to None.
        use_vesslignore(bool): True if local files matching glob patterns
            in .vesslignore files should be ignored. Patterns apply relative to
            the directory containing that .vesslignore file.
        root_volume_size(str): Root volume size. Defaults to "100Gi".
        ports(List[Dict[str, Any]]): Port numbers to expose. Defaults to None.
        init_script(str) Custom init script. Defaults to None.

    Example:
        ```python
        vessl.create_workspace(
            name="modern-kick",
            cluster_name="vessl-oci-sanjose",
            kernel_resource_spec_name="cpu-medium",
            kernel_image_url="public.ecr.aws/vessl/kernels:py36.full-cpu.jupyter",
        )
        ```
    """
    cluster = read_cluster(cluster_name)
    cluster_node_ids = (
        [n.id for n in list_cluster_nodes(cluster.id, **kwargs) if n.name in cluster_node_names]
        if cluster_node_names is not None
        else None
    )

    kernel_resource_spec = kernel_resource_spec_id = None
    if kernel_resource_spec_name:
        kernel_resource_spec_id = read_kernel_resource_spec(
            cluster.id,
            kernel_resource_spec_name,
        ).id
    else:
        kernel_resource_spec = _configure_custom_kernel_resource_spec(
            processor_type,
            cpu_limit,
            memory_limit,
            gpu_type,
            gpu_limit,
        )

    all_ports = [
        OrmWorkspacePort(expose_type="http", port=8888, name="jupyter"),
        OrmWorkspacePort(expose_type="tcp", port=22, name="ssh"),
    ]
    if ports is not None:
        for p in ports:
            all_ports.append(
                OrmWorkspacePort(
                    expose_type=p.get("expose_type"),
                    port=p.get("port"),
                    name=p.get("name"),
                )
            )

    requests = []
    if dataset_mounts is not None:
        requests.extend(
            _configure_volume_mount_request_datasets(
                dataset_mounts, should_add_project_datasets=False, **kwargs
            )
        )
    if local_files:
        requests.extend(
            _configure_volume_mount_request_local_files(local_files, use_vesslignore, **kwargs)
        )

    volume_mount_requests = OrmVolumeMountRequests(
        root_volume_size=root_volume_size,
        requests=requests,
    )

    return vessl_api.workspace_create_api(
        organization_name=_get_organization_name(**kwargs),
        workspace_create_api_input=WorkspaceCreateAPIInput(
            cluster_id=cluster.id,
            cluster_node_ids=cluster_node_ids,
            resource_spec=kernel_resource_spec,
            resource_spec_id=kernel_resource_spec_id,
            image_url=kernel_image_url,
            name=name,
            max_hours=max_hours,
            ports=all_ports,
            volumes=volume_mount_requests,
            init_script=init_script,
        ),
    )


def delete_workspace(
    workspace_id: int,
    **kwargs
):
    vessl_api.workspace_delete_api(
        workspace_id=workspace_id,
        organization_name=_get_organization_name(**kwargs),
    )

def list_workspace_logs(
    workspace_id: int,
    tail: int = 200,
    **kwargs,
) -> List[InfluxdbWorkloadLog]:
    """List experiment logs in the default organization. If you want to override
    the default organization, then pass `organization_name` as `**kwargs`.

    Args:
        workspace_id(int): Workspace ID.
        tail (int): The number of lines to display from the end. Display all if
            -1. Defaults to 200.

    Example:
        ```python
         vessl.list_workspace_logs(
            workspace_id=123456,
         )
        ```
    """
    if tail == -1:
        tail = None

    return vessl_api.workspace_logs_api(
        workspace_id=workspace_id,
        organization_name=_get_organization_name(**kwargs),
        limit=tail,
    ).logs


def start_workspace(workspace_id: int, **kwargs) -> ResponseWorkspaceDetail:
    """Start the workspace container in the default organization. If you want to
    override the default organization, then pass `organization_name` as
    `**kwargs`.

    Args:
        workspace_id(int): Workspace ID.

    Example:
        ```python
        vessl.start_workspace(
            workspace_id=123456,
        )
        ```
    """
    return vessl_api.workspace_start_api(
        organization_name=_get_organization_name(**kwargs), workspace_id=workspace_id
    )


def stop_workspace(workspace_id: int, **kwargs) -> ResponseWorkspaceDetail:
    """Stop the workspace container in the default organization. If you want to
    override the default organization, then pass `organization_name` as
    `**kwargs`.

    Args:
        workspace_id(int): Workspace ID.

    Example:
        ```python
        vessl.stop_workspace(
            workspace_id=123456,
        )
        ```
    """
    return vessl_api.workspace_stop_api(
        organization_name=_get_organization_name(**kwargs), workspace_id=workspace_id
    )


def terminate_workspace(workspace_id: int, **kwargs) -> ResponseWorkspaceDetail:
    """Terminate the workspace container in the default organization. If you
    want to override the default organization, then pass `organization_name` as
    `**kwargs`.

    Args:
        workspace_id(int): Workspace ID.

    Example:
        ```python
        vessl.terminate_workspace(
            workspace_id=123456,
        )
        ```
    """
    return vessl_api.workspace_terminate_api(
        organization_name=_get_organization_name(**kwargs), workspace_id=workspace_id
    )


def backup_workspace() -> None:
    """Backup the home directory of the workspace. This command should be called
    inside a workspace.

    Example:
        ```python
        vessl.backup_workspace()
        ```
    """
    workspace_id = VesslConfigLoader().workspace
    if workspace_id is None:
        raise InvalidWorkspaceError("Can only be called within a workspace.")

    workspace = read_workspace(workspace_id)

    filename = (
        f'workspace-backup-{datetime.utcnow().strftime("%Y%m%d%H%M%S")}-{random_string()}.tar.gz'
    )
    gzip_file_full_path = os.path.join(TEMP_DIR, filename)
    home_dir = str(Path.home())

    home_dir_size = -1
    try:
        # use -k option because there is no -b option on UNIX du
        home_dir_size_in_kilobytes = int(
            subprocess.check_output(["du", "-sk", home_dir]).decode().split("\t")[0]
        )
        home_dir_size = home_dir_size_in_kilobytes * 1024
    except Exception:
        pass

    if home_dir_size > WORKSPACE_BACKUP_MAX_SIZE:
        formatted_size = humanfriendly.format_size(home_dir_size, binary=True)
        raise VesslException(
            f"Failed: The size of {home_dir} should be less than {WORKSPACE_BACKUP_MAX_SIZE_FORMATTED} to backup. "
            f"Current size is {formatted_size}.",
        )

    print("Creating a backup file...")
    size = Tar.gzip(gzip_file_full_path, home_dir, exclude_paths=[".vessl"])
    if size > WORKSPACE_BACKUP_MAX_SIZE:
        formatted_size = humanfriendly.format_size(size, binary=True)
        raise VesslException(
            f"Failed: The size of {home_dir} should be less than {WORKSPACE_BACKUP_MAX_SIZE_FORMATTED} to backup. "
            f"Current size is {formatted_size}.",
        )

    print(f"Backup file size: {humanfriendly.format_size(size, binary=True)}")
    print("Uploading the backup file...")
    copy_volume_file(
        source_volume_id=None,
        source_path=gzip_file_full_path,
        dest_volume_id=workspace.backup_volume_id,
        dest_path=filename,
    )
    os.remove(gzip_file_full_path)

    vessl_api.cli_workspace_backup_create_api(
        workspace_id=workspace_id,
        cli_workspace_backup_create_api_input=CliWorkspaceBackupCreateAPIInput(
            filename=filename,
        ),
    )


def restore_workspace() -> None:
    """Restore the home directory from the previous backup. This command should
    be called inside a workspace.

    Example:
        ```python
        vessl.restore_workspace()
        ```
    """
    workspace_id = VesslConfigLoader().workspace
    if workspace_id is None:
        raise InvalidWorkspaceError("Can only be called within a workspace.")

    workspace = read_workspace(workspace_id=workspace_id)
    if workspace.last_backup is None:
        raise click.ClickException("This workspace does not have any backup.")

    if workspace.last_backup.filename.endswith(".zip"):
        dest_path = os.path.join(TEMP_DIR, "workspace-backup.zip")
        print("Downloading the backup file...")
        copy_volume_file(
            source_volume_id=workspace.backup_volume_id,
            source_path=workspace.last_backup.filename,
            dest_volume_id=None,
            dest_path=dest_path,
        )

        zipper = Zipper(dest_path, "r")
        size = zipper.size()
        print(f"Backup file size: {size}")
        extract_path = str(Path.home())
        if size > WORKSPACE_BACKUP_MAX_SIZE:
            extract_path = os.path.join(TEMP_DIR, "workspace-backup")
        print("Extracting...")
        zipper.extractall(extract_path)
        if size > WORKSPACE_BACKUP_MAX_SIZE:
            print(
                f"Restored to {extract_path} "
                f"since the size of the backup file is larger than {WORKSPACE_BACKUP_MAX_SIZE_FORMATTED}. "
            )
        zipper.close()
        zipper.remove()
    else:
        gzip_file_full_path = os.path.join(TEMP_DIR, "workspace-backup.tar.gz")
        print("Downloading the backup file...")
        copy_volume_file(
            source_volume_id=workspace.backup_volume_id,
            source_path=workspace.last_backup.filename,
            dest_volume_id=None,
            dest_path=gzip_file_full_path,
        )

        print("Extracting...")
        Tar.extract(gzip_file_full_path, str(Path.home()))
        os.remove(gzip_file_full_path)


def connect_workspace_ssh(private_key_path: str) -> None:
    """Connect to a running workspace via SSH.

    Args:
        private_key_path(str): SSH private key path

    Example:
        ```python
        vessl.connect_workspace_ssh(
            private_key_path="~/.ssh/key_path",
        )
        ```
    """
    running_workspaces = list_workspaces(statuses=["running"], mine=True)
    if len(running_workspaces) == 0:
        raise click.ClickException("There is no running workspace.")

    if len(running_workspaces) == 1:
        workspace = running_workspaces[0]
    else:
        workspace = inquirer.prompt(
            [
                inquirer.List(
                    "question",
                    message="Select workspace",
                    choices=[
                        (f"{w.name} (created {parse_time_to_ago(w.created_dt)})", w)
                        for w in running_workspaces
                    ],
                )
            ],
            raise_keyboard_interrupt=True,
        ).get("question")

    cmd = ssh_command_from_endpoint(workspace.endpoints.ssh.endpoint, private_key_path)
    print_info(cmd)
    ret = os.system(cmd)
    if ret != 0:
        print_warning("Try again with the correct private key path with --key-path option.")


def update_vscode_remote_ssh(private_key_path: str) -> None:
    """Update .ssh/config file for VSCode Remote-SSH plugin.

    Args:
        private_key_path(str): SSH private key path

    Example:
        ```python
        vessl.update_vscode_remote_ssh(
            private_key_path="~/.ssh/key_path",
        )
        ```
    """
    running_workspaces = list_workspaces(statuses=["running"], mine=True)
    if len(running_workspaces) == 0:
        raise click.ClickException("There is no running workspace.")

    ssh_config = paramiko.SSHConfig()
    try:
        with open(SSH_CONFIG_PATH, "r", encoding="utf-8") as fr:
            ssh_config.parse(fr)
    except FileNotFoundError:
        pass

    host_set = ssh_config.get_hostnames()
    for workspace in running_workspaces:
        host = f"{workspace.name}-{int(workspace.created_dt.timestamp())}"
        if host in host_set:
            continue

        ssh_endpoint = urllib.parse.urlparse(workspace.endpoints.ssh.endpoint)

        config_value = SSH_CONFIG_FORMAT.format(
            host=host,
            hostname=ssh_endpoint.hostname,
            port=ssh_endpoint.port,
        )
        if private_key_path:
            config_value += f"    IdentityFile {private_key_path}\n"

        with open(SSH_CONFIG_PATH, "a", encoding="utf-8") as f:
            f.write(config_value)
