#!python

# Copyright DataStax, Inc.
#
# 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 argparse
import hashlib
import logging
import os
import os.path
import sys
from base64 import urlsafe_b64encode
from datetime import datetime, tzinfo, timedelta
from functools import partial

try:
    from itertools import ifilterfalse as filterfalse
except ImportError:
    from itertools import filterfalse

try:
    from tempfile import TemporaryDirectory
except ImportError:
    from backports.tempfile import TemporaryDirectory

try:
    from datetime import timezone
    utc = timezone.utc
except ImportError:
    class UTC(tzinfo):
        def utcoffset(self, dt):
            return timedelta(0)
        def tzname(self, dt):
            return "UTC"
        def dst(self, dt):
            return timedelta(0)
    utc = UTC()


import click

from adelphi.anonymize import anonymize_keyspace
from adelphi.cql import export_cql_schema
from adelphi.gemini import export_gemini_schema
from adelphi.gh import build_github, build_origin_repo, build_branch, commit_schemas, build_pull_request
from adelphi.store import build_keyspace_objects
from adelphi.store import with_cluster as _with_cluster

logging.basicConfig(level=logging.INFO)
log = logging.getLogger('adelphi')


# Final version of PyGithub with Python 2.x support includes a bug that appears when users
# are part of a Github organization.  See https://github.com/datastax/adelphi/issues/70
# for more detail
#
# Source below is Repository.create_fork from v1.45 + a fix for the bug referenced above.
# See https://github.com/PyGithub/PyGithub/blob/v1.45/github/Repository.py#L2075-L2090 for the
# base impl from v1.45.
def patch_repo(repo):
    log.info("Adding safe_create_fork function to github.Repository")
    import six
    from github import GithubObject, Repository
    def safe_create_fork(self, organization=GithubObject.NotSet):
        # TODO: I'm not sure we need to support the non-default string case here... maintaining
        # it for now to preserve fidelity with the 1.45 impl
        assert organization is GithubObject.NotSet or isinstance(
            organization, (str, six.text_type)
        ), organization
        post_parameters = {}
        if organization is not GithubObject.NotSet:
            post_parameters["organization"] = organization
        headers, data = self._requester.requestJsonAndCheck(
            "POST", self.url + "/forks", input=post_parameters,
        )
        return Repository.Repository(self._requester, headers, data, completed=True)

    import types
    repo.safe_create_fork = types.MethodType(safe_create_fork, repo)


# Note: click assumes a type of string and a required value of False unless explicitly specified
@click.group()
@click.option('--hosts', default='127.0.0.1', show_default=True, help='Comma-separated list of contact points')
@click.option('--port', type=int, default=9042, show_default=True, help='Database RPC port')
@click.option('--username', help='Database username')
@click.option('--password', help='Database password')
@click.option('--keyspaces',
              help='Comma-separated list of keyspaces to include. If not specified all non-system keypaces will be included')
@click.option('--rf', type=int, help='Replication factor to override original setting. Optional.')
@click.option('--no-anonymize', help="Disable schema anonymization", is_flag=True)
@click.option('--output-dir', help='Directory schema files should be written to. If not specified, it will write to stdout')
@click.option('--purpose', help='Comments on the anticipated purpose of this schema.  Optional.')
@click.option('--maturity', help="The maturity of this schema.  Sample values would include 'alpha', 'beta', 'dev', 'test' or 'prod'.  Optional.")
@click.pass_context
def export(ctx, hosts, port, username, password, keyspaces, rf, no_anonymize, output_dir, purpose, maturity):

    ctx.ensure_object(dict)

    # If output-dir is specified make sure it actually exists, is a directory and is writable
    if output_dir is not None:
        if not os.path.exists(output_dir):
            log.error("Output directory " + output_dir + " does not exist")
            exit(3)
        if not os.path.isdir(output_dir):
            log.error("Specified output directory " + output_dir + " exists but is not a directory")
            exit(4)
        if not os.access(output_dir, os.W_OK):
            log.error("Output directory " + output_dir + " exists but is not writable")
            exit(5)

    ctx.obj['hosts'] = hosts.split(',')
    ctx.obj['port'] = port
    ctx.obj['username'] = username
    ctx.obj['password'] = password
    ctx.obj['keyspace-names'] = keyspaces.split(',') if keyspaces is not None else None
    ctx.obj['rf'] = rf
    ctx.obj['anonymize'] = not no_anonymize
    ctx.obj['output-dir'] = output_dir
    ctx.obj['purpose'] = purpose
    ctx.obj['maturity'] = maturity

# ============================ Utility functions for commands ============================
# unique_everseen from itertools recipes
def unique(iterable, key=None):
    "List unique elements, preserving order. Remember all elements ever seen."
    # unique_everseen('AAAABBBCCDAABBB') --> A B C D
    # unique_everseen('ABBCcAD', str.lower) --> A B C D
    seen = set()
    seen_add = seen.add
    if key is None:
        for element in filterfalse(seen.__contains__, iterable):
            seen_add(element)
            yield element
    else:
        for element in iterable:
            k = key(element)
            if k not in seen:
                seen_add(k)
                yield element


