import re
import os
import requests
import json
from datetime import datetime, date
import ibm_boto3
from ibm_botocore.client import Config, ClientError
from tqdm import tqdm
import getpass

def default(o):
    if isinstance(o, (date, datetime)):
        return o.strftime("%Y-%m-%dT%H:%M:%S.000+00:00")

def isValidFolder():
    return bool(re.match(r"^/resources/labs/cvstudio/[0-9a-z-]+/(runs|apps)/[0-9a-z-]+$", os.getcwd()))

def tokenToHeaders(token):
    return {
        'Authorization': 'Bearer ' + token,
        'Content-Type': 'application/json'
    }

class CVStudio(object):
    @staticmethod
    def interactive(token=None, ibm_api_key_id=None, base_url=None):
        if base_url == None:
            base_url = os.environ.get('CV_STUDIO_BASE_URL', 'https://vision.skills.network')

        if token is None:
            token = os.environ.get('CV_STUDIO_TOKEN')
        if token is None:
            # raise Exception('need to pass valid token or set CV_STUDIO_TOKEN environment variable')
            print('Need to validate your identity as no CV Studio token set in environment')
            print('Please visit the following url in your browser and copy the token: ' + base_url + '/api/token')
            token = getpass.getpass('Enter your CV Studio token: ')
        
        projects = CVStudio.getProjects(token, base_url)

        print('Projects:')
        i=1
        for project in projects:
            print(str(i) + ". " + project['name'])
            i+=1

        val = input("Choose a Project: ")
        val = int(val)

        project = projects[val-1]
        print()
        print("You've chosen: " + project['name'])
        print()

        print("Connect to an Training Run or Application")
        print("1. Training Run")
        print("2. Application")

        val = input("Choose a type: ")

        if val == "1":
            print("You've chosen: Training Runs")
            print()

            runs = CVStudio.getRuns(token, base_url, project['id'])
            print('Training Runs in selected Project:')
            i=1
            for run in runs:
                print(str(i) + ". " + run['name'])
                i+=1

            val = input("Choose a Training Run: ")
            val = int(val)

            run = runs[val-1]
            print()
            print("You've chosen: " + run['name'])

            token = CVStudio.getTokenInteractive(token, base_url, 'run', run['id'])
        
        elif val == "2":
            print("You've chosen: Applications")
            print()

            apps = CVStudio.getApplications(token, base_url, project['id'])
            print('Applications in selected Project:')
            i=1
            for app in apps:
                print(str(i) + ". " + app['name'])
                i+=1

            val = input("Choose an Application: ")
            val = int(val)

            app = apps[val-1]
            print()
            print("You've chosen: " + app['name'])

            token = CVStudio.getTokenInteractive(token, base_url, 'app', app['id'])
        
        return token

    @staticmethod
    def getProjects(token, base_url):
        data = {
            'type': 'projects'
        }
        x = requests.post(base_url + '/api/interactive', headers=tokenToHeaders(token), data=json.dumps(data, default=default))
        if not x.ok:
            raise Exception("Failed to setup CV Studio client")

        return json.loads(x.text)

    @staticmethod
    def getRuns(token, base_url, projectId):
        data = {
            'projectId': projectId,
            'type': 'runs'
        }
        x = requests.post(base_url + '/api/interactive', headers=tokenToHeaders(token), data=json.dumps(data, default=default))
        if not x.ok:
            raise Exception("Failed to setup CV Studio client")

        return json.loads(x.text)

    @staticmethod
    def getApplications(token, base_url, projectId):
        data = {
            'projectId': projectId,
            'type': 'apps'
        }
        x = requests.post(base_url + '/api/interactive', headers=tokenToHeaders(token), data=json.dumps(data, default=default))
        if not x.ok:
            raise Exception("Failed to setup CV Studio client")

        return json.loads(x.text)

    @staticmethod
    def getTokenInteractive(token, base_url, itemType, itemId):
        data = {
            'itemId': itemId,
            'type': itemType + 'Token'
        }
        x = requests.post(base_url + '/api/interactive', headers=tokenToHeaders(token), data=json.dumps(data, default=default))
        if not x.ok:
            raise Exception("Failed to setup CV Studio client")

        return json.loads(x.text)

    def __init__(self, interactive=False, token=None, ibm_api_key_id=None, base_url=None, public=False):
        self.public = public

        if base_url == None:
            self.base_url = os.environ.get('CV_STUDIO_BASE_URL', 'https://vision.skills.network')
        else:
            self.base_url = base_url

        if interactive:
            result = CVStudio.interactive(base_url=base_url)
            if 'token' in result:
                token = result['token']
            if 'apikey' in result:
                ibm_api_key_id = result['apikey']

        if token == None:
            self.token = self.getToken()
            if self.token is None:
                raise Exception('need to pass valid token or set CV_STUDIO_TOKEN environment variable')
        else:
            self.token = token

        if ibm_api_key_id == None:
            self.ibm_api_key_id = os.environ.get('IBMCLOUD_API_KEY')
            if self.ibm_api_key_id == None:
                raise Exception('need to pass valid ibm_api_key_id or set IBMCLOUD_API_KEY environment variable')
        else:
            self.ibm_api_key_id = ibm_api_key_id

        self.setup()

    def getToken(self):
        token = os.environ.get('CV_STUDIO_TOKEN')
        if token is None:
            return None

        if not isValidFolder():
            raise Exception('Need to run from a valid folder. e.g. /resources/labs/cvstudio/<project-name>/runs/<run-name>')

        rest = os.getcwd()
        rest, itemName = os.path.split(rest)
        rest, itemType = os.path.split(rest)
        rest, projectName = os.path.split(rest)

        data = {
            'projectName': projectName,
            'itemType': itemType,
            'itemName': itemName
        }

        headers = {
            'Authorization': 'Bearer ' + token,
            'Content-Type': 'application/json'
        }

        x = requests.post(self.base_url + '/api/token', headers=headers, data=json.dumps(data, default=default))

        if not x.ok:
            raise Exception("Failed to setup CV Studio client")

        item = json.loads(x.text)        
        return item['token']

    def setup(self):
        headers = {
            'Authorization': 'Bearer ' + self.token,
            'Content-Type': 'application/json'
        }
        
        x = requests.post(self.base_url + '/api/cos_credentials', headers=headers)

        if x.ok:
            item = json.loads(x.text)
            

            self.endpoint = item['endpoint']
            if self.public:
                self.endpoint = self.endpoint.replace('.private', '')

            self.bucket = item['bucket']

            config = Config(
                connect_timeout=1,
                retries={'max_attempts': 2},
                signature_version="oauth"
            )

            self.cos = ibm_boto3.client("s3",
                ibm_api_key_id=self.ibm_api_key_id,
                ibm_service_instance_id=self.ibm_api_key_id,
                config=config,
                endpoint_url=self.endpoint)
            
            if self.public:
                return x
            
            try:
                self.cos.head_bucket(Bucket=self.bucket)
            except Exception as e:
                self.endpoint = self.endpoint.replace('.private', '.direct')
            
            self.cos = ibm_boto3.client("s3",
                ibm_api_key_id=self.ibm_api_key_id,
                ibm_service_instance_id=self.ibm_api_key_id,
                config=config,
                endpoint_url=self.endpoint)

            try:
                self.cos.head_bucket(Bucket=self.bucket)
            except Exception as e:
                self.endpoint = self.endpoint.replace('.direct', '')
            
            self.cos = ibm_boto3.client("s3",
                ibm_api_key_id=self.ibm_api_key_id,
                ibm_service_instance_id=self.ibm_api_key_id,
                config=config,
                endpoint_url=self.endpoint)

        else:
            print("Failed to setup CV Studio client")

        return x

    def download_file_cos(self, key, folder=os.getcwd(), attempts=0):
        try:
            local_file_name=os.path.join(folder, key)
            res=self.cos.download_file(Bucket=self.bucket, Key=key, Filename=local_file_name)
        except Exception as e:
            print('Exception', e)
            if attempts < 3:
                print('Retrying download')
                self.download_file_cos(key, folder=folder, attempts=attempts+1)

    def get_annotations(self):
        try:
            return json.loads(self.cos.get_object(Bucket=self.bucket, Key='_annotations.json')['Body'].read())
        except Exception as e:
            print('Exception', e)
    
    def downloadAll(self):
        image_folder = os.path.join(os.getcwd(), 'images')
        if not os.path.exists(image_folder):
                os.mkdir(image_folder)

        for image in tqdm(self.get_annotations()['annotations'].keys()):
            self.download_file_cos(image, image_folder)
    
    def getDataset(self, transform=None,train_test=None,percentage_train=0.8,random_state=0,degrees=5):
        dataset=getDataset(annotations=self.get_annotations(), bucket_name=self.bucket, transform=transform,train_test=train_test,percentage_train=percentage_train,random_state=random_state,degrees=degrees)
        return dataset
    
    def uploadModel(self, key, details={}):
        self.upload_file_cos(key)
        details['ready'] = True
        details['filename'] = key
        details['filesize'] = os.stat(key).st_size
        
        self.report(model=details)

    def downloadModel(self, folder=os.getcwd()):
        headers = {
            'Authorization': 'Bearer ' + self.token,
            'Content-Type': 'application/json'
        }
        
        x = requests.post(self.base_url + '/api/modelDetails', headers=headers)

        if x.ok:
            model = json.loads(x.text)
            
            if 'external' in model:
                raise Exception("External models support coming soon.")
            elif 'ready' not in model or not model['ready']:
                raise Exception("Model isn't ready")

            self.download_file_cos(model['filename'], folder)

            return model

        else:
            raise Exception('Failed to get Model Details')

    def upload_file_cos(self, filename):
        try:
            res=self.cos.upload_file(filename, self.bucket, os.path.basename(filename))
        except Exception as e:
            print('Exception', e)
        else:
            print('File Uploaded')
    
    def report(self, started=None, completed=None, url=None, parameters=None, accuracy=None, model=None):
        data = {}

        if started is not None:
            data['started'] = started
        
        if completed is not None:
            data['completed'] = completed

        if parameters is not None:
            data['parameters'] = parameters

        if accuracy is not None:
            data['accuracy'] = accuracy

        if model is not None:
            data['model'] = model

        if url is not None:
            data['url'] = url
        
        if len(data) == 0:
            raise Exception('Nothing to report')

        headers = {
            'Authorization': 'Bearer ' + self.token,
            'Content-Type': 'application/json'
        }
        
        x = requests.post(self.base_url + '/api/report', headers=headers, data=json.dumps(data, default=default))

        return x

    def ping(self):
        headers = {
            'Authorization': 'Bearer ' + self.token,
            'Content-Type': 'application/json'
        }
        
        x = requests.post(self.base_url + '/api/ping', headers=headers)

        return x


