from __future__ import annotations

import base64
import hashlib
import os
import re
import shutil
import subprocess
import webbrowser
from typing import Optional, Sequence

import click
import importlib_metadata
from rich import print
from urllib3.util import parse_url

import coiled
from coiled.cli.cluster.ssh import add_key_to_agent, check_ssh
from coiled.compatibility import DISTRIBUTED_VERSION
from coiled.errors import DoesNotExist

from ..utils import CONTEXT_SETTINGS

# Path on VM to sync to.
# We use `/scratch` for now because it's already bind-mounted into docker.
SYNC_TARGET = "/scratch/synced"
MIN_DISTRIBUTED_VERSION = "2022.8.1"
MUTAGEN_NAME_FORMAT = "coiled-{cluster_id}"


def check_distributed_version() -> bool:
    if DISTRIBUTED_VERSION < MIN_DISTRIBUTED_VERSION:
        print(
            "[bold red]"
            f"distributed>{MIN_DISTRIBUTED_VERSION} is required to launch notebooks. "
            f"You have {DISTRIBUTED_VERSION}."
            "[/]"
        )
        return False
    return True


def check_jupyter() -> bool:
    try:
        importlib_metadata.distribution("jupyter_server")
    except ModuleNotFoundError:
        print("[bold red]Jupyter must be installed locally to launch notebooks.[/]")
        return False
    return True


def check_mutagen() -> bool:
    if not shutil.which("mutagen"):
        print(
            "[bold red]"
            "mutagen must be installed to synchronize files with notebooks.[/]\n"
            "Install via homebrew (on macOS, Linux, or Windows) with:\n\n"
            "brew install mutagen-io/mutagen/mutagen@0.16\n\n"
            "Or, visit https://github.com/mutagen-io/mutagen/releases/latest to download "
            "a static, pre-compiled binary for your system, and place it anywhere on your $PATH."
        )
        return False
    return True


def check_ssh_keygen() -> bool:
    if not shutil.which("ssh-keygen"):
        print("[bold red]Unable to find `ssh-keygen`, you may need to install OpenSSH or add it to your paths.[/]")
        return False
    return True


def cwd_cluster_name():
    "Get a name for the notebook based on the hash of the current path."
    # Based on https://github.com/python-poetry/poetry/blob/fa5543a6/src/poetry/utils/env.py#L1210-L1212
    cwd = os.getcwd()
    normalized_cwd = os.path.normcase(os.path.realpath(cwd))
    h_bytes = hashlib.sha256(normalized_cwd.encode()).digest()
    h_str = base64.urlsafe_b64encode(h_bytes).decode()[:8]

    basename = os.path.basename(normalized_cwd)
    sanitized_name = re.sub(r'[ $`!*@"/\\\r\n\t\.]', "_", basename)[:42]
    return f"notebook-{sanitized_name}-{h_str}"


def mutagen_session_exists(cluster_id: int) -> bool:
    sessions = (
        subprocess.run(
            [
                "mutagen",
                "sync",
                "list",
                "--label-selector",
                f"managed-by=coiled,cluster-id={cluster_id}",
                "--template",
                "{{range .}}{{.Name}}{{end}}",
            ],
            check=True,
            text=True,
            capture_output=True,
        )
        .stdout.strip()
        .splitlines()
    )

    if not sessions:
        return False
    if sessions == [MUTAGEN_NAME_FORMAT.format(cluster_id=cluster_id)]:
        return True

    if len(sessions) == 1:
        raise RuntimeError(
            f"Unexpected mutagen session name {sessions[0]!r}. "
            f"Expected {MUTAGEN_NAME_FORMAT.format(cluster_id=cluster_id)!r}."
        )

    raise RuntimeError(f"Multiple mutagen sessions found for cluster {cluster_id}: {sessions}")


