import io
from flask import Flask, request, send_file, jsonify
import numpy as np
from PIL import Image
import cv2

app = Flask(__name__)

try:
    from imwatermark import WatermarkEncoder
    encoder = WatermarkEncoder()
    WM_AVAILABLE = True
    WM_ERROR = None
except Exception as e:
    encoder = None
    WM_AVAILABLE = False
    WM_ERROR = str(e)


@app.post("/watermark")
def watermark():
    if not WM_AVAILABLE or encoder is None:
        return jsonify({"success": False, "error": f"Watermark encoder unavailable: {WM_ERROR}"}), 500
    image_file = request.files.get("image")
    payload = request.form.get("payload", "")

    if not image_file:
        return jsonify({"success": False, "error": "No image provided"}), 400
    if not payload:
        return jsonify({"success": False, "error": "No payload provided"}), 400

    try:
        img = Image.open(image_file.stream).convert("RGB")
        img_np = np.array(img)[:, :, ::-1]  # to BGR
    except Exception as e:
        return jsonify({"success": False, "error": f"Invalid image: {e}"}), 400

    try:
        wm_bytes = payload.encode("utf-8")
        encoder.set_watermark('bytes', wm_bytes)
        watermarked = encoder.encode(img_np, 'dwtDct')
    except Exception as e:
        return jsonify({"success": False, "error": f"Watermark failed: {e}"}), 500

    watermarked_rgb = cv2.cvtColor(watermarked, cv2.COLOR_BGR2RGB)
    out_img = Image.fromarray(watermarked_rgb)

    buf = io.BytesIO()
    out_img.save(buf, format="WEBP", quality=90)
    buf.seek(0)
    return send_file(buf, mimetype="image/webp")


if __name__ == "__main__":
    app.run(host="0.0.0.0", port=8000)