import os
import json
import sys
import subprocess
import requests
from concurrent.futures import ThreadPoolExecutor
import time

# Configuration
MAX_RETRIES = 3  # Maximum retry attempts for downloading files
NUM_WORKERS = 5  # Default number of workers
TIMEOUT = 30  # Timeout in seconds for each file download
PROCESSING_DIR = "/processing"  # Directory to store zip, .txt, and .json files

# Function to download the JSON file from the specified URL
def download_json(bucket_name, download_directory):
    try:
        url = f"https://livetimelapse.com.au/ai/chatgpt/internal/b2/servers/www.php?bucket_name={bucket_name}"
        response = requests.get(url, stream=True)
        response.raise_for_status()  # Raise an exception for HTTP errors

        # Save JSON to local file with new naming convention
        json_filename = f"{bucket_name}-files.json"
        json_path = os.path.join(PROCESSING_DIR, json_filename)
        with open(json_path, "w") as json_file:
            json.dump(response.json(), json_file, indent=4)

        print(f"JSON file downloaded and saved to: {json_path}")
        return json_path
    except Exception as e:
        print(f"Error downloading JSON file: {e}")
        sys.exit(1)

# Function to load the JSON file into memory
def load_json_file(json_path):
    try:
        with open(json_path, "r") as json_file:
            return json.load(json_file)
    except Exception as e:
        print(f"Error loading JSON file: {e}")
        sys.exit(1)

# Function to create directories and save a file
def download_file(file_info, base_dir, stats, image_index, total_images):
    file_path = file_info["path"]
    cdn_url = file_info["cdn"]

    print(f"\nImage {image_index}/{total_images}")
    print(f"Downloading: {cdn_url}")

    local_path = os.path.join(base_dir, file_path)
    os.makedirs(os.path.dirname(local_path), exist_ok=True)

    if os.path.exists(local_path):
        print(f"File already exists, skipping: {local_path}\n")
        stats["skipped"] += 1
        return

    for attempt in range(1, MAX_RETRIES + 1):
        try:
            response = requests.get(cdn_url, stream=True, timeout=TIMEOUT)
            response.raise_for_status()

            with open(local_path, "wb") as file:
                for chunk in response.iter_content(chunk_size=8192):
                    file.write(chunk)

            print(f"Saved to: {local_path}\n")
            stats["downloaded"] += 1
            return
        except requests.exceptions.RequestException as e:
            print(f"Attempt {attempt} failed for {cdn_url}: {e}")
            if attempt < MAX_RETRIES:
                time.sleep(2 ** attempt)  # Exponential backoff
            else:
                print(f"Failed to download {cdn_url} after {MAX_RETRIES} attempts.\n")
                stats["failed"] += 1
                return

# Function to compress files into a ZIP
def compress_with_7zip(directory, zip_name):
    try:
        zip_path = os.path.join(PROCESSING_DIR, zip_name)
        command = f'7z a "{zip_path}" "{directory}/*"'
        print(f"Compressing files into: {zip_path}")
        result = subprocess.run(command, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)

        # Save process output to a log file
        log_file = zip_path.replace(".zip", "-log.txt")
        with open(log_file, "w") as log:
            log.write(result.stdout)
            log.write("\n\nErrors (if any):\n")
            log.write(result.stderr)

        print(f"Compression completed: {zip_path}")
        return zip_path, log_file
    except subprocess.CalledProcessError as e:
        print(f"Error compressing {directory}: {e}")
        sys.exit(1)

