from flask import Flask, render_template, request
import io
import base64
from PIL import Image
import numpy as np
import random
# Import procedural generators
from generate_dataset import generate_wood, generate_marble, generate_fabric, generate_brick, generate_noise

app = Flask(__name__)

# Configuration
TEXTURE_TYPES = ["wood", "marble", "fabric", "brick", "noise"]

# Map types to generator functions
GENERATORS = {
    "wood": generate_wood,
    "marble": generate_marble,
    "fabric": generate_fabric,
    "brick": generate_brick,
    "noise": generate_noise,
}

@app.route('/')
def index():
    return render_template('index.html', texture_types=TEXTURE_TYPES)

@app.route('/generate', methods=['POST'])
def generate():
    try:
        texture_type = request.form.get('texture_type')
        seed = int(request.form.get('seed', 0))
        # Creativity affects randomness parameters
        creativity = float(request.form.get('creativity', 1.0))
        
        if texture_type not in GENERATORS:
            return {"error": f"Unknown texture type: {texture_type}"}, 400

        # Set seed for reproducibility
        random.seed(seed)
        np.random.seed(seed)
        
        # Call the procedural generator
        # Note: 'creativity' parameter is not supported by current generate_dataset functions directly,
        # but we can simulate it by just using the seed variation or updating generate_dataset.
        # For now, we rely on the seed.
        
        img_pil = GENERATORS[texture_type](size=256) # Generate larger image for UI
        
        # Convert to base64
        buf = io.BytesIO()
        img_pil.save(buf, format='PNG')
        buf.seek(0)
        img_b64 = base64.b64encode(buf.getvalue()).decode('utf-8')
        
        return {"image": f"data:image/png;base64,{img_b64}"}
        
    except Exception as e:
        print(f"Error generating image: {e}")
        return {"error": str(e)}, 500

if __name__ == '__main__':
    app.run(debug=True, port=5000)