from google import genai
import os


def enhance_prompt(image, prompt):
    input_caption_prompt = (
        "Please provide a prompt for a Diffusion Model text-to-image generative model for the image I will give you. "
        "The prompt should be a detailed description of the image, especially the main subject (i.e. the main character/asset/item), the environment, the pose, the lighting, the camera view, the style etc."
        "The prompt should be detailed enough to generate the target image. "
        "The prompt should be short and precise, in one-line format, and does not exceed 77 tokens."
        "The prompt should be individually coherent as a description of the image."
    )

    caption_model = genai.Client(
        vertexai=False, api_key=os.environ["GOOGLE_API_KEY"]
    )
    input_image_prompt = caption_model.models.generate_content(
        model='gemini-1.5-flash', contents=[input_caption_prompt, image]).text
    input_image_prompt = input_image_prompt.replace('\r', '').replace('\n', '')

    enhance_instruction = "Enhance this input text prompt: '"
    enhance_instruction += prompt
    enhance_instruction += "'. Please extract other details, especially description of the main subject from the following reference prompt: '"
    enhance_instruction += input_image_prompt
    enhance_instruction += "'. Please keep the details that are mentioned in the input prompt, and enhance the rest. "
    enhance_instruction += "Response with only the enhanced prompt. "
    enhance_instruction += "The enhanced prompt should be short and precise, in one-line format, and does not exceed 77 tokens."
    enhanced_prompt = caption_model.models.generate_content(
        model='gemini-1.5-flash', contents=[enhance_instruction]).text.replace('\r', '').replace('\n', '')
    print("input_image_prompt: ", input_image_prompt)
    print("prompt: ", prompt)
    print("enhanced_prompt: ", enhanced_prompt)
    return enhanced_prompt

def enhance_garment_prompt(image, prompt):
    """
    Enhances a prompt specifically for garment transformation tasks.
    
    Args:
        image: The garment image
        prompt: User provided prompt for the desired output
        
    Returns:
        Enhanced prompt that preserves garment details while incorporating user requirements
    """
    input_caption_prompt = (
        "Please provide a detailed description of this garment/clothing item I will show you. "
        "Focus on describing the garment's color, pattern, style, fabric, cut, and unique details. "
        "Be specific about the type of garment (e.g., t-shirt, dress, jacket, pants) and its defining characteristics. "
        "The description should be detailed enough for an image generation model to recreate this exact garment. "
        "The description should be short and precise, in one-line format."
    )
    
    caption_model = genai.Client(
        vertexai=False, api_key=os.environ["GOOGLE_API_KEY"]
    )
    
    # Get detailed garment description
    garment_description = caption_model.models.generate_content(
        model='gemini-1.5-flash', contents=[input_caption_prompt, image]).text
    garment_description = garment_description.replace('\r', '').replace('\n', '')
    
    # Enhance user prompt to include garment details
    enhance_instruction = (
        f"I need to generate an image of a model wearing a specific garment. "
        f"The garment is described as: '{garment_description}'. "
        f"The user wants: '{prompt}'. "
        f"Create a detailed prompt that combines these elements, ensuring the garment description is preserved exactly while "
        f"incorporating the user's requirements for the model (person wearing it) and setting/background. "
        f"Focus on describing a photorealistic scene with a model wearing this specific garment. "
        f"The enhanced prompt should be short and precise, in one-line format, and should not exceed 77 tokens."
    )
    
    enhanced_prompt = caption_model.models.generate_content(
        model='gemini-1.5-flash', contents=[enhance_instruction]).text.replace('\r', '').replace('\n', '')
    
    print("garment_description: ", garment_description)
    print("user prompt: ", prompt)
    print("enhanced_prompt: ", enhanced_prompt)
    
    return enhanced_prompt