# app.py

import os
from flask import Flask, request, send_file, jsonify
from rembg import remove
from PIL import Image, ImageEnhance, ImageFilter
import io
import cv2
import numpy as np
from diffusers import DiffusionPipeline
import torch

# --- Create the Flask App ---
app = Flask(__name__)

# --- Configuration ---
# Get the API Key from an environment variable for security.
API_KEY = os.environ.get("BG_REMOVER_API_KEY") 

# --- API Endpoints ---

# A simple root endpoint to check if the server is running.
@app.route('/')
def index():
    return "Background Remover API is running!"

# The main endpoint for removing the background.
@app.route('/remove-bg', methods=['POST'])
def remove_background_api():
    # 1. --- API Key Authentication ---
    api_key_header = request.headers.get('x-api-key')
    if not api_key_header or api_key_header != API_KEY:
        return jsonify({"error": "Unauthorized. Invalid or missing API Key."}), 401

    # 2. --- Image Validation ---
    if 'file' not in request.files:
        return jsonify({"error": "No file part in the request"}), 400

    file = request.files['file']

    if file.filename == '':
        return jsonify({"error": "No selected file"}), 400

    # 3. --- Image Processing ---
    if file:
        try:
            input_image_bytes = file.read()
            output_image_bytes = remove(input_image_bytes)
            
            # 4. --- Send the Response ---
            return send_file(
                io.BytesIO(output_image_bytes),
                mimetype='image/png',
                as_attachment=True,
                download_name='background_removed.png'
            )
        except Exception as e:
            return jsonify({"error": "Failed to process image", "details": str(e)}), 500

    return jsonify({"error": "An unknown error occurred"}), 500

# --- IMAGE COMPRESSOR API ENDPOINT ---

@app.route('/compress-image', methods=['POST'])
def compress_image_api():
    # 1. --- API Key Authentication (Bilkul pehle jaisa) ---
    api_key_header = request.headers.get('x-api-key')
    if not api_key_header or api_key_header != API_KEY:
        return jsonify({"error": "Unauthorized. Invalid or missing API Key."}), 401

    # 2. --- File aur Quality Level ko Check Karna ---
    if 'file' not in request.files:
        return jsonify({"error": "No file part in the request"}), 400

    file = request.files['file']
    
    # User se quality level haasil karna (default 85)
    # request.form se text data liya jaata hai
    quality = int(request.form.get('quality', 85))

    if file.filename == '':
        return jsonify({"error": "No selected file"}), 400

    # 3. --- Image Processing (Pillow Library ka istemal) ---
    if file:
        try:
            # File ko memory mein bytes ki tarah parhna
            input_image_bytes = file.read()
            
            # Pillow ka istemal karke image ko kholna
            img = Image.open(io.BytesIO(input_image_bytes))
            
            # Ek khali "in-memory" file banana jahan hum compressed image save karenge
            output_buffer = io.BytesIO()

            # Image ko 'RGB' mode mein convert karna (JPG ke liye zaroori)
            if img.mode in ("RGBA", "P"):
                img = img.convert("RGB")

            # Image ko buffer mein save karna, lekin is baar quality ke saath
            # quality parameter 1 se 95 tak hota hai (95 best quality)
            img.save(output_buffer, format='JPEG', quality=quality)
            
            # Compressed image ke bytes haasil karna
            output_image_bytes = output_buffer.getvalue()

            # 4. --- Send the Response ---
            # Compressed image ko user ko wapas bhejna
            return send_file(
                io.BytesIO(output_image_bytes),
                mimetype='image/jpeg',
                as_attachment=True,
                download_name=f'compressed_quality_{quality}.jpg'
            )
        except Exception as e:
            return jsonify({"error": "Failed to process image", "details": str(e)}), 500

    return jsonify({"error": "An unknown error occurred"}), 500

# --- AI IMAGE ENHANCER API ENDPOINT ---

