import os
import sys
import gzip
from io import BytesIO
import json
import shutil
import requests
import tempfile
import tarfile
import urllib3
import re
from aimodelshare.data_sharing.utils import redo_with_write

urllib3.disable_warnings()

def get_auth_head_no_aws_auth(auth_url, registry, repository, type):
	resp = requests.get('{}?service={}&scope=repository:{}:pull'.format(auth_url, registry, repository), verify=False)
	access_token = resp.json()['token']
	auth_head = {'Authorization':'Bearer '+ access_token, 'Accept': type}
	return auth_head

def progress_bar(layer_label, nb_traits):
	sys.stdout.write('\r' + layer_label + 'Downloading [')
	for i in range(0, nb_traits):
		if i == nb_traits - 1:
			sys.stdout.write('>')
		else:
			sys.stdout.write('=')
	for i in range(0, 49 - nb_traits):
		sys.stdout.write(' ')
	sys.stdout.write(']')
	sys.stdout.flush()

def get_auth_url(registry): # to do with auth
	return 'https://' + registry + '/token/' # no aws auth

def get_auth_head(auth_url, registry, repository):
    # Broaden Accept header to allow manifest list / OCI fallbacks
    return get_auth_head_no_aws_auth(
        auth_url,
        registry,
        repository,
        ('application/vnd.docker.distribution.manifest.v2+json,'
         'application/vnd.docker.distribution.manifest.list.v2+json,'
         'application/vnd.oci.image.manifest.v1+json')
    )

def _fetch_concrete_manifest(registry, repository, tag_or_digest, auth_head):
    """Fetch a concrete image manifest (not a list)."""
    resp = requests.get(
        f'https://{registry}/v2/{repository}/manifests/{tag_or_digest}',
        headers=auth_head,
        verify=False
    )
    if not resp.ok:
        raise RuntimeError(
            f"Failed to fetch manifest {tag_or_digest} (status {resp.status_code}): {resp.text[:300]}"
        )
    return resp


def download_layer(layer, layer_count, tmp_img_dir, blobs_resp):

	ublob = layer['digest']
	layer_id = 'layer_' + str(layer_count) + '_' + ublob[7:]
	layer_label = ""
	layer_dir = tmp_img_dir + '/' + layer_id

	# Creating layer.tar file
	sys.stdout.write(layer_label + 'Downloading...')
	sys.stdout.flush()

	# Stream download and follow the progress
	unit = int(blobs_resp.headers['Content-Length']) / 50
	acc = 0
	nb_traits = 0
	progress_bar(layer_label, nb_traits)

	os.mkdir(layer_dir)
	with open(layer_dir + '/layer_gzip.tar', "wb") as file:
		for chunk in blobs_resp.iter_content(chunk_size=8192): 
			if chunk:
				file.write(chunk)
				acc = acc + 8192
				if acc > unit:
					nb_traits = nb_traits + 1
					progress_bar(layer_label, nb_traits)
					acc = 0

	sys.stdout.flush()

	with open(layer_dir + '/layer.tar', "wb") as file:
		unzip_layer = gzip.open(layer_dir + '/layer_gzip.tar','rb')
		shutil.copyfileobj(unzip_layer, file)
		unzip_layer.close()
	os.remove(layer_dir + '/layer_gzip.tar')

	return layer_id, layer_dir

