# Copyright 2017-2023 Posit Software, PBC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import logging
import os
import pprint
import shutil
import subprocess
import sys
import zipfile

from guild import remote as remotelib
from guild import remote_util
from guild import util

from . import meta_sync

log = logging.getLogger("guild.remotes.gist")


class NoSuchGist(remotelib.OperationError):
    pass


class MissingRequiredEnv(remotelib.OperationError):
    pass


class GistRemoteType(remotelib.RemoteType):
    def __init__(self, _ep):
        pass

    def remote_for_config(self, name, config):
        return GistRemote(name, config)

    def remote_for_spec(self, spec):
        name = f"gist:{spec}"
        user, gist_name = _parse_spec(spec)
        config = remotelib.RemoteConfig({
            "user": user,
            "gist-name": gist_name,
        })
        return GistRemote(name, config)


def _parse_spec(spec):
    parts = spec.split("/", 1)
    if len(parts) == 1:
        try:
            return _required_gist_user_env({}), parts[0]
        except MissingRequiredEnv as e:
            raise remotelib.InvalidRemoteSpec(str(e))
    return parts


def _required_gist_user_env(env):
    try:
        return _required_env("GIST_USER", [env, os.environ])
    except KeyError as e:
        raise MissingRequiredEnv(
            "gist remotes must be specified as USER/GIST_NAME if GIST_USER "
            "environment variable is not defined"
        ) from e


def _required_env(name, sources):
    for src in sources:
        try:
            return src[name]
        except KeyError:
            pass
    raise KeyError(name)


