# -*- coding: utf-8 -*-

import vertexai
import os
from vertexai.preview.vision_models import ImageGenerationModel

from const import TEMP_FOLDER
from main_rest.app.utils.parsing_utils import prompt_to_filenames
from main_rest.app.utils.date_utils import D
from utils.custom_sftp import sftp_client
from config import rest_config


class ImagenConst:
    project_id = "glasses-imagen"
    location = "asia-east1"
    model = "imagen-3.0-generate-001"
    # model = "imagen-3.0-fast-generate-001"


def imagen_generate_image(prompt,download_count=1):
    vertexai.init(project=ImagenConst.project_id, location=ImagenConst.location)

    model = ImageGenerationModel.from_pretrained(ImagenConst.model)

    _file_name = prompt_to_filenames(prompt)
    _folder = os.path.join(TEMP_FOLDER)
    _remote_folder = os.path.join(rest_config.remote_folder,"imagen")
    _datetime = D.date_file_name()
    
    # if not os.path.isdir(_folder):
    #     raise FileExistsError(f"FileExistsError: {_folder}")
    
    images = model.generate_images(
        prompt=prompt,
        # Optional parameters
        number_of_images=download_count,
        language="ko",
        # You can't use a seed value and watermark at the same time.
        # add_watermark=False,
        # seed=100,
        aspect_ratio="1:1",
        safety_filter_level="block_some",
        person_generation="dont_allow",
    )
    
    if len(images.images) <=1 :
        images[0].save(location=os.path.join(_folder,f"imagen_{_file_name}_{_datetime}.png"), include_generation_parameters=False)
        sftp_client.remote_copy_data(os.path.join(_folder,f"imagen_{_file_name}_{_datetime}.png"), os.path.join(_remote_folder,f"imagen_{_file_name}_{_datetime}.png"))
        return 1
    
    else:
        for i in range(len(images.images)):
            images[i].save(location=os.path.join(_folder,f"imagen_{_file_name}_{i+1}_{_datetime}.png"), include_generation_parameters=False)
            sftp_client.remote_copy_data(os.path.join(_folder,f"imagen_{_file_name}_{i+1}_{_datetime}.png"), os.path.join(_remote_folder,f"imagen_{_file_name}_{i+1}_{_datetime}.png"))
        return len(images.images)
    
def imagen_generate_image_data(prompt,download_count=1):
    vertexai.init(project=ImagenConst.project_id, location=ImagenConst.location)

    model = ImageGenerationModel.from_pretrained(ImagenConst.model)
    
    images = model.generate_images(
        prompt=prompt,
        # Optional parameters
        number_of_images=download_count,
        language="ko",
        # You can't use a seed value and watermark at the same time.
        # add_watermark=False,
        # seed=100,
        aspect_ratio="1:1",
        safety_filter_level="block_some",
        person_generation="dont_allow",
    )
    
    return images.images[0]._pil_image

def imagen_generate_image_path(image_prompt):
    
    MODEL = "imagen"
    QUERY = "query"
    create_time = D.date_file_name()
    
    folder_name = os.path.join(TEMP_FOLDER,f"{MODEL}_{QUERY}_{image_prompt}_{create_time}")
    if not os.path.exists(folder_name):
        os.makedirs(folder_name)
    
    generate_img = imagen_generate_image_data(image_prompt)
    
    generate_img.save(os.path.join(folder_name,f"query.png"))
       
    return os.path.join(folder_name,f"query.png")

def imagen_generate_temp_image_path(image_prompt):
    
    create_time = D.date_file_name()
    
    if not os.path.exists(TEMP_FOLDER):
        os.makedirs(TEMP_FOLDER)
    
    generate_img = imagen_generate_image_data(image_prompt)
    
    generate_img.save(os.path.join(TEMP_FOLDER,f"query_{create_time}.png"))
       
    return os.path.join(TEMP_FOLDER,f"query_{create_time}.png")

if __name__ == '__main__':
    pass
    # imagen_generate_image_data("cat")