import os
import io
import uuid
import logging
from PIL import Image

try:
    from rembg import remove as _rembg_remove
    REMBG_AVAILABLE = True
except ImportError:
    REMBG_AVAILABLE = False
    _rembg_remove = None

logger = logging.getLogger(__name__)

BG_REMOVED_DIR = os.path.join('static', 'bg-removed')
os.makedirs(BG_REMOVED_DIR, exist_ok=True)

ALLOWED_EXTENSIONS = {'jpg', 'jpeg', 'png', 'webp'}
MAX_FILE_SIZE = 10 * 1024 * 1024
MAX_PIXELS = 25_000_000
MAX_DIMENSION = 8000


def validate_image(file_storage):
    if not file_storage or not file_storage.filename:
        return "No file uploaded"

    ext = file_storage.filename.rsplit('.', 1)[-1].lower() if '.' in file_storage.filename else ''
    if ext not in ALLOWED_EXTENSIONS:
        return "Unsupported format. Please upload JPG, PNG, or WEBP"

    file_storage.seek(0, os.SEEK_END)
    size = file_storage.tell()
    file_storage.seek(0)
    if size > MAX_FILE_SIZE:
        return f"File too large ({size // (1024*1024)}MB). Maximum is 10MB"

    try:
        img = Image.open(file_storage)
        img.verify()
        file_storage.seek(0)
    except Exception:
        return "File is not a valid image"

    img = Image.open(file_storage)
    w, h = img.size
    file_storage.seek(0)

    if w > MAX_DIMENSION or h > MAX_DIMENSION:
        return f"Image dimensions too large ({w}x{h}). Maximum is {MAX_DIMENSION}px per side"
    if w * h > MAX_PIXELS:
        return f"Image has too many pixels ({w*h:,}). Maximum is {MAX_PIXELS:,}"

    return None


def remove_background(file_storage):
    if not REMBG_AVAILABLE:
        raise RuntimeError(
            "Background removal is not available: the 'rembg' package is not installed. "
            "Install it with: pip install rembg[gpu] or pip install rembg"
        )

    ext = file_storage.filename.rsplit('.', 1)[-1].lower()
    uid = uuid.uuid4().hex[:12]

    orig_filename = f'orig_{uid}.{ext}'
    orig_path = os.path.join(BG_REMOVED_DIR, orig_filename)
    file_storage.save(orig_path)

    with open(orig_path, 'rb') as f:
        input_data = f.read()

    output_data = _rembg_remove(input_data)

    result_filename = f'result_{uid}.png'
    result_path = os.path.join(BG_REMOVED_DIR, result_filename)
    with open(result_path, 'wb') as f:
        f.write(output_data)

    return {
        'original_url': f'/static/bg-removed/{orig_filename}',
        'result_url': f'/static/bg-removed/{result_filename}',
    }