class GistRemote(meta_sync.MetaSyncRemote):
    def __init__(self, name, config):
        self.name = name
        self.user = config["user"]
        self.gist_name = config["gist-name"]
        self._gist_readme_name = _gist_readme_name(self.gist_name)
        self.local_env = remote_util.init_env(config.get("local-env"))
        self.local_sync_dir = meta_sync.local_meta_dir(
            _remote_full_name(self.user, self.gist_name), ""
        )
        self._local_gist_repo = os.path.join(self.local_sync_dir, "gist")
        runs_dir = os.path.join(self.local_sync_dir, "runs")
        super().__init__(runs_dir, None)

    def status(self, verbose=False):
        remote_util.remote_activity("Getting %s status", self.name)
        gist = self._repo_gist()
        sys.stdout.write(f"{self.name} (gist {gist['id']}) is available\n")
        if verbose:
            sys.stdout.write(pprint.pformat(gist))
            sys.stdout.write("\n")

    def start(self):
        remote_util.remote_activity("Getting %s status", self.name)
        try:
            gist = self._repo_gist()
        except NoSuchGist:
            log.info("Creating gist")
            gist = self._create_gist()
            log.info(
                "Created %s (gist %s) for user %s",
                self.name,
                gist["id"],
                self.user,
            )
            self._sync_runs_meta()
        else:
            raise remotelib.OperationError(
                f"{self.name} (gist {gist['id']}) already exists for user {self.user}"
            )

    def stop(self):
        self._delete_gist()
        self._clear_gist_cache()

    def _delete_gist(self):
        gist = self._repo_gist()
        log.info("Deleting gist %s", gist["id"])
        _delete_gist(gist, self.local_env)

    def _clear_gist_cache(self):
        log.info("Clearing local cache")
        log.debug("deleting %s", self.local_sync_dir)
        util.ensure_safe_rmtree(self.local_sync_dir)

    def stop_details(self):
        remote_util.remote_activity("Getting %s status", self.name)
        try:
            gist = self._repo_gist()
        except NoSuchGist:
            return None
        else:
            return f"gist {gist['id']} will be deleted - THIS CANNOT BE UNDONE!"

    def _sync_runs_meta(self, force=False):
        remote_util.remote_activity(f"Refreshing run info for {self.name}")
        self._ensure_local_gist_repo()
        self._sync_runs_meta_for_gist(force)

    def _ensure_local_gist_repo(self):
        if _is_git_repo(self._local_gist_repo):
            log.debug("gist local repo found at %s", self._local_gist_repo)
            return
        log.debug("initializing gist local repo at %s", self._local_gist_repo)
        gist = self._repo_gist()
        _sync_gist_repo(gist, self._local_gist_repo, self.local_env)

    def _repo_gist(self):
        gist = _find_gist_with_file(self.user, self._gist_readme_name, self.local_env)
        if not gist:
            raise NoSuchGist(
                f"cannot find gist remote '{self.gist_name}' (denoted by the "
                f"file '{self._gist_readme_name}') for user {self.user}\n"
                "If the gist is private, you must specify a valid access token with "
                "GIST_ACCESS_TOKEN.\nFor more information see "
                "https://my.guild.ai/docs/gists."
            )
        return gist

    def _sync_runs_meta_for_gist(self, force):
        try:
            _pull_gist_repo(self._local_gist_repo, self.local_env)
        except NoSuchGist:
            self._clear_gist_cache()
        else:
            git_commit = self._gist_repo_current_commit()
            if not force and self._meta_current(git_commit):
                return
            _refresh_runs_meta(
                self._local_gist_repo,
                self._runs_dir,
                git_commit,
                self.local_sync_dir,
            )

    def _meta_current(self, git_commit):
        return meta_sync.meta_current(self.local_sync_dir, lambda: git_commit)

    def _gist_repo_current_commit(self):
        return _git_current_commit(self._local_gist_repo)

    def _delete_runs(self, runs, permanent):
        assert permanent  # gist remotes only support permanent delete
        _delete_gist_runs(runs, self._local_gist_repo, self._runs_dir)
        _commit_and_push_gist_repo_for_delete(
            self._local_gist_repo,
            _delete_commit_msg(),
            self.local_env,
            self.name,
        )

    def _restore_runs(self, runs):
        raise NotImplementedError()

    def _purge_runs(self, runs):
        raise NotImplementedError()

    def push(self, runs, delete=False):
        self._ensure_synced_gist_repo()
        _export_runs_to_gist_archives(runs, self._local_gist_repo)
        _commit_and_push_gist_repo_for_push(
            self._local_gist_repo,
            _push_commit_msg(),
            self.local_env,
            self.name,
        )
        self._sync_runs_meta_for_gist(True)

    def _ensure_synced_gist_repo(self):
        try:
            self._sync_runs_meta()
        except NoSuchGist:
            self._init_gist_repo()

    def _init_gist_repo(self):
        gist = self._create_gist()
        _sync_gist_repo(gist, self._local_gist_repo, self.local_env)

    def _create_gist(self):
        return _create_gist(
            self.user, self.gist_name, self._gist_readme_name, self.local_env
        )

    def pull(self, runs, delete=False):
        from guild import var

        # That we have `runs` means we've sync'd runs meta. "Meta" in
        # this case also contains the runs themselves as zip
        # archives. At this point we need only extract the run
        # archives to the runs dir.
        _extract_runs(runs, self._local_gist_repo, var.runs_dir(), self.name)


def _remote_full_name(user, gist_name):
    return f"gist-{user}-{gist_name}"


def _gist_readme_name(gist_name):
    return f"[Guild AI] {_ensure_md_ext(gist_name)}"


def _ensure_md_ext(s):
    if s.lower().endswith(".md"):
        return s
    return s + ".md"


def _find_gist_with_file(user, filename, env):
    import urllib.request
    import json

    page = 1
    while True:
        url = f"https://api.github.com/users/{ user }/gists?page={ page }&per_page=100"
        req = urllib.request.Request(
            url,
            headers=_github_auth_headers(env),
            method="GET",
        )
        resp = urllib.request.urlopen(req)
        gists = json.loads(resp.read().decode("utf-8"))

        if not gists:
            return None
        for gist in gists:
            for name in gist["files"]:
                if name == filename:
                    return gist
        page += 1


def _github_auth_headers(env):
    try:
        access_token = _required_gist_access_token(env)
    except MissingRequiredEnv:
        return {}
    else:
        return {"Authorization": f"token {access_token}"}


