from coegil.ApiService import set_configuration, get_environment, get_sdk
import pkg_resources
import tempfile
import subprocess
import sys
import tarfile
import os
import click
import json


@click.group()
def cli():
    pass


@cli.command('configure', help="Configure your environment.")
@click.option('--key', prompt=True, hide_input=True, required=True, help="Your Coegil API key.")
@click.option('--env', required=False, default='prod', help="Optionally, set an environment.  "
                                                            "If you don't know what this is, leave it blank.")
def configure(key: str, env: str):
    set_configuration(env, key)


@cli.command('validate', help="Validate your environment.")
def validate():
    result = get_sdk().validate()
    environment = get_environment()
    if environment != 'prod':
        result['environment'] = environment
    _write_to_std_out(result)


@cli.command('version', help="Lists The version of the CLI.")
def version():
    _write_to_std_out(pkg_resources.get_distribution("coegilcli").version)


@cli.command('generate-temporary-credentials', help="Generate temporary credentials.")
@click.option('--duration_hours', required=False, help="The number of hours to keep the credential active.  "
                                                       "The default is 36 hours.  The max is 36 hours.", type=int)
def generate_temporary_credentials(duration_hours: int):
    result = get_sdk().generate_temporary_api_key(duration_hours)
    _write_to_std_out(result)


@cli.command('query', help="Run an ad hoc database query.")
@click.argument('query_text')
@click.option('--asset_id', required=True, help="The asset's unique identifier.")
@click.option('--database_name', required=True, help="Your database artifact's name.")
@click.option('--engine', default='mysql', type=click.Choice(['mysql', 'athena'], case_sensitive=False),
              help='The database engine (defaults to mysql)')
def query(query_text: str, engine: str, asset_id: str, database_name: str):
    result = get_sdk().query_database(query_text, asset_id, database_name, engine)
    _write_to_std_out(result)


@cli.command('saved-query', help="Run a saved database query.")
@click.option('--asset_id', required=True, help="The asset's unique identifier.")
@click.option('--query_name', required=True, help="Your saved query's name.")
@click.option('--parameters', required=False, help="Parameter overrides - passed as a JSON string of a dictionary.")
def saved_query(asset_id: str, query_name: str, parameters: str):
    result = get_sdk().run_saved_database_query(asset_id, query_name,
                                                None if parameters is None else json.loads(parameters))
    _write_to_std_out(result)


@cli.command('queue-schedule', help="Queues a job or pipeline schedule to be invoked.  "
                                    "The queue is drained in batches.")
@click.option('--asset_id', required=True, help="The asset's unique identifier.")
@click.option('--job_name', required=True, help="The name of the job artifact (pipeline, notebook, or trigger).")
@click.option('--schedule_name', required=False, help="Required unless you are invoking a trigger.")
@click.option('--variable_override', required=False, help="Passed as a JSON string of a dictionary.")
def queue_schedule(asset_id: str, job_name: str, schedule_name: str, variable_override: str):
    get_sdk().queue_schedule(asset_id, job_name, schedule_name,
                             None if variable_override is None else json.loads(variable_override))


@cli.command('invoke-schedule', help="Invoke a job or pipeline schedule.")
@click.option('--asset_id', required=True, help="The asset's unique identifier.")
@click.option('--job_name', required=True, help="The name of the job artifact (pipeline, notebook, or trigger).")
@click.option('--schedule_name', required=False, help="Required unless you are invoking a trigger.")
@click.option('--variable_override', required=False, help="Passed as a JSON string of a dictionary.")
def invoke_schedule(asset_id: str, job_name: str, schedule_name: str, variable_override: str):
    result = get_sdk().invoke_schedule(asset_id, job_name, schedule_name,
                                       None if variable_override is None else json.loads(variable_override))
    _write_to_std_out(result)


@cli.command('get-schedule-status', help="Gets the status of an invoked job")
@click.option('--job_id', required=True, help="The job run id returned from the invoke call.")
def get_schedule_status(job_id: str):
    result = get_sdk().get_schedule_status(job_id)
    _write_to_std_out(result)


@cli.command('list-schedules', help="List all of the schedules in a given asset")
@click.option('--asset_id', required=True, help="The asset's unique identifier.")
def list_schedules(asset_id: str):
    result = get_sdk().list_schedules(asset_id)
    _write_to_std_out(result)


@cli.command('list-artifacts', help="Lists all of the artifacts in a given asset.")
@click.option('--asset_id', required=True, help="The asset's unique identifier.")
def list_artifacts(asset_id: str):
    result = get_sdk().list_artifacts(asset_id)
    _write_to_std_out(result)


@cli.command('download-file', help="Download a file artifact to your local drive.")
@click.option('--asset_id', required=True, help="The asset's unique identifier.")
@click.option('--file', required=True, type=click.Path(exists=False, file_okay=True, dir_okay=False),
              help="The location to save the file.")
@click.option('--artifact_name', required=True, help="The artifact's name.")
def download_file(asset_id: str, file: str, artifact_name: str):
    contents = get_sdk().read_file(asset_id, artifact_name)
    with open(file, 'wb+') as f:
        f.write(contents)


@cli.command('save-file', help="Saves a file to an asset.")
@click.option('--asset_id', required=True, help="The asset's unique identifier.")
@click.option('--file', required=True, type=click.Path(exists=True, file_okay=True, dir_okay=False),
              help="The location of the file to upload.")
@click.option('--artifact_name', required=False, help="The artifact's name.  If not specified, "
                                                      "it will be the file's name.")
def save_file(asset_id: str, file: str, artifact_name: str):
    artifact_name = os.path.basename(file) if artifact_name is None else artifact_name
    _upload_file(asset_id, artifact_name, file)


@cli.command('publish', help="Publish a python library to an asset.")
@click.option('--asset_id', required=True, help="The asset's unique identifier.")
@click.option('--package_name', required=True, help="Your package's name.")
@click.option('--library', required=True, type=click.Path(exists=True),
              help="The location of the Python library to package.")
def publish(asset_id: str, package_name: str, library: str):
    with tempfile.TemporaryDirectory() as wheel_destination:
        commands = [
            sys.executable,
            "-m",
            "pip",
            "wheel",
            library,
            f'--wheel-dir={wheel_destination}'
        ]
        process = subprocess.Popen(commands, stderr=subprocess.PIPE)
        out, err = process.communicate()
        output = out.decode() if out is not None else ''
        error = err.decode()
        success = error == ''
        _write_to_std_out(output)
        if not success:
            raise Exception(error)
        tar_file_name = os.path.join(wheel_destination, f'{package_name}.tar')
        with tarfile.open(tar_file_name, "w:gz") as tar:
            tar.add(wheel_destination, arcname=os.path.basename(wheel_destination))
        _upload_file(asset_id, package_name, tar_file_name, 'python_package')


def _upload_file(asset_id: str, artifact_name: str, file_location: str, artifact_sub_type: str = None):
    with open(file_location, 'rb') as contents:
        get_sdk().save_file(asset_id, artifact_name, contents, artifact_sub_type)


def _write_to_std_out(val, is_error=False):
    click.echo(json.dumps(val), err=is_error)


if __name__ == '__main__':
    cli()