def pull_image(image_uri):
    image_uri_parts = image_uri.split('/')
    registry = image_uri_parts[0]
    image, tag = image_uri_parts[2].split(':')
    repository = '/'.join([image_uri_parts[1], image])

    auth_url = get_auth_url(registry)
    auth_head = get_auth_head(auth_url, registry, repository)

    # 1. Fetch initial manifest (may be list or concrete)
    resp = _fetch_concrete_manifest(registry, repository, tag, auth_head)
    manifest_json = resp.json()

    # 2. Handle manifest list fallback
    if 'config' not in manifest_json:
        if 'manifests' in manifest_json:
            # Choose amd64 if available, else first
            chosen = None
            for m in manifest_json['manifests']:
                arch = (m.get('platform') or {}).get('architecture')
                if arch in ('amd64', 'x86_64'):
                    chosen = m
                    break
            if chosen is None:
                chosen = manifest_json['manifests'][0]
            digest = chosen['digest']
            # Re-auth to avoid token expiry
            auth_head = get_auth_head(auth_url, registry, repository)
            resp = _fetch_concrete_manifest(registry, repository, digest, auth_head)
            manifest_json = resp.json()
        else:
            raise KeyError(
                f"Manifest does not contain 'config' or 'manifests'. Keys: {list(manifest_json.keys())}"
            )

    if 'config' not in manifest_json or 'layers' not in manifest_json:
        raise KeyError(
            f"Unexpected manifest shape. Keys: {list(manifest_json.keys())}"
        )

    config = manifest_json['config']['digest']
    config_resp = requests.get(
        f'https://{registry}/v2/{repository}/blobs/{config}',
        headers=auth_head,
        verify=False
    )
    if not config_resp.ok:
        raise RuntimeError(
            f"Failed to fetch config blob {config} (status {config_resp.status_code}): {config_resp.text[:300]}"
        )

    tmp_img_dir = tempfile.gettempdir() + '/' + f'tmp_{image}_{tag}'
    os.mkdir(tmp_img_dir)

    with open(f'{tmp_img_dir}/{config[7:]}.json', 'wb') as f:
        f.write(config_resp.content)

    content = [{
        'Config': config[7:] + '.json',
        'RepoTags': [image_uri],
        'Layers': []
    }]

    layer_count = 0
    layers = manifest_json['layers']  # removed [6:] slicing

    for layer in layers:
        layer_count += 1
        # Refresh auth (avoid expiry)
        auth_head = get_auth_head(auth_url, registry, repository)
        blobs_resp = requests.get(
            f'https://{registry}/v2/{repository}/blobs/{layer["digest"]}',
            headers=auth_head,
            stream=True,
            verify=False
        )
        if not blobs_resp.ok:
            raise RuntimeError(
                f"Failed to stream layer {layer['digest']} status {blobs_resp.status_code}: {blobs_resp.text[:200]}"
            )
        layer_id, layer_dir = download_layer(layer, layer_count, tmp_img_dir, blobs_resp)
        content[0]['Layers'].append(layer_id + '/layer.tar')

        # Create layer json
        with open(layer_dir + '/json', 'w') as fjson:
            if layers[-1]['digest'] == layer['digest']:
                json_obj = json.loads(config_resp.content)
                json_obj.pop('history', None)
                json_obj.pop('rootfs', None)
            else:
                json_obj = {}
            json_obj['id'] = layer_id
            fjson.write(json.dumps(json_obj))

    with open(tmp_img_dir + '/manifest.json', 'w') as mf:
        mf.write(json.dumps(content))

    # repositories file
    repositories_json = {
        '/'.join(image_uri_parts[:-1]) + '/' + image: {tag: layer_id}
    }
    with open(tmp_img_dir + '/repositories', 'w') as rf:
        rf.write(json.dumps(repositories_json))

    docker_tar = tempfile.gettempdir() + '/' + '_'.join([repository.replace('/', '_'), tag]) + '.tar'
    tar = tarfile.open(docker_tar, "w")
    tar.add(tmp_img_dir, arcname=os.path.sep)
    tar.close()
    shutil.rmtree(tmp_img_dir, onerror=redo_with_write)
    return docker_tar


def extract_data_from_image(image_name, file_name, location):
    tar = tarfile.open(image_name, 'r')
    files = []
    for t in tar.getmembers():
        if('.tar' not in t.name):
            continue
        tar_layer = tarfile.open(fileobj=tar.extractfile(t))
        for tl in tar_layer.getmembers():
            if(re.match("var/task/"+file_name, tl.name)):
                files.append(tl)
        if(len(files)>0):
            break
    tar_layer.extractall(members=files, path=tempfile.gettempdir())
    if(os.path.isdir(file_name)):
        shutil.rmtree(file_name, onerror=redo_with_write)
    shutil.copytree(tempfile.gettempdir()+'/var/task/'+file_name, os.path.join(location, file_name))
    shutil.rmtree(tempfile.gettempdir()+'/var', onerror=redo_with_write)