def _sync_gist_repo(gist, local_repo, env):
    repo_url = _gist_repo_url(gist, env)
    if _is_git_repo(local_repo):
        _pull_gist_repo(local_repo, env)
    else:
        _clone_gist_repo(repo_url, local_repo, env)


def _gist_repo_url(gist, env):
    if _gist_urltype(env) == "ssh":
        return f"git@gist.github.com:{gist['id']}.git"
    return gist["git_pull_url"]


def _gist_urltype(env):
    try:
        return _required_env("GIST_URLTYPE", [env, os.environ])
    except KeyError:
        return None


def _clone_gist_repo(repo_url, local_repo, env):
    cmd = [_git_cmd(), "clone", "--quiet", repo_url, local_repo]
    log.debug("cloning %s to %s", repo_url, local_repo)
    _subprocess_tty(cmd, extra_env=env)


def _git_cmd():
    cmd = util.which("git")
    if not cmd:
        raise remotelib.OperationError(
            "git command is not available\n"
            "Refer to https://git-scm.com/book/en/v2/Getting-Started-Installing-Git "
            "for help installing it."
        )
    return cmd


def _pull_gist_repo(local_repo, env):
    cmd = [_git_cmd(), "-C", local_repo, "pull", "--quiet", "--rebase"]
    log.debug("pulling for %s", local_repo)
    code = _subprocess_tty(cmd, extra_env=env, allowed_returncodes=(0, 1))
    if code == 1:
        raise NoSuchGist()


def _refresh_runs_meta(gist_repo, runs_dir, meta_id, local_sync_dir):
    for archive in _run_archives(gist_repo):
        _unpack_meta(archive, runs_dir)
    meta_sync.write_local_meta_id(meta_id, local_sync_dir)


def _run_archives(dir):
    for name in os.listdir(dir):
        if _is_guild_run(name):
            yield os.path.join(dir, name)


def _is_guild_run(name):
    return name.startswith("guildai-run-") and name.endswith(".zip")


def _unpack_meta(archive, runs_dir):
    log.debug("unpacking %s meta to %s", archive, runs_dir)
    with zipfile.ZipFile(archive, "r") as zf:
        for name in zf.namelist():
            if meta_sync.is_meta_file(name):
                zf.extract(name, runs_dir)


def _is_git_repo(dir):
    return os.path.exists(os.path.join(dir, ".git"))


def _git_current_commit(git_repo):
    if not _is_git_repo(git_repo):
        return None
    cmd = [_git_cmd(), "-C", git_repo, "log", "-1", "--format=%H"]
    out = subprocess.check_output(cmd, stderr=subprocess.STDOUT)
    return out.decode("utf-8").strip()


def _extract_runs(runs, archive_dir, dest_dir, gist_name):
    for run in runs:
        archive = os.path.join(archive_dir, _run_archive_filename(run.id))
        if not os.path.exists(archive):
            log.error(
                "%s archive for gist does not exist (%s), skipping", run.id, archive
            )
            continue
        log.info("Copying %s from %s", run.id, gist_name)
        _replace_run(archive, run.id, dest_dir)


def _run_archive_filename(run_id):
    return f"guildai-run-{run_id}.zip"


def _replace_run(archive, run_id, dest_dir):
    with util.TempDir("guild-gist-run-") as tmp:
        _extract_archive(archive, tmp.path)
        extracted_run_dir = _validate_extracted_run(tmp.path, run_id, archive)
        dest_run_dir = os.path.join(dest_dir, run_id)
        _replace_run_dir(dest_run_dir, extracted_run_dir)


def _extract_archive(archive, dest_dir):
    with zipfile.ZipFile(archive, "r") as zf:
        for name in zf.namelist():
            zf.extract(name, dest_dir)


def _validate_extracted_run(dir, run_id, archive):
    # RUN_DIR/.guild/opref is required for a run.
    extracted_run_dir = os.path.join(dir, run_id)
    opref_path = os.path.join(extracted_run_dir, ".guild", "opref")
    if not os.path.exists(opref_path):
        log.error("%s does not contain expected run %s", archive, run_id)
        raise remotelib.OperationError("invalid run archive in gist")
    return extracted_run_dir


