import json
import base64
import random
from PIL import Image, ImageEnhance
from io import BytesIO
import numpy as np
import time

def load_image_data(json_path='AgPixelGen.json'):
    with open(json_path, 'r') as f:
        return json.load(f)

def decode_image_from_base64(base64_string):
    image_bytes = base64.b64decode(base64_string)
    image = Image.open(BytesIO(image_bytes))
    return image

def random_transform(image):
    random.seed(time.time())
    transforms = ['color_change', 'noise']
    random.shuffle(transforms)

    for transform_type in transforms:
        if transform_type == 'color_change':
            image = color_change(image)
        
        elif transform_type == 'noise':
            image = add_noise(image)
    
    return image

def color_change(image):
    enhancer = ImageEnhance.Color(image)
    image = enhancer.enhance(random.uniform(0.2, 3.0))
    
    enhancer = ImageEnhance.Brightness(image)
    image = enhancer.enhance(random.uniform(0.5, 2.0))

    enhancer = ImageEnhance.Contrast(image)
    image = enhancer.enhance(random.uniform(0.5, 2.0))
    
    return image

def add_noise(image):
    width, height = image.size
    pixels = np.array(image)
    
    noise_factor = random.uniform(0.02, 0.1)
    noise = np.random.normal(scale=noise_factor, size=(height, width, 3))
    
    noisy_pixels = pixels + noise * 255
    noisy_pixels = np.clip(noisy_pixels, 0, 255).astype(np.uint8)
    
    return Image.fromarray(noisy_pixels)

def decode_and_transform_image(image_data, key):
    image = decode_image_from_base64(image_data[key])
    transformed_image = random_transform(image)
    return transformed_image

def save_image(image, output_path='output.jpg'):
    image.save(output_path)

def generate_image(user_prompt=""):
    image_data = load_image_data()
    if user_prompt:
        key = user_prompt.replace(' ', '+')
    else:
        matched_keys = [k for k in image_data.keys() if all(word.lower() in k.lower() for word in user_prompt.split())]
        if matched_keys:
            key = random.choice(matched_keys)
        else:
            key = None

    if not key or key not in image_data:
        matched_keys = [k for k in image_data.keys() if all(word.lower() in k.lower() for word in user_prompt.split())]
        
        if matched_keys:
            key = random.choice(matched_keys)
        else:
            return "<error>"

    transformed_image = decode_and_transform_image(image_data, key)
    save_image(transformed_image)
    return "Image generated and saved."

if __name__ == "__main__":
    user_prompt = input("Enter your prompt: ")
    result = generate_image(user_prompt)
    print(result)