import os
import zipfile
import tempfile
from concurrent.futures import ThreadPoolExecutor, as_completed

from image_functions import generate_image, is_valid_image
from r2_uploader import upload_image_to_r2
from mongo_logger import create_log, update_log_status

def process_zip_and_generate_images(image_path, template_path, category, size, quality, user_prompt, sentiment):
    try:
        image_files = []
        image_files.append((os.path.basename(template_path), template_path.name, os.path.basename(image_path.name), image_path.name))
        results = process_image_files(image_files, category, size, quality, user_prompt, sentiment)
        all_urls = [url for entry in results for url in entry["urls"]]

        new_base = "https://adgen.statics.lookfinity.net/"

        modified_urls = [new_base + url.rsplit('/', 1)[-1] for url in all_urls]
        return modified_urls

    except Exception as e:
        log_error(f"Global error in zip processing: {e}")
        return [], [f"Global error: {e}"]


def extract_zip_file(zip_path):

    temp_dir = tempfile.TemporaryDirectory()
    with zipfile.ZipFile(zip_path, 'r') as zip_ref:
        zip_ref.extractall(temp_dir.name)

    return temp_dir


def get_valid_image_files(temp_dir):

    valid_files = []
    zip_contents = os.listdir(temp_dir.name)

    for file in zip_contents:
        file_path = os.path.join(temp_dir.name, file)
        if is_valid_image(file) and "__MACOSX" not in file:
            valid_files.append((file, file_path))

    return valid_files


def process_image_files(image_files, category, size, quality, user_prompt, sentiment):
    final_results = []
    logs = []

    with ThreadPoolExecutor(max_workers=5) as executor:
        futures = []

        for template_name, template_path, file_name, file_path in image_files:
            log_id = create_log(category, file_name)
            logs.append(f"Processing {file_name}")

            futures.append(
                executor.submit(
                    process_single_image,
                    template_path,
                    file_name,
                    file_path,
                    category,
                    size,
                    quality,
                    user_prompt,
                    sentiment,
                    log_id,
                    logs
                )
            )

        for future in as_completed(futures):
            result = future.result()
            if result:
                final_results.append(result)

    return final_results


def process_single_image(template_path, file_name, file_path, category, size, quality, user_prompt, sentiment, log_id, logs):
    try:
        # base64_img = encode_image_to_base64(file_path)
        # logs.append(f"Prompts generated for {file_name}")
        #
        # prompts = get_prompts(base64_img, category, user_prompt, sentiment, negative_prompt)
        # logs.append(f"Generating images for {file_name}")

        image_urls = generate_images_from_prompts(template_path, file_path, size, quality, category, sentiment, user_prompt)

        status = "completed" if image_urls else "failed"
        message = "Completed successfully" if image_urls else "No images generated"
        update_log_status(log_id, status, urls=image_urls, message=message)

        if image_urls:
            return {
                "file_name": file_name,
                "urls": image_urls
            }
        return None

    except Exception as e:
        logs.append(f"Processing failed for {file_name}: {e}")
        update_log_status(log_id, "failed", urls=[], message=str(e))
        return None


def generate_images_from_prompts(template_path, file_path, size, quality, category, sentiment, user_prompt):
    image_urls = []

    for _ in range(10):
        try:
            img_url = generate_image(template_path, file_path, size, quality, category, sentiment, user_prompt)

            s3_url = upload_image_to_r2(img_url)
            image_urls.append(s3_url)

        except Exception as e:
            log_error(f"Image generation failed: {e}")

    return image_urls


def log_error(message):
    print(f"ERROR: {message}")