def getDataset(annotations, bucket_name, transform,train_test=None,percentage_train=0.8,random_state=0,degrees=5):
    from torch.utils.data import Dataset
    import pandas as pd
    from PIL import Image
    



    class Dataset(Dataset):

        def __init__(self, annotations, bucket_name, transform,train_test=None,percentage_train=percentage_train,random_state=random_state,degrees=degrees):
            
            self.train_test=train_test
            self.transform=transform
            if self.transform is None:
                print("defult transform for pretrained model resnet18")
                self.setDefaultTransform()
                
            self.train_test=train_test
            self.degrees=degrees

            problem_type=annotations['type']
            
            if problem_type != 'classification':
                raise Exception('Can only get Dataset of Classification models')
            
            labels_types=annotations['labels']
            
            labels_filename=os.path.join(os.getcwd(),bucket_name+"_lables.csv")

            if os.path.exists(labels_filename):
                self.data=pd.read_csv(labels_filename)
            else:
                data_={'label':[],'y':[],'file_name':[],'key':[]}
                for key,label_dict in annotations['annotations'].items():                
                    label=label_dict[0]['label']

                    data_['label'].append(label)
                    data_['y'].append(labels_types.index(label))
                    data_['key'].append(key)
                    data_['file_name'].append(os.path.join(os.getcwd(),'images',key))

                self.data=pd.DataFrame(data_)
                #train set
            if self.train_test=='train':
                print("this is the training set")
                self.data=self.data.sample(frac=percentage_train,random_state=random_state).reset_index()
                labels_filename=os.path.join(os.getcwd(),bucket_name+"_train_lables.csv")
                #self.data.to_csv(labels_filename)
            #test set    
            if self.train_test=='test':
                print("this is the test set")
                labels_filename=os.path.join(os.getcwd(),bucket_name+"_test_lables.csv")    
                temp=self.data.sample(frac=percentage_train,random_state=random_state) 
                self.data=self.data.drop(temp.index).reset_index()           
                #self.data.to_csv(labels_filename)
                
                
            
            self.n_classes=len(self.data['y'].unique())

        def setDefaultTransform(self):
            from torchvision import transforms
            

            mean = [0.485, 0.456, 0.406]
            std = [0.229, 0.224, 0.225]
            if self.train_test=='train':
            # Data augmentation and normalization for training
                self.transform = transforms.Compose([
                    transforms.Resize((224, 224)),
                    transforms.RandomHorizontalFlip(),
                    transforms.RandomRotation(degrees=5),
                    transforms.ToTensor(),
                    transforms.Normalize(mean, std)])
            else:
            # Data augmentation and normalization for training
                self.transform = transforms.Compose([
                    transforms.Resize((224, 224)),
                    transforms.ToTensor(),
                    transforms.Normalize(mean, std)])

        def __len__(self):
            return self.data.shape[0]

        def __getitem__(self, idx):
            image=Image.open(self.data.loc[idx, 'file_name'])
            y=self.data.loc[idx,'y']
            if self.transform:
                image= self.transform(image)
            
            return image, y

    dataset=Dataset(annotations=annotations, bucket_name=bucket_name, transform=transform,train_test=train_test,percentage_train=percentage_train,random_state=random_state,degrees=degrees)
    return dataset