import requests
import os
import boto3
from boto3.exceptions import ClientError
from mimetypes import MimeTypes

class LucidicAPI:
    def __init__(self, api_key: str):
        self.api_key = api_key
        self.headers = {"Authorization": f"Api-Key {self.api_key}"}
        self.base_url = "https://dashboard.lucidic.ai/demo/api/v1"
        self.endpoints = {
            "verifyAPIKey": "verifyAPIkey",
            "initializejob": "initializejob",
            "assumeAWSS3Role": "getS3creds",
        }

    def _makeRequest(self, endpoint, params=None):
        try:
            url = f'{self.base_url}/{self.endpoints[endpoint]}'
            response = requests.get(
                url,
                headers=self.headers,
                params=params
            )
            response.raise_for_status()
            return response
        except requests.exceptions.RequestException as e:
            print(f"Error during API Call, status {response.status_code}: {e}")
            raise
        except KeyError as e:
            print(f"Error: Specified endpoint not found.")
            raise


    def _tryCreateS3Bucket(self, creds, bucket_name):
        s3_client = boto3.client(
            's3',
            aws_access_key_id=creds['AccessKeyId'],
            aws_secret_access_key=creds['SecretAccessKey'],
            aws_session_token=creds['SessionToken'],
            region_name='us-west-2'
        )

        try:
            s3_client.head_bucket(Bucket=bucket_name)
            bucket_exists = True
            print(f"S3 bucket '{bucket_name}' already exists.")
            return False
        except ClientError as e:
            error_code = int(e.response['Error']['Code'])
            if error_code == 404:
                print(f"S3 bucket '{bucket_name}' does not exist. Creating the bucket...")
            else:
                print(f"Error checking bucket existence: {e}")
                raise
        
        try:
            create_bucket_params = {
                'Bucket': bucket_name,
                'CreateBucketConfiguration': {
                    'LocationConstraint': 'us-west-2'
                    }
                }
            s3_client.create_bucket(**create_bucket_params)
            print(f"S3 bucket '{bucket_name}' created successfully.")
            return True
        except Exception as e:
            print(f"Error creating bucket: {e}")
            raise


    def _queueFiles(self, creds, bucket_name, pathToDataFolder):
        s3_client = boto3.client(
            's3',
            aws_access_key_id=creds['AccessKeyId'],
            aws_secret_access_key=creds['SecretAccessKey'],
            aws_session_token=creds['SessionToken'],
            region_name='us-west-2'
        )
        mime = MimeTypes()
        for root, _, files in os.walk(pathToDataFolder):
            for file in files:
                file_path = os.path.join(root, file)
                key = os.path.relpath(file_path, pathToDataFolder)  # Key is relative path in S3
                mime_type, _ = mime.guess_type(file_path)

                # Check if the file is an image or video
                if mime_type and (mime_type.startswith('image/') or mime_type.startswith('video/')):
                    try:
                        # Check if the file already exists in the bucket
                        try:
                            s3_client.head_object(Bucket=bucket_name, Key=key)
                            print(f"File already exists: s3://{bucket_name}/{key}. Skipping upload.")
                            continue
                        except ClientError as e:
                            if int(e.response['Error']['Code']) == 404:
                                # File does not exist, proceed with upload
                                pass
                            else:
                                raise

                        # Upload the file
                        s3_client.upload_file(file_path, bucket_name, key)
                        print(f"Uploaded: {file_path} -> s3://{bucket_name}/{key}")
                    except Exception as e:
                        print(f"Error uploading {file_path}: {e}")


    def verifyAPIKey(self):
        response = self._makeRequest('verifyAPIKey')
        return response.json().get('project', None)
    

    def startJob(self, pathToDataFolder):
        print(f"Starting job, verifying API key...")
        project = self.verifyAPIKey()
        assert project is not None
        print(f"API Key verified for project {project}!")

        print(f"Initializing job...")
        jobInitializationResponseJSON = self._makeRequest('initializejob').json() 
        print(f"{jobInitializationResponseJSON}")
        jobID = jobInitializationResponseJSON.get('jobID', None)
        assert jobID is not None
        print(f"Job Initialized with jobID: {jobID}")

        print(f"Issuing temporary AWS S3 Credentials...")
        creds = self._makeRequest('assumeAWSS3Role').json()
        assert 'AccessKeyId' in creds and 'SecretAccessKey' in creds and 'Expiration' in creds
        print(f"Temporary AWS Credentials Issued!")

        print(f"Creating AWS S3 Bucket...")
        bucketName = project
        if len(project) > 25:
            bucketName = project[:25]
        bucketName += '.' + jobID
        self._tryCreateS3Bucket(creds, bucketName)

        print(f"Uploading files...")
        self._queueFiles(creds, bucketName, pathToDataFolder)

    