@app.route('/enhance-image', methods=['POST'])
def enhance_image_api():
    # 1. --- API Key Authentication (Bilkul pehle jaisa) ---
    api_key_header = request.headers.get('x-api-key')
    if not api_key_header or api_key_header != API_KEY:
        return jsonify({"error": "Unauthorized. Invalid or missing API Key."}), 401

    # 2. --- File Check Karna ---
    if 'file' not in request.files:
        return jsonify({"error": "No file part in the request"}), 400

    file = request.files['file']

    if file.filename == '':
        return jsonify({"error": "No selected file"}), 400

    # 3. --- Image Processing (Pillow Library ka Jadoo) ---
    if file:
        try:
            # File ko memory mein bytes ki tarah parhna
            input_image_bytes = file.read()
            
            # Pillow ka istemal karke image ko kholna
            img = Image.open(io.BytesIO(input_image_bytes))
            
            # Enhancement Step 1: Contrast ko behtar karna
            # 1.0 original contrast hai. 1.2 ka matlab hai 20% zyada contrast.
            enhancer_contrast = ImageEnhance.Contrast(img)
            img = enhancer_contrast.enhance(1.2)
            
            # Enhancement Step 2: Sharpness ko behtar karna
            # 1.0 original sharpness hai. 1.4 ka matlab hai 40% zyada sharp.
            enhancer_sharpness = ImageEnhance.Sharpness(img)
            img = enhancer_sharpness.enhance(1.4)
            
            # Ek khali "in-memory" file banana
            output_buffer = io.BytesIO()
            
            # Enhanced image ko buffer mein save karna
            # PNG format mein save karein taake transparency (agar ho) barqarar rahe
            img.save(output_buffer, format='PNG')
            
            # Enhanced image ke bytes haasil karna
            output_image_bytes = output_buffer.getvalue()

            # 4. --- Send the Response ---
            # Enhanced image ko user ko wapas bhejna
            return send_file(
                io.BytesIO(output_image_bytes),
                mimetype='image/png',
                as_attachment=True,
                download_name='enhanced_image.png'
            )
        except Exception as e:
            return jsonify({"error": "Failed to enhance image", "details": str(e)}), 500

    return jsonify({"error": "An unknown error occurred"}), 500

# --- POWERFUL AI IMAGE ENHANCER API ENDPOINT ---

@app.route('/auto-enhance-image', methods=['POST'])
def auto_enhance_image_api():
    # 1. --- API Key Authentication ---
    api_key_header = request.headers.get('x-api-key')
    if not api_key_header or api_key_header != API_KEY:
        return jsonify({"error": "Unauthorized. Invalid or missing API Key."}), 401

    # 2. --- File Check ---
    if 'file' not in request.files:
        return jsonify({"error": "No file part in the request"}), 400

    file = request.files['file']
    if file.filename == '':
        return jsonify({"error": "No selected file"}), 400

    # 3. --- Professional Image Enhancement with OpenCV ---
    if file:
        try:
            # Image ko bytes se parh kar ek NumPy array mein convert karna
            filestr = file.read()
            npimg = np.frombuffer(filestr, np.uint8)
            img = cv2.imdecode(npimg, cv2.IMREAD_COLOR)

            # --- Enhancement Step 1: CLAHE for Perfect Contrast ---
            # Image ko LAB color space mein convert karna
            lab = cv2.cvtColor(img, cv2.COLOR_BGR2LAB)
            l, a, b = cv2.split(lab)
            # Lightness (L) channel per CLAHE apply karna
            clahe = cv2.createCLAHE(clipLimit=3.0, tileGridSize=(8, 8))
            cl = clahe.apply(l)
            # Channels ko wapas milana
            limg = cv2.merge((cl, a, b))
            # Image ko wapas BGR color space mein convert karna
            img_clahe = cv2.cvtColor(limg, cv2.COLOR_LAB2BGR)

            # --- Enhancement Step 2: Unsharp Masking for Natural Sharpness ---
            # Image ko thoda blur karna
            gaussian_blur = cv2.GaussianBlur(img_clahe, (0, 0), 3.0)
            # Original image mein sharpness add karna
            # Is se details bohot behtar ho jaati hain
            img_enhanced = cv2.addWeighted(img_clahe, 1.5, gaussian_blur, -0.5, 0)

            # --- Final Result ---
            # Processed image (NumPy array) ko wapas PNG bytes mein convert karna
            is_success, buffer = cv2.imencode(".png", img_enhanced)
            if not is_success:
                raise Exception("Could not encode enhanced image.")
            
            output_image_bytes = buffer.tobytes()

            # 4. --- Send the Response ---
            return send_file(
                io.BytesIO(output_image_bytes),
                mimetype='image/png',
                as_attachment=True,
                download_name='auto_enhanced_image.png'
            )
        except Exception as e:
            return jsonify({"error": "Failed to enhance image", "details": str(e)}), 500

    return jsonify({"error": "An unknown error occurred"}), 500

