from boto3 import session
from botocore.client import Config
import os
ACCESS_ID = os.getenv('ACCESS_ID', '')
SECRET_KEY = os.getenv('SECRET_KEY', '')

LOCAL_WEIGHTS_DIR = 'weights'
DO_SPACE='sing'
def upload_to_do(file_path):
    
    boto_session=session.Session()
    client = boto_session.client('s3',
                            region_name='nyc3',
                            endpoint_url='https://nyc3.digitaloceanspaces.com',
                            aws_access_key_id=ACCESS_ID,
                            aws_secret_access_key=SECRET_KEY)
    filename_only = os.path.basename(file_path)
    # Upload a file to your Space
    response=client.upload_file(file_path, 'sing', filename_only)

    
    return response
    
def download_from_do(file_key):
    boto_session = session.Session()
    client = boto_session.client('s3',
                                 region_name='nyc3',
                                 endpoint_url='https://nyc3.digitaloceanspaces.com',
                                 aws_access_key_id=ACCESS_ID,
                                 aws_secret_access_key=SECRET_KEY)
    
      # Ensure the downloads directory exists
    downloads_dir = 'downloads'
    if not os.path.exists(downloads_dir):
        os.makedirs(downloads_dir)
    
    # Set the full local path for the download
    full_local_path = os.path.join(downloads_dir, file_key)
    
    # Download the file from your Space
    client.download_file('sing', file_key, full_local_path)
    
    # Verify the download
    if os.path.exists(full_local_path):
        print(f"File downloaded successfully to {full_local_path}")
        return full_local_path
    else:
        print("Download failed.")
        return None

def get_local_models(prefix):
    """Get list of model files starting with prefix in the local directory."""
    models = [f for f in os.listdir(LOCAL_WEIGHTS_DIR) if f.startswith(prefix) and f.endswith('.pth')]
    return models

def get_do_models(client, prefix):
    """Get list of model files starting with prefix in the DO space."""
    paginator = client.get_paginator('list_objects')
    page_iterator = paginator.paginate(Bucket=DO_SPACE, Prefix=prefix)

    models = []
    for page in page_iterator:
        models.extend([obj['Key'] for obj in page['Contents'] if obj['Key'].endswith('.pth')])
    return models

def sync_missing_models(client, local_models, do_models):
    """Download missing model files from DO space."""
    missing_models = set(do_models) - set(local_models)
    print('missing models:',missing_models)
    for model in missing_models:
        client.download_file(DO_SPACE, model, os.path.join(LOCAL_WEIGHTS_DIR, model))
        print(f"Downloaded {model} from DO space to local weights directory.")


def list_models(email_prefix):
    
    

    #ensure_local_directory_exists()
    local_models = get_local_models(email_prefix)

    # Initialize DO S3 client
    boto_session = session.Session()
    client = boto_session.client('s3',
                                 region_name='nyc3',
                                 endpoint_url='https://nyc3.digitaloceanspaces.com',
                                 aws_access_key_id=ACCESS_ID,
                                 aws_secret_access_key=SECRET_KEY)

    
    do_models = get_do_models(client, email_prefix)
    sync_missing_models(client, local_models, do_models)

    # Return the updated list of local models after syncing
    updated_local_models = get_local_models(email_prefix)
    print(updated_local_models)
    #return jsonify(updated_local_models)



def download_from_do_with_prefix(prefix):
    boto_session = session.Session()
    client = boto_session.client('s3',
                                 region_name='nyc3',
                                 endpoint_url='https://nyc3.digitaloceanspaces.com',
                                 aws_access_key_id=ACCESS_ID,
                                 aws_secret_access_key=SECRET_KEY)
    
    # Ensure the downloads directory exists
    downloads_dir = 'downloads'
    if not os.path.exists(downloads_dir):
        os.makedirs(downloads_dir)
    
    # List objects in the Space with the specified prefix
    response = client.list_objects(Bucket='sing', Prefix=prefix)
    print(response)
    downloaded_files = []
    if 'Contents' in response:
        for obj in response['Contents']:
            file_key = obj['Key']
            
            # Set the full local path for the download
            full_local_path = os.path.join(downloads_dir, os.path.basename(file_key))
            
            # Download the file from your Space
            client.download_file('sing', file_key, full_local_path)
            
            # Verify the download and add to the list if successful
            if os.path.exists(full_local_path):
                print(f"File downloaded successfully to {full_local_path}")
                downloaded_files.append(full_local_path)
            else:
                print(f"Download failed for {file_key}.")
    else:
        print("No files found with the specified prefix.")
    
    return downloaded_files if downloaded_files else None
# Initiate session

def ensure_model_in_weights_dir(model_name):
    weights_dir = 'weights'
    model_path = os.path.join(weights_dir, model_name)

    # Check if the model already exists
    if os.path.exists(model_path):
        print(f"Model {model_name} already exists in {weights_dir}.")
        return True

    # If the model does not exist, attempt to download it
    print(f"Model {model_name} not found in {weights_dir}. Attempting to download...")

    # Initialize a session using DigitalOcean Spaces
    boto_session = session.Session()
    client = boto_session.client('s3',
                                 region_name='nyc3',
                                 endpoint_url='https://nyc3.digitaloceanspaces.com',
                                 aws_access_key_id=ACCESS_ID,
                                 aws_secret_access_key=SECRET_KEY)

    # Ensure the weights directory exists
    if not os.path.exists(weights_dir):
        os.makedirs(weights_dir)

    # Attempt to download the model file
    try:
        client.download_file('sing', f"{model_name}", model_path)
        print(f"Model {model_name} downloaded successfully to {model_path}.")
        return True
    except Exception as e:
        print(f"Failed to download {model_name}: {e}")
        return False