from rembg import remove, new_session
from PIL import Image, ImageFilter, ImageOps
import io

# Load model only once (important for performance)
session = new_session("u2net_human_seg")


def replace_background_service(image_bytes: bytes, background_bytes: bytes) -> bytes:
    """
    High-quality background replacement service
    """

    # -------------------------
    # 1. Load input image
    # -------------------------
    input_image = Image.open(io.BytesIO(image_bytes)).convert("RGBA")

    # Ensure clean alpha edges before rembg
    input_image = ImageOps.exif_transpose(input_image)

    # -------------------------
    # 2. Remove background (AI)
    # -------------------------
    foreground = remove(
        input_image,
        session=session,
        alpha_matting=True,
        alpha_matting_foreground_threshold=240,
        alpha_matting_background_threshold=10,
        alpha_matting_erode_size=10,
    )

    # Optional: smooth edges (important for realism)
    alpha = foreground.split()[-1]
    alpha = alpha.filter(ImageFilter.GaussianBlur(radius=0.6))
    foreground.putalpha(alpha)

    # -------------------------
    # 3. Load background
    # -------------------------
    background = Image.open(io.BytesIO(background_bytes)).convert("RGBA")

    # Maintain aspect ratio instead of forced resize distortion
    background = ImageOps.fit(
        background,
        foreground.size,
        method=Image.Resampling.LANCZOS,
        centering=(0.5, 0.5)
    )

    # Optional: slight blur for depth effect (like DSLR bokeh)
    # background = background.filter(ImageFilter.GaussianBlur(radius=1.5))

    # -------------------------
    # 4. Composite (better blending)
    # -------------------------
    result = Image.alpha_composite(background, foreground)

    # -------------------------
    # 5. Export optimized output
    # -------------------------
    output = io.BytesIO()
    result.save(
        output,
        format="PNG",
        optimize=True,
        quality=95
    )

    return output.getvalue()