@click.command(context_settings=CONTEXT_SETTINGS)
@click.option(
    "--name",
    default=cwd_cluster_name(),
    help="Cluster name. If not given, defaults to a hash based on current working directory.",
)
@click.option(
    "--sync",
    default=False,
    is_flag=True,
    help="Sync the working directory with the filesystem on the notebook. Requires mutagen.",
)
@click.option(
    "--software",
    default=None,
    help=(
        "Software environment name to use. If not given, all the currently-installed "
        "Python packages are replicated on the VM using package sync."
    ),
)
@click.option(
    "--vm-type",
    default=[],
    multiple=True,
    help="VM type to use. Specify multiple times to provide multiple options.",
)
@click.option(
    "--open",
    default=True,
    is_flag=True,
    help="Whether to open the notebook in the default browser once it's launched",
)
def start_notebook(name: str, sync: bool, software: Optional[str], vm_type: Sequence[str], open: bool):
    """
    Launch or re-open a notebook session, with optional file syncing.

    If a notebook session with the same ``name`` already exists, it's not re-created.
    If file sync was initially not enabled, running ``coiled notebook-beta up --sync``
    will begin file sync without re-launching the notebook.
    """
    if not (check_distributed_version() and check_jupyter()):
        return

    if sync and not (check_mutagen() and check_ssh() and check_ssh_keygen()):
        return

    with coiled.Cloud() as cloud:
        # TODO how can we get the widget to show up during CLI commands?
        cluster = coiled.Cluster(
            name=name,
            cloud=cloud,
            n_workers=0,
            software=software,
            jupyter=True,
            scheduler_options={"idle_timeout": None},
            scheduler_vm_types=list(vm_type) if vm_type else None,
            allow_ssh=True,
        )

        url = cluster.jupyter_link
        cluster_id = cluster.cluster_id
        assert cluster_id is not None
        if sync:
            url = parse_url(url)._replace(path="/jupyter/lab/tree/synced").url

            if mutagen_session_exists(cluster_id):
                print("[bold]File sync session already active; reusing it.[/]")
            else:
                print("[bold]Launching file synchronization...[/]")
                ssh_info = cloud.get_ssh_key(cluster_id)

                scheduler_address = ssh_info["scheduler_public_address"]
                target = f"ubuntu@{scheduler_address}"

                add_key_to_agent(scheduler_address, key=ssh_info["private_key"])

                # Update known_hosts. We can't specify SSH options to mutagen so we can't pass
                # `-o StrictHostKeyChecking=no`. Could alternatively add an entry in `~/.ssh/config`,
                # but that feels more intrusive.
                # TODO get public key from Coiled
                subprocess.run(
                    f"ssh-keyscan {scheduler_address} >> ~/.ssh/known_hosts",
                    shell=True,
                    check=True,
                    capture_output=True,
                )

                # Start mutagen
                subprocess.run(
                    [
                        "mutagen",
                        "sync",
                        "create",
                        "--name",
                        MUTAGEN_NAME_FORMAT.format(cluster_id=cluster_id),
                        "--label",
                        "managed-by=coiled",
                        "--label",
                        f"cluster-id={cluster_id}",
                        "--ignore-vcs",
                        "--max-staging-file-size=1 GiB",
                        ".",
                        f"{target}:{SYNC_TARGET}",
                    ],
                    check=True,
                )

                # Within the docker container, symlink the sync directory (`/scratch/sync`)
                # into the working directory for Jupyter, so you can actually see the synced
                # files in the Jupyter browser. We use a symlink since the container doesn't
                # have capabilities to make a bind mount.
                # TODO if we don't like the symlink, Coiled could see what the workdir is for
                # the image before running, and bind-mount `/sync` on the host to `$workdir/sync`
                # in the container? Custom docker images make this tricky; we can't assume anything
                # about the directory layout or what the working directory will be.
                subprocess.run(
                    [
                        "ssh",
                        target,
                        f"docker exec tmp-dask-1 bash -c 'mkdir -p {SYNC_TARGET} && ln -s {SYNC_TARGET} .'",
                    ],
                    check=True,
                )

    print(f"[bold]Jupyter available at {url}[/]")
    if open:
        webbrowser.open(url, new=2)


@click.command(context_settings=CONTEXT_SETTINGS)
@click.option(
    "--name",
    default=cwd_cluster_name(),
    help=("Cluster name. If not given, defaults to a hash based on current working directory."),
)
def stop_notebook(name: str):
    "Shut down a notebook session."
    with coiled.Cloud() as cloud:
        try:
            cluster_id = cloud.get_cluster_by_name(name)
        except DoesNotExist:
            print(f"[bold red]Cluster {name!r} does not exist[/]")
            return  # TODO exit 1

        if shutil.which("mutagen") and mutagen_session_exists(cluster_id):
            # NOTE: we can't tell if the user asked for `--sync` or not at creation.
            # Best we can do is check if mutagen is installed and the session exists.
            if not (check_ssh() and check_ssh_keygen()):
                return

            # Stop mutagen
            print(f"Stopping sync with cluster {name!r} ({cluster_id})")
            subprocess.run(
                ["mutagen", "sync", "terminate", MUTAGEN_NAME_FORMAT.format(cluster_id=cluster_id)],
                check=True,
            )

            ssh_info = cloud.get_ssh_key(cluster_id)
            scheduler_address = ssh_info["scheduler_public_address"]
            add_key_to_agent(scheduler_address, key=ssh_info["private_key"], delete=True)

            # Remove `known_hosts` entries.
            # TODO don't like touching the user's `known_hosts` file like this.
            subprocess.run(
                [
                    "ssh-keygen",
                    "-f",
                    os.path.expanduser("~/.ssh/known_hosts"),
                    "-R",
                    scheduler_address,
                ]
            )

        print(f"Stopping cluster {name!r} ({cluster_id})...")
        cloud.delete_cluster(cluster_id)


@click.command(context_settings=CONTEXT_SETTINGS)
@click.option(
    "--name",
    default=cwd_cluster_name(),
    help=("Cluster name. If not given, defaults to a hash based on current working directory."),
)
def monitor_sync(name: str):
    "Monitor file sync status for a notebook session."
    if not check_mutagen():
        return

    with coiled.Cloud() as cloud:
        try:
            cluster_id = cloud.get_cluster_by_name(name)
        except DoesNotExist:
            print(f"[bold red]Cluster {name!r} does not exist[/]")
            return  # TODO exit 1

    if not mutagen_session_exists(cluster_id):
        print(f"[bold red]No file synchronization session for cluster {name!r} ({cluster_id})[/]")
        return  # TODO exit 1

    subprocess.run(["mutagen", "sync", "monitor", MUTAGEN_NAME_FORMAT.format(cluster_id=cluster_id)])
