import nest_asyncio
import os
from dataclasses import dataclass
from pathlib import Path

from custom_apps.bingimagecreator.BingImageCreator import *
from custom_apps.utils import cookie_manager
from config import rest_config


@dataclass
class DallEArgument:
    """
    U : Auth cookie from browser
    cookie_file : File containing auth cookie
    prompt: Prompt to generate images for
    output_dir: Output directory
    download_count: Number of images to download, value must be less than five
    debug_file: Path to the file where debug information will be written.
    quiet: Disable pipeline messages
    asyncio: Run ImageGen using asyncio
    version: Print the version number
    """
    U = cookie_manager.get_cookie()
    prompt: str
    cookie_file: str|None = None
    output_dir: str = os.path.join(rest_config.remote_folder,"dalle","img")
    download_count: int = 1
    debug_file: str|None = None
    quiet: bool = False
    asyncio: bool = True
    version: bool = False
    

class Counter:
    """
    생성된 이미지 카운트
    """
    
    def __init__(self):
        self.counter = 0
        self.error_messages = []
        
    def reset(self):
        self.counter = 0
        self.error_messages = []
        
    def count(self):
        self.counter += 1
        
    def get_counter(self):
        return self.counter

    def add_error_messages(self,error):
        self.error_messages.append(error)
        
    def get_error_messages(self):
        return self.error_messages
    

def dalle3_generate_image(args):   
    
    nest_asyncio.apply()
    
    counter = Counter()
    
    Path(args.output_dir).mkdir(parents=True, exist_ok=True)
    
    if args.version:
        print(pkg_resources.get_distribution("BingImageCreator").version)
        sys.exit()

    # Load auth cookie
    cookie_json = None
    if args.cookie_file is not None:
        with contextlib.suppress(Exception):
            with open(args.cookie_file, encoding="utf-8") as file:
                cookie_json = json.load(file)

    if args.U is None and args.cookie_file is None:
        raise Exception("Could not find auth cookie")

    if args.download_count > 4:
        raise Exception("The number of downloads must be less than five")


    asyncio.run(
        async_image_gen(
            prompt=args.prompt,
            download_count=args.download_count,
            output_dir=args.output_dir,
            counter=counter,
            u_cookie=args.U,
            debug_file=args.debug_file,
            quiet=args.quiet,
            all_cookies=cookie_json,
        ),
    )
    
    return counter

