from flask import Flask, request, jsonify, send_file
import torch
import cv2
import numpy as np
import tempfile
import os
from PIL import Image
from io import BytesIO

app = Flask(__name__)

# Load model
print("🔄 Loading AnimeGANv2 model...")
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = torch.hub.load("bryandlee/animegan2-pytorch:main", "generator", pretrained="face_paint_512_v2")
model = model.to(device).eval()
print("✅ Model loaded on:", device)

def image_to_anime(input_image, size=512):
    """Convert single image to anime style"""
    img = input_image.resize((size, size))
    img_np = np.array(img).astype(np.float32) / 127.5 - 1
    img_tensor = torch.from_numpy(img_np).permute(2, 0, 1).unsqueeze(0).to(device)
    
    with torch.no_grad():
        out = model(img_tensor)
    
    out_np = out.squeeze(0).permute(1, 2, 0).cpu().numpy()
    out_np = (out_np + 1) * 127.5
    out_np = np.clip(out_np, 0, 255).astype(np.uint8)
    result = Image.fromarray(out_np).resize(input_image.size, Image.LANCZOS)
    return result

@app.route('/health', methods=['GET'])
def health():
    return jsonify({
        "status": "healthy", 
        "device": str(device),
        "supported_formats": ["png", "jpg", "jpeg", "webp", "bmp"]
    })

@app.route('/convert/image', methods=['POST'])
def convert_image():
    """Convert image to anime style - Supports PNG, JPG, JPEG, WEBP, BMP"""
    
    # Check if image file is present
    if 'image' not in request.files:
        return jsonify({"error": "No image file provided"}), 400
    
    file = request.files['image']
    
    # Check if filename is empty
    if file.filename == '':
        return jsonify({"error": "No file selected"}), 400
    
    # Check file extension (optional validation)
    allowed_extensions = {'png', 'jpg', 'jpeg', 'webp', 'bmp'}
    file_ext = file.filename.rsplit('.', 1)[1].lower() if '.' in file.filename else ''
    
    if file_ext not in allowed_extensions:
        return jsonify({
            "error": f"Unsupported format. Use: {', '.join(allowed_extensions)}"
        }), 400
    
    try:
        # Open and convert image (PIL handles multiple formats automatically)
        img = Image.open(file).convert('RGB')
        
        # Get original format for response
        original_format = file_ext if file_ext else 'png'
        
        # Convert to anime style
        result = image_to_anime(img)
        
        # Save to bytes buffer as PNG (quality preserved)
        buf = BytesIO()
        result.save(buf, format='PNG')
        buf.seek(0)
        
        # Return as PNG (consistent output format)
        return send_file(
            buf, 
            mimetype='image/png', 
            as_attachment=True, 
            download_name=f'anime_output_{original_format}.png'
        )
        
    except Exception as e:
        return jsonify({"error": f"Conversion failed: {str(e)}"}), 500

@app.route('/convert/image/base64', methods=['POST'])
def convert_image_base64():
    """Convert image to anime style and return as base64 (good for n8n)"""
    
    if 'image' not in request.files:
        return jsonify({"error": "No image file provided"}), 400
    
    file = request.files['image']
    
    try:
        import base64
        img = Image.open(file).convert('RGB')
        result = image_to_anime(img)
        
        buf = BytesIO()
        result.save(buf, format='PNG')
        buf.seek(0)
        
        encoded = base64.b64encode(buf.getvalue()).decode('utf-8')
        
        return jsonify({
            "status": "success",
            "image_base64": encoded,
            "format": "png"
        })
        
    except Exception as e:
        return jsonify({"error": str(e)}), 500

if __name__ == "__main__":
    port = int(os.environ.get("PORT", 7860))
    app.run(host='0.0.0.0', port=port)