# --- AI IMAGE GENERATOR (TEXT-TO-IMAGE) API ENDPOINT (FINAL, GUARANTEED LAZY LOADING) ---

# Hum model ke liye ek global "box" (variable) banayenge, lekin usko start mein khali rakhenge.
pipe = None

def get_pipeline():
    """
    Yeh ek helper function hai jo model ko sirf ek baar load karega.
    """
    global pipe
    if pipe is None:
        print("Model not loaded. Loading Tiny Stable Diffusion model for the first time...")
        # Note: Secrets (like HUGGING_FACE_HUB_CACHE) will be available here.
        pipe = DiffusionPipeline.from_pretrained(
            "hf-internal-testing/tiny-stable-diffusion-torch"
        )
        print("Model loaded successfully into memory.")
    return pipe

@app.route('/generate-image', methods=['POST'])
def generate_image_api():
    # API Key Authentication
    api_key_header = request.headers.get('x-api-key')
    if not api_key_header or api_key_header != API_KEY:
        return jsonify({"error": "Unauthorized. Invalid or missing API Key."}), 401

    # User ka Text Prompt Haasil Karna
    if not request.is_json:
        return jsonify({"error": "Invalid request: JSON expected"}), 400
    data = request.get_json()
    prompt = data.get('prompt', '')
    if not prompt:
        return jsonify({"error": "Prompt is required"}), 400

    # AI Model se Image Generate Karna
    try:
        # Har request per, hum helper function ko call karke model haasil karenge.
        # Agar model pehle se loaded hai, to woh foran mil jayega.
        # Agar nahi, to woh ab load hoga.
        pipeline = get_pipeline()
        
        print(f"Generating image for prompt: {prompt}")
        image = pipeline(prompt, num_inference_steps=10).images[0]
        print("Image generated.")
        
        output_buffer = io.BytesIO()
        image.save(output_buffer, format='PNG')
        output_image_bytes = output_buffer.getvalue()

        return send_file(
            io.BytesIO(output_image_bytes),
            mimetype='image/png',
            as_attachment=True,
            download_name='generated_image.png'
        )
    except Exception as e:
        return jsonify({"error": "Failed to generate image", "details": str(e)}), 500

# --- Run the App ---
if __name__ == '__main__':
    # For local testing, ensure the environment variable is set.
    if not API_KEY:
        # If you're using the temporary hard-coded key method, you can ignore this error.
        # Otherwise, this reminds you to set the variable.
        print("WARNING: BG_REMOVER_API_KEY is not set. The API will not be secured.")
        print("For local testing, run 'set BG_REMOVER_API_KEY=your-key' (Windows) or 'export BG_REMOVER_API_KEY=your-key' (Mac/Linux) first.")
    
    app.run(debug=True, port=5000)