# Function to process zipping for multiple folders
def create_multiple_zips(bucket_name, bucket_dir):
    # Generate modified bucket name
    modified_bucket_name = bucket_name.replace("-", "_")

    # Paths for ZIP files
    originals_zip = f"{bucket_name}-originals.zip"
    www_zip = f"{bucket_name}-www.zip"
    logs = []

    # Combine 'originals' and 'fullsize' into one ZIP while retaining structure
    originals_paths = []
    originals_dir = os.path.join(bucket_dir, "originals")
    fullsize_dir = os.path.join(bucket_dir, "fullsize")
    if os.path.exists(originals_dir):
        originals_paths.append(originals_dir)
    if os.path.exists(fullsize_dir):
        originals_paths.append(fullsize_dir)

    if originals_paths:
        try:
            # Use 7z to add each directory into the ZIP while preserving structure
            command = f'7z a "{os.path.join(PROCESSING_DIR, originals_zip)}" ' + " ".join([f'"{path}"' for path in originals_paths])
            print(f"Compressing originals/fullsize into: {originals_zip}")
            result = subprocess.run(command, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)

            # Save process output to a log file
            log_file = os.path.join(PROCESSING_DIR, originals_zip.replace(".zip", "-log.txt"))
            with open(log_file, "w") as log:
                log.write(result.stdout)
                log.write("\n\nErrors (if any):\n")
                log.write(result.stderr)
            logs.append(log_file)

            print(f"Compression completed for: {originals_zip}")
        except subprocess.CalledProcessError as e:
            print(f"Error compressing originals/fullsize: {e}")
            sys.exit(1)

    # Compress the modified bucket folder
    modified_bucket_path = os.path.join(bucket_dir, modified_bucket_name)
    if os.path.exists(modified_bucket_path):
        zip_path, log_path = compress_with_7zip(modified_bucket_path, www_zip)
        logs.append(log_path)

    return [os.path.join(PROCESSING_DIR, originals_zip), os.path.join(PROCESSING_DIR, www_zip)], logs

# Function to move files to Google Drive in a screen session
def move_to_gdrive_with_screen(file_path, rclone_target):
    try:
        time.sleep(30)  # Delay to ensure ZIP file is finalized
        log_file = f"/tmp/{os.path.basename(file_path)}.log"
        session_name = f"rclone_{os.path.basename(file_path)}"
        command = f'screen -dmS {session_name} sh -c "rclone move {file_path} {rclone_target} -P > {log_file} 2>&1"'
        print(f"Starting rclone move for {file_path} in screen session: {session_name}, logging to {log_file}")
        os.system(command)
    except Exception as e:
        print(f"Error starting rclone move command for {file_path}: {e}")
        sys.exit(1)

# Function to copy files to Google Drive
def copy_to_gdrive(file_path, rclone_target):
    try:
        command = f'rclone copy {file_path} {rclone_target} -P'
        print(f"Copying {file_path} to {rclone_target}...")
        os.system(command)
    except Exception as e:
        print(f"Error copying {file_path} to {rclone_target}: {e}")
        sys.exit(1)

# Main function to download JSON, process files, and compress
def main(num_workers, bucket_name, download_directory, rclone_target):
    os.makedirs(download_directory, exist_ok=True)
    os.makedirs(PROCESSING_DIR, exist_ok=True)
    summary = {"downloaded": 0, "skipped": 0, "failed": 0, "download_time": 0}

    # Step 1: Download JSON
    json_path = download_json(bucket_name, PROCESSING_DIR)

    # Step 2: Load JSON and prepare bucket directory
    file_list = load_json_file(json_path)
    bucket_dir = os.path.join(download_directory, bucket_name)
    os.makedirs(bucket_dir, exist_ok=True)

    # Step 3: Download files
    print(f"Downloading {len(file_list)} images to: {bucket_dir}")
    with ThreadPoolExecutor(max_workers=num_workers) as executor:
        for idx, file_info in enumerate(file_list, start=1):
            executor.submit(download_file, file_info, bucket_dir, summary, idx, len(file_list))
    print(f"Finished downloading files for bucket: {bucket_name}")

    # Step 4: Create multiple ZIPs
    zip_files, log_files = create_multiple_zips(bucket_name, bucket_dir)

    # Step 5: Move ZIPs to Google Drive using screen
    for zip_file in zip_files:
        if os.path.exists(zip_file):
            move_to_gdrive_with_screen(zip_file, rclone_target)

    # Step 6: Copy log and JSON files to Google Drive
    for log_file in log_files:
        if os.path.exists(log_file):
            copy_to_gdrive(log_file, rclone_target)
    copy_to_gdrive(json_path, rclone_target)

    # Step 7: Remove directories
    print(f"Cleaning up files in {bucket_dir}")
    subprocess.run(f"rm -Rf {bucket_dir}", shell=True, check=True)

if __name__ == "__main__":
    if len(sys.argv) != 5:
        print("Usage: python3 script.py <num_workers> <bucket_name> <download_directory> <rclone_target>")
        sys.exit(1)

    num_workers = int(sys.argv[1])
    bucket_name = sys.argv[2]
    download_directory = sys.argv[3]
    rclone_target = sys.argv[4]
    main(num_workers, bucket_name, download_directory, rclone_target)