def with_cluster(cluster_fn, props):
    return _with_cluster(cluster_fn, **({k:props[k] for k in ["hosts", "port", "username", "password"]}))


def build_keyspace_id(ks):
    m = hashlib.sha256()
    m.update(ks.name.encode("utf-8"))
    # Leverage the urlsafe base64 encoding defined in RFC 4648, section 5 to provide an ID which can
    # safely be used for filenames as well
    return urlsafe_b64encode(m.digest()).decode('ascii')


def get_keyspaces(props, cluster):
    keyspace_names = props["keyspace-names"]
    metadata = cluster.metadata
    keyspaces = build_keyspace_objects(keyspace_names, metadata)

    if len(keyspaces) == 0:
        log.info("No valid keyspaces selected")
        exit(1)

    log.info("Processing the following keyspaces: %s", ','.join((ks.name for ks in keyspaces)))
    base_map = { ks : build_keyspace_id(ks) for ks in keyspaces}
    return {anonymize_keyspace(ks) if props['anonymize'] else ks : keyspace_id for (ks, keyspace_id) in base_map.items()}


def get_cluster_metadata(cluster):
    hosts = cluster.metadata.all_hosts()
    unique_dcs = unique((host.datacenter for host in hosts))
    return {"host-count": len(hosts), "dc-count": sum(1 for i in unique_dcs)}


def build_keyspace_map_and_metadata(props):
    def build_fn(cluster):
        metadata = {k : props[k] for k in ["purpose", "maturity"] if k in props}
        metadata.update(get_cluster_metadata(cluster))
        metadata["creation_timestamp"] = datetime.now(utc).isoformat()
        return (get_keyspaces(props, cluster), {k : metadata[k] for k in metadata.keys() if metadata[k]})

    return with_cluster(build_fn, props)


def export_keyspaces(props, export_fn, keyspace_map, metadata):
    options = {'rf': props['rf']}
    export_dir = props["output-dir"]

    if export_dir is None:
        print(export_fn(iter(keyspace_map.keys()), metadata, options))
    else:
        for (keyspace_obj, keyspace_id) in keyspace_map.items():
            # Output filename should use the file-safe alphabet for b64 encoding
            output_file_name = keyspace_id if props["anonymize"] else keyspace_obj.name
            with open(os.path.join(export_dir, output_file_name),'w') as output_file:
                log.info("Writing schema for keyspace " + keyspace_obj.name + " to file " + output_file.name)
                metadata['keyspace_id'] = keyspace_id
                output_file.write(export_fn(iter([keyspace_obj]), metadata, options))
                del metadata['keyspace_id']


# ============================ Command implementations ============================
@export.command()
@click.pass_context
def export_cql(ctx):
    """Export a schema as raw CQL statements"""
    (keyspace_map, metadata) = build_keyspace_map_and_metadata(ctx.obj)
    export_keyspaces(ctx.obj, export_cql_schema, keyspace_map, metadata)


@export.command()
@click.pass_context
def export_gemini(ctx):
    """Export a schema in a format suitable for use with the the Gemini data integrity test framework"""
    (keyspace_map, metadata) = build_keyspace_map_and_metadata(ctx.obj)

    if len(keyspace_map) > 1:
        log.error("Gemini schema doesn't support multiple keyspaces.")
        exit(2)

    export_keyspaces(ctx.obj, export_gemini_schema, keyspace_map, metadata)


@export.command()
@click.option('--token', help='Personal access token for Github user')
@click.pass_context
def contribute(ctx, token):
    """Contribute schemas to Adelphi"""
    if not ctx.obj["anonymize"]:
        log.error("All schemas contributed to Adelphi must be anonymized")
        exit(6)

    if ctx.obj["output-dir"]:
        log.info("Output directory argument is ignored when contributing a schema to Adelphi")

    gh = build_github(token)
    origin_repo = build_origin_repo(gh)

    with TemporaryDirectory() as tmpdir:
        props = ctx.obj.copy()
        props["output-dir"] = tmpdir

        (keyspace_map, metadata) = build_keyspace_map_and_metadata(ctx.obj)
        metadata["github_user"] = gh.get_user().login

        export_keyspaces(props, export_cql_schema, keyspace_map, metadata)

        if not sys.version_info[0] == 3:
            patch_repo(origin_repo)
            fork_repo = origin_repo.safe_create_fork()
        else:
            fork_repo = origin_repo.create_fork()
        branch_name = build_branch(gh, fork_repo)
        commit_schemas(gh, fork_repo, tmpdir, branch_name)
        build_pull_request(gh, origin_repo, branch_name)


if __name__ == '__main__':
    export(obj={}, auto_envvar_prefix="ADELPHI")