def download_data(repository, location="./"):
	data_zip_name = repository.split('/')[2].split('-repository')[0]
	docker_tar = pull_image(repository)
	extract_data_from_image(docker_tar, data_zip_name, location)
	os.remove(docker_tar)
	print('\n\nData downloaded successfully.')


def import_quickstart_data(tutorial, section="modelplayground"):
    from aimodelshare.data_sharing.download_data import download_data
    import tensorflow as tf
    import os
    import pickle
    import shutil 
    import pandas as pd
    
   
    #Download Quick Start materials
    if all([tutorial == "flowers", section == "modelplayground"]):
        quickstart_repository = "public.ecr.aws/z5w0c9e9/quickstart_materials-repository:latest"
        existing_folder = 'flower_competition_data'

    if all([tutorial == "flowers", section == "competition"]):
        quickstart_repository = "public.ecr.aws/z5w0c9e9/quickstart_flowers_competition-repository:latest"
        existing_folder = 'flower_competition_data'
    
    if all([tutorial == "mnist", section == "modelplayground"]):
        quickstart_repository = "public.ecr.aws/y2e2a1d6/fashion_mnist_quickstart_materials-repository:latest"   
        existing_folder = 'fashion_mnist_competition_data'
    if all([tutorial == "mnist", section == "competition"]):
        quickstart_repository = "public.ecr.aws/y2e2a1d6/quickstart_mnist_competition-repository:latest"
        existing_folder = 'fashion_mnist_competition_data'
        
    if all([tutorial == "titanic", section == "modelplayground"]):
        quickstart_repository = "public.ecr.aws/y2e2a1d6/titanic_quickstart-repository:latest" 
        existing_folder = 'titanic_competition_data'
        
    if all([tutorial == "cars", section == "modelplayground"]):
        quickstart_repository = "public.ecr.aws/z5w0c9e9/quickstart_car_sales_competition-repository:latest" 
        existing_folder = 'used_car_competition_data'
        
    if all([tutorial == "clickbait", section == "modelplayground"]):
        quickstart_repository = "public.ecr.aws/y2e2a1d6/quickstart_clickbait_materials-repository:latest" 
        existing_folder = 'clickbait_competition_data'

    if all([tutorial == "covid_tweets", section == "modelplayground"]):
        quickstart_repository = "public.ecr.aws/y2e2a1d6/quickstart_covid_competition-repository:latest" 
        existing_folder = 'covid_tweet_competition_data'
        
    if all([tutorial == "sports", section == "modelplayground"]):
        quickstart_repository = "public.ecr.aws/y2e2a1d6/sports_quick_start_materials-repository:latest" 
        existing_folder = 'sports_clips_competition_data'
    if all([tutorial == "sports", section == "competition"]):
        quickstart_repository = "public.ecr.aws/y2e2a1d6/quickstart_sports_competition-repository:latest"
        existing_folder = 'sports_clips_competition_data'

    if all([tutorial == "dogs", section == "modelplayground"]):
        quickstart_repository = "public.ecr.aws/y2e2a1d6/dog_breed_quickstart_materials-repository:latest"   
        existing_folder = 'dog_competition_data'
    if all([tutorial == "dogs", section == "competition"]):
        quickstart_repository = "public.ecr.aws/y2e2a1d6/quickstart_dog_breed_competition-repository:latest"
        existing_folder = 'dog_competition_data'

    if all([tutorial == "imdb", section == "modelplayground"]):
        quickstart_repository = "public.ecr.aws/z5w0c9e9/imdb_quickstart_materials-repository:latest"
        existing_folder = 'imdb_competition_data'

    download_data(quickstart_repository)
    
    #Delete pre-existing tutorial folders
    if os.path.exists(existing_folder):
        shutil.rmtree(existing_folder, onerror=redo_with_write)

    #{{{ Prepare modelplayground materials
    if section == "modelplayground": 
        print("\nPreparing downloaded files for use...")
        
        if tutorial == "dogs":
           #instantiate model
            model = tf.keras.models.load_model('dog_breed_quickstart_materials/model.h5')
            
            #unpack data
            y_train = pd.read_csv("dog_breed_quickstart_materials/y_train.csv")
        
        if tutorial == "covid_tweets":
            #unpack data 
            X_train = pd.read_csv("quickstart_covid_competition/X_train.csv").squeeze("columns")
            X_test = pd.read_csv("quickstart_covid_competition/X_test.csv").squeeze("columns")
            y_test_labels = pd.read_csv("quickstart_covid_competition/y_test_labels.csv").squeeze("columns")
            y_train_labels = pd.read_csv("quickstart_covid_competition/y_train_labels.csv").squeeze("columns")
            # example data
            example_data = X_train[50:55]

            #move data files to cometition folder
            os.mkdir('covid_tweet_competition_data')
                
            files = ['quickstart_covid_competition/X_train.csv', 
                      'quickstart_covid_competition/X_test.csv',
                     'quickstart_covid_competition/y_train_labels.csv']
                
            for f in files:
                shutil.move(f, 'covid_tweet_competition_data')

        if tutorial == "imdb":
            #unpack data 
            X_train = pd.read_csv("imdb_quickstart_materials/X_train.csv").squeeze("columns")
            X_test = pd.read_csv("imdb_quickstart_materials/X_test.csv").squeeze("columns")
            with open("imdb_quickstart_materials/y_train_labels.json", "r") as f:
                y_train_labels = json.load(f)
            with open("imdb_quickstart_materials/y_test_labels.json", "r") as f:
                y_test_labels = json.load(f)
            import pandas as pd
            y_train_labels=pd.Series(y_train_labels)
            y_test_labels=pd.Series(y_test_labels)
            # example data
            example_data = X_train[50:55]

            #instantiate models
            lstm_model = tf.keras.models.load_model('imdb_quickstart_materials/model_1.h5')
            lstm_model2 = tf.keras.models.load_model('imdb_quickstart_materials/model_2.h5')

            #move data files to cometition folder
            os.mkdir('imdb_competition_data')
                
            files = ['imdb_quickstart_materials/X_train.csv', 
                      'imdb_quickstart_materials/X_test.csv',
                      'imdb_quickstart_materials/y_train_labels.csv']
                
            for f in files:
                shutil.move(f, 'imdb_competition_data')

        if tutorial == "flowers":
           #instantiate model
            model = tf.keras.models.load_model('quickstart_materials/flowermodel.h5')
            
            #unpack data
            with open("quickstart_materials/y_train_labels.txt", "rb") as fp:  
                y_train_labels = pickle.load(fp)

        if tutorial == "mnist":
           #instantiate model
            model = tf.keras.models.load_model('fashion_mnist_quickstart_materials/mnist_model_1.h5')
            
            #unpack data
            with open("fashion_mnist_quickstart_materials/y_train_labels.pkl", "rb") as fp:  
                y_train_labels = pickle.load(fp)
        
        if tutorial == "sports":
           #instantiate model
            model = tf.keras.models.load_model('sports_quick_start_materials/video_1.h5')
            
            #unpack data
            y_train_labels = pd.read_csv("sports_quick_start_materials/y_train.csv")
                
        if tutorial == "clickbait":
            import pandas as pd
            # suppress tf warning
            import tensorflow as tf
            tf.compat.v1.logging.set_verbosity(tf.compat.v1.logging.ERROR)
            
            #instantiate models
            lstm_model = tf.keras.models.load_model('quickstart_clickbait_materials/lstm_model1.h5')
            lstm_model2 = tf.keras.models.load_model('quickstart_clickbait_materials/lstm_model2.h5')
           
            # bring in data 
            clickbait = pd.read_csv('quickstart_clickbait_materials/clickbait_data',  sep="\n", header = None)
            clickbait['label'] = "clickbait"
            clickbait.columns = ['headline', 'label']
        
            not_clickbait = pd.read_csv('quickstart_clickbait_materials/non_clickbait_data',  sep="\n", header = None)
            not_clickbait['label'] = "not clickbait"
            not_clickbait.columns = ['headline', 'label']

            # train/test/split
            from sklearn.model_selection import train_test_split
            X = clickbait.append(not_clickbait)
            y = X['label']
            X = X.drop(['label'], axis=1)
            X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=942)
            X_test.reset_index(drop=True, inplace=True)
            X_test = X_test.squeeze()
            X_train.reset_index(drop=True, inplace=True)
            X_train = X_train.squeeze()

            # exampledata
            example_data = X_train[0:5]
            
            # Create data directory for competition 
            X_train.to_csv("X_train.csv", index=False)
            X_test.to_csv("X_test.csv", index=False)
            y_train.to_csv("y_train.csv", index=False)
            
            os.mkdir('clickbait_competition_data')
            
            files = ['X_train.csv', 'X_test.csv', 'y_train.csv']
                        
            for f in files:
                shutil.move(f, 'clickbait_competition_data')

        if tutorial == "titanic":
            from sklearn.model_selection import train_test_split
            import pandas as pd
            #read in data
            data = pd.read_csv("titanic_quickstart/titanic_data.csv")
            y = data['survived']
            y = y.map({0: 'died', 1: 'survived'}) 
            X = data.drop(['survived','sibsp','parch','ticket','name','cabin','boat','body','home.dest'], axis=1)
            #create subset as exampledata 
            example_data = pd.DataFrame(X[0:4])
            # create data directory for competition
            X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2)
            training_data = pd.merge(X_train, y_train, left_index=True, right_index=True)
            training_data = training_data.drop(list(training_data.filter(like='Unnamed')), axis=1)
            training_data.to_csv("training_data.csv", index=False)

            test_data = X_test
            test_data = test_data.drop(list(test_data.filter(like='Unnamed')), axis=1)
            test_data.to_csv("test_data.csv", index=False)
            
            os.mkdir('titanic_competition_data')
            files = ['training_data.csv', 
                   'test_data.csv']
            
            for f in files:
                shutil.move(f, 'titanic_competition_data')
            
            #make y_test_labels for competition
            y_test_labels = y_test.to_list()

        if tutorial == "cars":
            from sklearn.model_selection import train_test_split

            # read in data
            import pandas as pd
            data = pd.read_csv("quickstart_car_sales_competition/used_car_dataset.csv")
            y = data['selling_price']
            X = data.drop(['selling_price', 'torque', 'name'], axis=1)

            #Data Prep:
                # convert rupees to $ (for smaller MSEs)
            y = y.mul(.014)
                # A: Split units from mileage and convert units 
            Correct_Mileage= []
            for i in X.mileage:
                if str(i).endswith('km/kg'):
                    i = i[:-6]
                    i = float(i)*1.40
                    Correct_Mileage.append(float(i))
                elif str(i).endswith('kmpl'):
                    i = i[:-5]
                    Correct_Mileage.append(float(i))
                else: 
                    Correct_Mileage.append(None)
            X['mileage']=Correct_Mileage

                  #B: Split units from engine, & max_pwer
            X['engine'] = X['engine'].str.replace(' CC', '')
            X['engine'] = pd.to_numeric(X['engine'])

            X['max_power'] = X['max_power'].str.replace(' bhp', '')
            X['max_power'] = pd.to_numeric(X['max_power'])

            #create subset as exampledata 
            example_data = pd.DataFrame(X[0:4])

            #create data directory for competition
            X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2)
            training_data = X_train
            training_data = pd.merge(X_train, y_train, left_index=True, right_index=True)
            training_data.to_csv("training_data.csv", index=False)

            test_data = X_test
            test_data.to_csv("test_data.csv", index=False)

            os.mkdir('used_car_competition_data')
            files = ['training_data.csv', 
                        'test_data.csv']
            for f in files:
                shutil.move(f, 'used_car_competition_data')
    #}}}

    #{{{ prepare competition materials
    if section == "competition":
        print("\nPreparing downloaded files for use...")

        if tutorial == "flowers":
            #Instantiate Model 
            model_2 = tf.keras.models.load_model('quickstart_flowers_competition/flowermodel_2.h5')
        
            #unpack data
            with open("quickstart_flowers_competition/y_test_labels.txt", "rb") as fp:  
                y_test_labels = pickle.load(fp)
                
            #move data files to folder to upload with create_competiton
            os.mkdir('flower_competition_data')
                
            folders = ['quickstart_flowers_competition/test_images', 
                      'quickstart_flowers_competition/train_images']
                
            for f in folders:
                shutil.move(f, 'flower_competition_data')
        
        if tutorial == "mnist":
            #Instantiate Model 
            model_2 = tf.keras.models.load_model('quickstart_mnist_competition/mnist_model_2.h5')
        
            #unpack data
            with open("quickstart_mnist_competition/y_test_labels.pkl", "rb") as fp:  
                y_test_labels = pickle.load(fp)
                
            #move data files to folder to upload with create_competiton
            os.mkdir('fashion_mnist_competition_data')
                
            folders = ['quickstart_mnist_competition/test_data', 
                      'quickstart_mnist_competition/training_data', 
                       'fashion_mnist_quickstart_materials/y_train_labels.pkl']
                
            for f in folders:
                shutil.move(f, 'fashion_mnist_competition_data')

        if tutorial == "dogs":
            #Instantiate Model 
            model_2 = tf.keras.models.load_model('quickstart_dog_breed_competition/model_2.h5')
        
            #unpack data
            with open("quickstart_dog_breed_competition/y_test_labels.txt", "rb") as fp:  
                y_test_labels = pickle.load(fp)
                
            #move data files to folder to upload with create_competiton
            os.mkdir('dog_competition_data')
                
            folders = ['quickstart_dog_breed_competition/dog_breed_competition_data/test_images', 
                      'quickstart_dog_breed_competition/dog_breed_competition_data/train_images']
                
            for f in folders:
                shutil.move(f, 'dog_competition_data')
        
        if tutorial == "sports":
            model_2 = tf.keras.models.load_model('quickstart_sports_competition/video_2.h5')
            y_test = pd.read_csv("quickstart_sports_competition/y_test.csv")
            y_test_labels = y_test.idxmax(axis=1)
            os.mkdir('sports_clips_competition_data')
            folders = ['quickstart_sports_competition/clips_test.zip', 
                      'quickstart_sports_competition/clips_train.zip']
            for f in folders:
                shutil.move(f, 'sports_clips_competition_data')
    #}}}

    success_message = ("\nSuccess! Your Quick Start materials have been downloaded. \n"
                       "You are now ready to run the tutorial.")
    
    print(success_message)

    if all([tutorial == "flowers", section == "modelplayground"]):
        return model, y_train_labels

    if all ([tutorial == "flowers", section == "competition"]): 
        return model_2, y_test_labels
    
    if all([tutorial == "mnist", section == "modelplayground"]):
        return model, y_train_labels

    if all([tutorial == "mnist", section == "competition"]): 
        return model_2, y_test_labels

    if all([tutorial == "dogs", section == "modelplayground"]):
        return model, y_train

    if all ([tutorial == "dogs", section == "competition"]): 
        return model_2, y_test_labels
    
    if all([tutorial == "sports", section == "modelplayground"]):
        return model, y_train_labels
    
    if all ([tutorial == "sports", section == "competition"]): 
        return model_2, y_test, y_test_labels
    
    if tutorial == "titanic":
        return X_train, X_test, y_train, y_test, example_data, y_test_labels

    if tutorial == "cars":
        return X_train, X_test, y_train, y_test, example_data
    
    if tutorial == "clickbait":
        return X_train, X_test, y_train, y_test, example_data, lstm_model, lstm_model2	

    if tutorial == "imdb":
        return X_train, X_test, y_train_labels, y_test_labels, example_data, lstm_model, lstm_model2	

    if tutorial == "covid_tweets":
        return X_train, X_test, y_train_labels, y_test_labels, example_data
