# -*- coding: utf-8 -*-

"""Main module."""
from __future__ import print_function
import os
import uuid
import tarfile
from pprint import pprint
import shutil
import boto3
from botocore.client import Config
import kubernetes.client
import kubernetes.config
from kubernetes.client.rest import ApiException
import pandas as pd
import tempfile

#Global definition strings
group = 'argoproj.io' # str | The custom resource's group name
version = 'v1alpha1' # str | The custom resource's version
namespace = 'argo' # str | The custom resource's namespace
plural = 'workflows' # str | The custom resource's plural name. For TPRs this would be lowercase plural kind.

AWS_HOST = os.getenv('AWS_HOST')
AWS_ENDPOINT = os.getenv('AWS_ENDPOINT')
AWS_ACCESS_KEY_ID = os.getenv('AWS_ACCESS_KEY_ID')
AWS_SECRET_ACCESS_KEY = os.getenv('AWS_SECRET_ACCESS_KEY')
hostname = os.environ.get('HOSTNAME')


def setup_k8s_api():
    """
    Common actions to setup Kubernetes API access to Argo workflows
    """
    kubernetes.config.load_incluster_config() #Only works inside of JupyterLab Pod
    configuration = kubernetes.client.Configuration() #Create configuration
    
    return kubernetes.client.CustomObjectsApi(kubernetes.client.ApiClient(configuration)) 

def get_jupyter_username():
    try:
        assert(hostname!=None) #check that env variable was read correctly
    except AssertionError:
        print("Error: Hostname cannot be retrieved. Are you running in JupyterLab?")
        return
    return hostname[8:] #Remove prefix 'jupyter-' from hostmane to get username

def list_argo_jobs():
    """
    This will return a list of all jobs associated with the current user
    """
    
    api_instance = setup_k8s_api()
    username = get_jupyter_username()
    
    try:
        api_response = api_instance.list_namespaced_custom_object(group, version, namespace, plural, label_selector = f'username={username},hidden=false')
        data = []
        for item in api_response['items']:
            notebook_name = ''
            if 'notebook-name' in item['metadata']['labels']:
                notebook_name = f"{item['metadata']['labels']['notebook-name']}"
            job_uuid = ''
            if 'uuid' in item['metadata']['labels']:
                job_uuid = f"{item['metadata']['labels']['uuid']}"
            data.append([f"{job_uuid}", notebook_name, f"{item['status']['phase']}", f"{item['status']['startedAt']}"])
        df = pd.DataFrame(data, columns=["ID", "NOTEBOOK NAME", "STATUS", "STARTED AT"])
        df['STARTED AT'] = pd.to_datetime(df['STARTED AT'])
        df['STARTED AT'] = df['STARTED AT'].dt.strftime('%b %d, %H:%M')
        df.set_index('ID', inplace=True)
        df = df.sort_values(by=['STARTED AT'], ascending=False)
        if len(df.index)>0:
            pprint(df)
    except ApiException as e:
        print("Error: Cannot list workflows")
    return
    
def create_argo_job(body):
    api_instance = setup_k8s_api()
            
    try: 
        api_response = api_instance.create_namespaced_custom_object(group, version, namespace, plural, body)
        print(f'Workflow {api_response["metadata"]["labels"]["uuid"]} created')
    except ApiException as e:
        print("Error: Cannot create workflow")
                         
def schedule_notebook(filename):
    """
    Send the Notebook with the given filename to Argo cluster for execution
    """
                         
    job_uuid = uuid.uuid4().hex
    
    # Put Notebook file into local temporary archive
    local_temp_archive_handle, local_temp_archive_path = tempfile.mkstemp(suffix='.tgz')
    os.close(local_temp_archive_handle)
    tar = tarfile.open(local_temp_archive_path, "w:gz")
    tar.add(filename)
    tar.close()
       
    # Upload archive to S3
    #Setup S3 resource
    try:
        assert(AWS_HOST!=None)
        assert(AWS_ENDPOINT!=None)
        assert(AWS_ACCESS_KEY_ID!=None)
        assert(AWS_SECRET_ACCESS_KEY!=None)
    except AssertionError:
        print("Error: S3 credentials are not set up")
        return
    
    s3 = boto3.resource('s3',
                        endpoint_url=AWS_ENDPOINT,
                        aws_access_key_id=AWS_ACCESS_KEY_ID,
                        aws_secret_access_key=AWS_SECRET_ACCESS_KEY,
                        config=Config(signature_version='s3v4'))

    #Create bucket object
    bucket = s3.Bucket('rookbucket')
                         
    username = get_jupyter_username()

    # Archived notebook is stored in S3 bucket at the following path:
    # '/<username>/inputs/<workflow_name>.tgz'
    # <workflow_name> is unique random string generated by uuid
    s3_archive_uuid = job_uuid
    s3_archive_key = username + "/inputs/" + s3_archive_uuid + ".tgz"
    bucket.upload_file(local_temp_archive_path, s3_archive_key)
    
    os.remove(local_temp_archive_path)
    #Check that upload was successful and delete local archive
    try:
        [bucket_object.key for bucket_object in bucket.objects.all()].index(s3_archive_key)
    except ValueError:
        print('Error: Notebook upload failed')
        return
         
    # Define the JSON schema of the Workflow Custom Resource Definition to create.
    body = {
      "apiVersion": "argoproj.io/v1alpha1",
      "kind": "Workflow",
      "metadata": {
        "name": username + "-workflow-" + job_uuid,
        "namespace": "argo",
        "labels": {
            "username": username,
            "uuid": job_uuid,
            "notebook-name": filename,
            "hidden": "false"
        }
      },
      "spec": {
        "serviceAccountName": "workflow",
        "arguments": {
          "artifacts": [
            {
              "name": "notebook-in",
              "s3": {
                "accessKeySecret": {
                  "key": "AccessKey",
                  "name": "rook-ceph-object-user-my-store-my-user"
                },
                "bucket": "rookbucket",
                "endpoint": "rook-ceph-rgw-my-store.rook-ceph",
                "insecure": True,
                "key": s3_archive_key,
                "secretKeySecret": {
                  "key": "SecretKey",
                  "name": "rook-ceph-object-user-my-store-my-user"
                }
              }
            }
          ]
        },
        "entrypoint": "papermill",
        "templates": [
          {
            "container": {
              "command": [
                "sh",
                "-c",
                "papermill /tmp/in.ipynb /tmp/out.ipynb"
              ],
              "image": "ktaletsk/polyglot-jupyter:latest"
            },
            "inputs": {
              "artifacts": [
                {
                  "name": "notebook-in",
                  "path": "/tmp/in.ipynb"
                }
              ]
            },
            "name": "papermill",
            "outputs": {
              "artifacts": [
                {
                  "name": "notebook-out",
                  "path": "/tmp/out.ipynb"
                }
              ]
            }
          }
        ]
      }
    }    
    create_argo_job(body)