def _replace_run_dir(run_dir, src_dir):
    log.debug("moving %s to %s", src_dir, run_dir)
    util.ensure_safe_rmtree(run_dir)
    shutil.move(src_dir, run_dir)


def _create_gist(gist_remote_user, gist_remote_name, gist_readme_name, env):
    import json
    import urllib.request

    access_token = _required_gist_access_token(env)
    content = _gist_readme_content(gist_remote_user, gist_remote_name)
    data = {
        "accept": "application/vnd.github.v3+json",
        "description": "Guild AI Repository",
        "public": True,
        "files": {
            gist_readme_name: {
                "filename": gist_readme_name,
                "type": "text/markdown",
                "language": "Markdown",
                "content": content,
            }
        },
    }
    headers = {
        "Authorization": f"token {access_token}",
    }
    req = urllib.request.Request(
        "https://api.github.com/gists",
        data=json.dumps(data).encode(),
        headers=headers,
        method="POST",
    )
    resp = urllib.request.urlopen(req)
    resp_data = resp.read().decode("utf-8")
    if resp.getcode() not in (200, 201):
        raise remotelib.OperationError(
            f"error creating gist: ({resp.getcode()}) {resp_data}"
        )
    return json.loads(resp_data)


def _required_gist_access_token(env):
    try:
        return _required_env("GIST_ACCESS_TOKEN", [env, os.environ])
    except KeyError as e:
        raise MissingRequiredEnv(
            "missing required environment variable GIST_ACCESS_TOKEN\n"
            "This operation requires a GitHub personal access token for "
            "creating gists.\n"
            "See https://my.guild.ai/docs/gists for more information."
        ) from e


def _gist_readme_content(user, remote_name):
    return (
        "This is a Guild AI runs repository. To access runs, "
        "[install Guild AI](https://guild.ai/install) and run "
        f"`guild pull gist:{user}/{remote_name}`. "
        "For more information about Guild AI Gist based repositories, see "
        "[Guild AI - Gists](https://my.guild.ai/docs/gists)."
    )


def _export_runs_to_gist_archives(runs, gist_repo):
    with util.TempDir("guild-runs-export-") as tmp:
        archives = [_run_export_archive(run, tmp.path) for run in runs]
        _export_runs(zip(runs, archives))
        for archive_src in archives:
            archive_dest = os.path.join(gist_repo, os.path.basename(archive_src))
            util.ensure_deleted(archive_dest)
            shutil.move(archive_src, archive_dest)


def _run_export_archive(run, export_dir):
    return os.path.join(export_dir, _run_archive_filename(run.id))


def _export_runs(runs_with_dest):
    from guild import run_util

    for run, dest in runs_with_dest:
        log.info("Compressing %s", run.id)
        run_util.export_runs([run], dest, copy_resources=False, quiet=True)


def _push_commit_msg():
    import guild

    return (
        f"`guild push` by {util.user()}@{util.hostname()} with "
        f"version {guild.__version__}"
    )


def _commit_and_push_gist_repo_for_push(repo, commit_msg, env, remote_name):
    _git_add_all(repo, env)
    try:
        _git_commit(repo, commit_msg, env)
    except _NoChanges:
        pass
    log.info("Copying runs to %s", remote_name)
    _git_push(repo, env)


def _git_add_all(local_repo, env, update=False):
    cmd = [_git_cmd(), "-C", local_repo, "add", "."]
    if update:
        cmd.append("-u")
    log.debug("adding files for %s", local_repo)
    _subprocess_quiet(cmd, extra_env=env)


class _NoChanges(Exception):
    pass


def _git_commit(local_repo, msg, env):
    cmd = [_git_cmd(), "-C", local_repo, "commit", "-m", msg]
    log.debug("commiting for %s", local_repo)
    result = _subprocess_quiet(cmd, extra_env=env, allowed_returncodes=(0, 1))
    if result == 1:
        raise _NoChanges()


