import requests
import time
import os
import zipfile
import io
import argparse
import sys

def get_api_key():
    """
    Retrieves the API key from the environment variable.
    """
    api_key = os.getenv('CODESENT_API_KEY')
    if not api_key:
        print("Error: CODESENT_API_KEY environment variable is not set.")
        sys.exit(1)
    return api_key

def zip_directory(directory_path):
    """
    Compresses the specified directory into a ZIP file stored in memory.
    
    :param directory_path: Path to the directory to be zipped.
    :return: BytesIO object containing the ZIP file.
    """
    if not os.path.isdir(directory_path):
        raise Exception(f"The provided path '{directory_path}' is not a directory or does not exist.")
    
    zip_buffer = io.BytesIO()
    with zipfile.ZipFile(zip_buffer, 'w', zipfile.ZIP_DEFLATED) as zip_file:
        for root, dirs, files in os.walk(directory_path):
            for file in files:
                file_path = os.path.join(root, file)
                # Compute the archive name by removing the base directory from the file path
                arcname = os.path.relpath(file_path, start=directory_path)
                zip_file.write(file_path, arcname)
    zip_buffer.seek(0)
    return zip_buffer

def upload_zip(zip_buffer, api_base_url, headers):
    """
    Uploads the ZIP file to the API and returns the proxy UUID.
    
    :param zip_buffer: BytesIO object containing the ZIP file.
    :param api_base_url: Base URL of the API.
    :param headers: Dictionary containing the request headers.
    :return: proxy_uuid as a string.
    """
    url = f'{api_base_url}/upload'
    files = {'file': ('proxy.zip', zip_buffer, 'application/zip')}
    response = requests.post(url, headers=headers, files=files)
    
    if response.status_code == 200:
        proxy_uuid = response.json()['proxy_uuid']
        print(f"File uploaded successfully. Proxy UUID: {proxy_uuid}")
        return proxy_uuid
    else:
        error_message = response.json().get('error', 'Unknown error')
        raise Exception(f"File upload failed: {error_message}")

def validate_proxy(proxy_uuid, api_base_url, headers):
    """
    Initiates validation for the specified proxy and returns the task UUID.
    
    :param proxy_uuid: UUID of the uploaded proxy.
    :param api_base_url: Base URL of the API.
    :param headers: Dictionary containing the request headers.
    :return: task_uuid as a string.
    """
    url = f'{api_base_url}/{proxy_uuid}/validate'
    response = requests.post(url, headers=headers)
    
    if response.status_code == 200:
        task_uuid = response.json()['task_uuid']
        print(f"Validation initiated. Task UUID: {task_uuid}")
        return task_uuid
    else:
        error_message = response.json().get('error', 'Unknown error')
        raise Exception(f"Validation initiation failed: {error_message}")

def check_status(proxy_uuid, task_uuid, api_base_url, headers):
    """
    Checks the current status of the validation task.
    
    :param proxy_uuid: UUID of the uploaded proxy.
    :param task_uuid: UUID of the validation task.
    :param api_base_url: Base URL of the API.
    :param headers: Dictionary containing the request headers.
    :return: Current status as a string.
    """
    url = f'{api_base_url}/{proxy_uuid}/{task_uuid}/status'
    response = requests.post(url, headers=headers)
    
    if response.status_code == 200:
        status = response.json()['status']
        print(f"Current status: {status}")
        return status
    else:
        error_message = response.json().get('error', 'Unknown error')
        raise Exception(f"Status check failed: {error_message}")

def get_results(proxy_uuid, task_uuid, api_base_url, headers):
    """
    Retrieves the results of the completed validation task.
    
    :param proxy_uuid: UUID of the uploaded proxy.
    :param task_uuid: UUID of the validation task.
    :param api_base_url: Base URL of the API.
    :param headers: Dictionary containing the request headers.
    :return: Dictionary containing the validation results.
    """
    url = f'{api_base_url}/{proxy_uuid}/{task_uuid}/results'
    response = requests.post(url, headers=headers)
    
    if response.status_code == 200:
        results = response.json()
        online_report = results['online_report']
        print(f"Validation Results is available at {online_report}")
        return results
    else:
        error_message = response.json().get('error', 'Unknown error')
        raise Exception(f"Results retrieval failed: {error_message}")

def main():
    # Parse command-line arguments
    parser = argparse.ArgumentParser(description='Apigee Buddy API Client')
    parser.add_argument('proxy_directory', default='.', help='Path to the directory containing the proxy configuration')
    parser.add_argument('--api-url', default='https://codesent.io/api/scan/v1', help='Base URL of the API')
    
    args = parser.parse_args()
    
    # Retrieve the API key from environment variables
    api_key = get_api_key()
    
    headers = {
        'Authorization': f'Bearer {api_key}'
    }
    
    try:
        print("Zipping the proxy directory...")
        zip_buffer = zip_directory(args.proxy_directory)
        print("Proxy directory zipped successfully.")
        
        print("Uploading the ZIP file...")
        proxy_uuid = upload_zip(zip_buffer, args.api_url, headers)
        
        print("Initiating validation...")
        task_uuid = validate_proxy(proxy_uuid, args.api_url, headers)
        
        # Periodically check the status of the validation task
        while True:
            status = check_status(proxy_uuid, task_uuid, args.api_url, headers)
            if status.lower() == 'done':
                print("Validation completed successfully.")
                break
            elif status.lower() in ['failed', 'error']:
                print("Validation completed with errors.")
                return
            time.sleep(5)  # Wait for 5 seconds before the next status check
        
        print("Retrieving validation results...")
        get_results(proxy_uuid, task_uuid, args.api_url, headers)
    
    except Exception as e:
        print(f"An error occurred: {e}")
        sys.exit(1)

if __name__ == "__main__":
    main()