def get_workflow(job_uuid, destination):
    username = get_jupyter_username()
    jobname = username + '-workflow-' + job_uuid
    
    #Setup S3 resource
    try:
        assert(AWS_HOST!=None)
        assert(AWS_ENDPOINT!=None)
        assert(AWS_ACCESS_KEY_ID!=None)
        assert(AWS_SECRET_ACCESS_KEY!=None)
    except AssertionError:
        print("Error: S3 credentials are not set up")
        return
    
    s3 = boto3.resource('s3',
                        endpoint_url=AWS_ENDPOINT,
                        aws_access_key_id=AWS_ACCESS_KEY_ID,
                        aws_secret_access_key=AWS_SECRET_ACCESS_KEY,
                        config=Config(signature_version='s3v4'))

    #Create bucket object
    bucket = s3.Bucket('rookbucket')
    
    # Download Argo atrifact (archive) to a temporary location
    local_temp_archive_handle, local_temp_archive_path = tempfile.mkstemp(suffix='.tgz')
    os.close(local_temp_archive_handle)

    # S3 location of the output is standardized
    result_key = jobname + '/' + jobname + '/' + 'notebook-out.tgz'
    
    # Download archive to a temporary place
    bucket.download_file(result_key, local_temp_archive_path)

    tar = tarfile.open(local_temp_archive_path)
    job_uuid = jobname.split('-')[-1]

    # Extract the resulting notebook to a temp location
    local_temp_result_folder = tempfile.mkdtemp()
    for member in tar.getmembers():
        if member.name == 'out.ipynb':
            tar.extract(member, local_temp_result_folder) # extract to a temp location

    # Rename resulting notebook and move it to the desired location
    shutil.move(local_temp_result_folder + '/out.ipynb', destination)

def delete_workflow(job_uuid):
    #Get CRD for uuid
    username = get_jupyter_username()
    jobname = username + '-workflow-' + job_uuid
    api_instance = setup_k8s_api()
    try: 
        api_response = api_instance.get_namespaced_custom_object(group, version, namespace, plural, jobname)
    except ApiException as e:
        print(f'Cannot find job {job_uuid}')

    #Patch CRD to change label
    api_response['metadata']['labels']['hidden'] = 'true'
    try: 
        api_instance.patch_namespaced_custom_object(group, version, namespace, plural, jobname, api_response)
    except ApiException as e:
        print("Cannot delete job {job_uuid}")

    # Clean up S3 bucket from artifacts
    #Setup S3 resource
    try:
        assert(AWS_HOST!=None)
        assert(AWS_ENDPOINT!=None)
        assert(AWS_ACCESS_KEY_ID!=None)
        assert(AWS_SECRET_ACCESS_KEY!=None)
    except AssertionError:
        print("Error: S3 credentials are not set up")
        return
    
    s3 = boto3.resource('s3',
                        endpoint_url=AWS_ENDPOINT,
                        aws_access_key_id=AWS_ACCESS_KEY_ID,
                        aws_secret_access_key=AWS_SECRET_ACCESS_KEY,
                        config=Config(signature_version='s3v4'))

    #Create bucket object
    bucket = s3.Bucket('rookbucket')
    
    s3_input_archive_key = username + "/inputs/" + job_uuid + ".tgz"
    s3_output_archive_key = jobname + '/' + jobname + '/' + 'notebook-out.tgz'
    
    bucket.delete_objects(
        Delete={
            'Objects': [
                {'Key': s3_input_archive_key},
                {'Key': s3_output_archive_key},
            ],
        }
    )