def _git_push(local_repo, env):
    cmd = [_git_cmd(), "-C", local_repo, "push", "--quiet"]
    _maybe_apply_askpass(local_repo, env)
    log.debug("pushing for %s", local_repo)
    _subprocess_tty(cmd, extra_env=env)


def _maybe_apply_askpass(local_repo, env):
    if not _gist_access_token_defined(env):
        return
    askpass_path = _maybe_gist_access_token_script(local_repo)
    if not askpass_path:
        return
    env["GIT_ASKPASS"] = askpass_path


def _gist_access_token_defined(env):
    try:
        _required_env("GIST_ACCESS_TOKEN", [env, os.environ])
    except KeyError:
        return False
    else:
        return True


def _maybe_gist_access_token_script(local_repo):
    if util.get_platform() == "Windows":
        return None
    script_path = _gist_access_token_script(local_repo)
    if os.path.exists(script_path):
        return script_path
    _write_gist_access_token_script(script_path)
    return script_path


def _gist_access_token_script(local_repo):
    return os.path.join(local_repo, ".git", "gist-access-token")


def _write_gist_access_token_script(path):
    with open(path, "w") as f:
        f.write("echo $GIST_ACCESS_TOKEN\n")
    util.make_executable(path)


def _delete_gist(gist, env):
    import urllib.request
    import json

    access_token = _required_gist_access_token(env)
    data = {
        "accept": "application/vnd.github.v3+json",
        "gist_id": gist["id"],
    }
    headers = {
        "Authorization": f"token {access_token}",
    }
    req = urllib.request.Request(
        url=f"https://api.github.com/gists/{ gist['id'] }",
        method="DELETE",
        data=json.dumps(data).encode(),
        headers=headers,
    )
    with urllib.request.urlopen(req) as resp:
        if resp.getcode() not in (200, 204):
            resp_data = json.loads(resp.read().decode("utf-8"))
            raise remotelib.OperationError(
                f"error creating gist: ({resp.getcode()}) {resp_data}"
            )


def _delete_gist_runs(runs, gist_repo, runs_dir):
    for run in runs:
        log.info("Deleting %s", run.id)
        _delete_gist_repo_run_archive(gist_repo, run.id)
        _delete_run(run, runs_dir)


def _delete_gist_repo_run_archive(gist_repo, run_id):
    run_archive = os.path.join(gist_repo, _run_archive_filename(run_id))
    log.debug("deleting %s", run_archive)
    util.ensure_deleted(run_archive)


def _delete_run(run, runs_dir):
    run_dir = os.path.join(runs_dir, run.id)
    log.debug("deleting %s", run_dir)
    util.ensure_safe_rmtree(run_dir)


def _commit_and_push_gist_repo_for_delete(repo, commit_msg, env, remote_name):
    _git_add_all(repo, env, update=True)
    try:
        _git_commit(repo, commit_msg, env)
    except _NoChanges:
        log.info("Nothing to update for %s - gist is up-to-date", remote_name)
    else:
        log.info("Updating runs on %s", remote_name)
        _git_push(repo, env)


def _delete_commit_msg():
    import guild

    return (
        f"`guild runs rm` by {util.user()}@{util.hostname()} "
        f"with version {guild.__version__}"
    )


def _subprocess_tty(cmd, extra_env, allowed_returncodes=(0,)):
    env = dict(os.environ)
    if extra_env:
        env.update(extra_env)
    log.debug("%r", cmd)
    p = subprocess.Popen(cmd, env=env)
    p.wait()
    if p.returncode not in allowed_returncodes:
        log.debug("exit code for %r is %i", cmd, p.returncode)
        raise SystemExit(f"error running {cmd[0]} - see above for details")
    return p.returncode


def _subprocess_quiet(cmd, extra_env, allowed_returncodes=(0,)):
    log.debug("%r", cmd)
    return remote_util.subprocess_call(
        cmd,
        extra_env=extra_env,
        quiet=True,
        allowed_returncodes=allowed_returncodes,
    )
