import os
import base64
import json
from flask import Flask, render_template, request, jsonify
from flask_cors import CORS
import requests
from werkzeug.utils import secure_filename
import tempfile

app = Flask(__name__)
CORS(app)

# Configuration
app.config['MAX_CONTENT_LENGTH'] = 25 * 1024 * 1024  # 25MB max file size
ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg', 'webp'}

# Get API key from environment variable
OPENAI_API_KEY = os.environ.get('OPENAI_API_KEY')
API_BASE_URL = os.environ.get('API_BASE_URL', 'https://api.openai.com/v1')

def allowed_file(filename):
    return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS

@app.route('/')
def index():
    return render_template('index.html')

@app.route('/generate', methods=['POST'])
def generate_image():
    try:
        data = request.json
        prompt = data.get('prompt', '')
        
        if not prompt:
            return jsonify({'error': 'Prompt is required'}), 400
        
        headers = {
            'Content-Type': 'application/json',
            'Authorization': f'Bearer {OPENAI_API_KEY}'
        }
        
        payload = {
            'model': 'dall-e-3',
            'prompt': prompt,
            'n': 1,
            'size': data.get('size', '1024x1024'),
            'quality': data.get('quality', 'standard'),
            'style': data.get('style', 'vivid')
        }
        
        response = requests.post(
            f'{API_BASE_URL}/images/generations',
            headers=headers,
            json=payload
        )
        
        if response.status_code == 200:
            result = response.json()
            # For dall-e-3, the response contains URLs
            image_url = result['data'][0]['url']
            
            # Download the image and convert to base64
            image_response = requests.get(image_url)
            if image_response.status_code == 200:
                image_base64 = base64.b64encode(image_response.content).decode('utf-8')
                return jsonify({
                    'success': True,
                    'image': f'data:image/png;base64,{image_base64}',
                    'usage': result.get('usage', {})
                })
            else:
                return jsonify({'error': 'Failed to download generated image'}), 500
        else:
            return jsonify({'error': response.json()}), response.status_code
            
    except Exception as e:
        return jsonify({'error': str(e)}), 500

@app.route('/edit', methods=['POST'])
def edit_image():
    try:
        # Check if file was uploaded
        if 'image' not in request.files:
            return jsonify({'error': 'No image provided'}), 400
        
        image_file = request.files['image']
        prompt = request.form.get('prompt', '')
        
        if not prompt:
            return jsonify({'error': 'Prompt is required'}), 400
        
        if not image_file or image_file.filename == '':
            return jsonify({'error': 'Image is required'}), 400
        
        if not allowed_file(image_file.filename):
            return jsonify({'error': 'Invalid image format'}), 400
        
        # Prepare multipart form data
        headers = {
            'Authorization': f'Bearer {OPENAI_API_KEY}'
        }
        
        # Read image content
        image_content = image_file.read()
        
        # Build files list for multipart upload
        files = {
            'model': (None, 'dall-e-2'),
            'prompt': (None, prompt),
            'n': (None, '1'),
            'size': (None, request.form.get('size', '1024x1024')),
            'image': (image_file.filename, image_content, image_file.content_type)
        }
        
        # Add mask if provided
        if 'mask' in request.files:
            mask_file = request.files['mask']
            if mask_file and mask_file.filename != '':
                mask_content = mask_file.read()
                files['mask'] = (mask_file.filename, mask_content, mask_file.content_type)
        
        response = requests.post(
            f'{API_BASE_URL}/images/edits',
            headers=headers,
            files=files
        )
        
        if response.status_code == 200:
            result = response.json()
            # For dall-e-2, the response format depends on response_format parameter
            # Default is URL format
            if 'b64_json' in result['data'][0]:
                image_data = f'data:image/png;base64,{result["data"][0]["b64_json"]}'
            else:
                # Download image from URL
                image_url = result['data'][0]['url']
                image_response = requests.get(image_url)
                if image_response.status_code == 200:
                    image_base64 = base64.b64encode(image_response.content).decode('utf-8')
                    image_data = f'data:image/png;base64,{image_base64}'
                else:
                    return jsonify({'error': 'Failed to download generated image'}), 500
            
            return jsonify({
                'success': True,
                'image': image_data,
                'usage': result.get('usage', {})
            })
        else:
            error_msg = 'Request failed'
            try:
                error_data = response.json()
                error_msg = error_data.get('error', {}).get('message', error_msg)
            except:
                error_msg = response.text
            return jsonify({'error': error_msg}), response.status_code
            
    except Exception as e:
        return jsonify({'error': str(e)}), 500

if __name__ == '__main__':
    if not OPENAI_API_KEY:
        print("Warning: OPENAI_API_KEY environment variable not set!")
    app.run(host='0.0.0.0', port=7860, debug=True)