import argparse
import json
import os
import secrets
from werkzeug.security import generate_password_hash
from duckduckgo_search import DDGS
import requests
import shutil

def set_login_info(username, password):
    salt = secrets.token_hex(16)
    password_hash = generate_password_hash(password + salt)
    key = secrets.token_hex(32)

    login_info = {
        "login_info": {
            "username": username,
            "password_hash": password_hash,
            "salt": salt
        },
        "key": key
    }

    os.makedirs('db', exist_ok=True)
    with open('db/keys.json', 'w') as f:
        json.dump(login_info, f, indent=4)

def get_next_batch_number():
    existing_files = os.listdir('images')
    batch_numbers = [int(f.split('_')[0]) for f in existing_files if '_' in f]
    return max(batch_numbers, default=0) + 1

def download_images(query, total_images=100, batch_size=10):
    os.makedirs('images', exist_ok=True)
    ddgs = DDGS()
    results = ddgs.images(query, max_results=total_images)
    batch_number = get_next_batch_number()
    downloaded_images = 0
    valid_extensions = {'jpg', 'jpeg', 'png', 'gif'}

    for i, result in enumerate(results):
        image_url = result['image']
        try:
            response = requests.get(image_url)
            if response.status_code == 200:
                ext = image_url.split('.')[-1].lower()
                if ext not in valid_extensions:
                    ext = 'jpg'
                filename = f'{batch_number:04d}_{(downloaded_images % batch_size) + 1:05d}.{ext}'
                with open(f'images/{filename}', 'wb') as f:
                    f.write(response.content)
                downloaded_images += 1
            else:
                print(f"Failed to download image from {image_url}. Status code: {response.status_code}")
        except requests.RequestException as e:
            print(f"An error occurred while downloading image from {image_url}: {e}")
        except Exception as e:
            print(f"An unexpected error occurred: {e}")

        if downloaded_images > 0 and downloaded_images % batch_size == 0:
            batch_number += 1

    print(f"Downloaded {downloaded_images} images for query '{query}'.")

def clear_cache():
    cache_folder = 'cache'
    if os.path.exists(cache_folder):
        shutil.rmtree(cache_folder)
        print(f"Cache directory '{cache_folder}' removed successfully.")
    else:
        print(f"Cache directory '{cache_folder}' does not exist.")

def main():
    parser = argparse.ArgumentParser(
        description='Set username and password for the application and optionally download images using DuckDuckGo.',
        formatter_class=argparse.ArgumentDefaultsHelpFormatter
    )

    # Add an argument to enter user setting mode
    parser.add_argument('--user-setting', action='store_true', help='Enter user setting mode to set username and password.')

    # Create a userset argument group
    userset_group = parser.add_argument_group('userset', 'Set username and password')
    userset_group.add_argument('-u', '--username', type=str, help='The username to set.')
    userset_group.add_argument('-p', '--password', type=str, help='The password to set.')

    parser.add_argument('-q', '--query', type=str, help='The search query to download images for testing.')
    parser.add_argument('-t', '--total_images', type=int, default=10, help='The total number of images to download.')
    parser.add_argument('-b', '--batch_size', type=int, default=5, help='The number of images to download in each batch.')
    parser.add_argument('-c', '--clear-cache', action='store_true', help='Remove the cache directory.')

    args = parser.parse_args()

    if args.user_setting:
        if not args.username or not args.password:
            parser.error('--user-setting requires --username and --password.')
        set_login_info(args.username, args.password)
        print(f"Username and password set successfully. Key generated and saved in db/keys.json.")
    elif args.clear_cache:
        clear_cache()
    elif args.query:
        download_images(args.query, args.total_images, args.batch_size)
    else:
        parser.print_help()

if __name__ == '__main__':
    main()