from fastapi import FastAPI, UploadFile, File, Form, HTTPException, APIRouter
from fastapi.responses import Response
from starlette.responses import JSONResponse
from trustmark import TrustMark
from PIL import Image
import io

tm = TrustMark(model_type="Q", verbose=False)
appTmp = FastAPI(title="TrustMark Encode API", version="1.0")
app = APIRouter(prefix="/watermark")

@app.post("/encode")
async def encode_watermark(
    image: UploadFile = File(...),
    payload: str = Form(...),
    outputType: str = Form(...),
    strength: float = Form(1.2)
):
    """
    给图片嵌入TrustMark盲水印
    :param image: 原图文件上传
    :param payload: 最多8个ASCII字符
    :param outputType: 输出类型，必须是png，jpg，jpeg
    :param strength: 取值一般0.0~1.5
    :return: 带水印图片二进制
    """
    img_bytes = await image.read()
    img = Image.open(io.BytesIO(img_bytes)).convert("RGB")
    watermarked_img = tm.encode(img, payload, WM_STRENGTH=strength)
    buf = io.BytesIO()
    watermarked_img.save(buf, format=outputType, quality=90)
    buf.seek(0)
    return Response(
        content=buf.read(),
        media_type=f"image/{outputType}"
    )

@app.post("/decode", summary="上传图片提取TrustMark盲水印")
async def extract_watermark(image: UploadFile = File(...)):
    MAX_FILE_SIZE = 10 * 1024 * 1024
    if image.size and image.size > MAX_FILE_SIZE:
        raise HTTPException(status_code=413, detail="文件过大，最大支持10MB")
    allowed_types = ["image/png", "image/jpeg", "image/jpg"]
    if image.content_type not in allowed_types:
        raise HTTPException(status_code=400, detail="仅支持 png / jpg 图片")

    try:
        img_bytes = await image.read()
        img = Image.open(io.BytesIO(img_bytes)).convert("RGB")
        payload, has_watermark, schema = tm.decode(img)

        return JSONResponse({
            "hasWatermark": has_watermark,
            "payload": payload if has_watermark else None,
            "schema": schema if has_watermark else None
        })

    except Exception as e:
        raise HTTPException(status_code=500, detail=f"解析失败: {str(e)}")

if __name__ == "__main__":
    import uvicorn
    uvicorn.run("main:app", host="localhost", port=8081, workers=1)