import requests
import json
import uuid
from fastapi import FastAPI, UploadFile, File, Query
from fastapi.responses import JSONResponse, Response
from fastapi.middleware.cors import CORSMiddleware
from rembg import remove
import uvicorn

# --------------------------
# Temp Memory Storage
# --------------------------
TEMP_FILES = {}

# --------------------------
# Pretty JSON Response
# --------------------------
class PrettyJSONResponse(JSONResponse):
    def render(self, content) -> bytes:
        return json.dumps(content, indent=4, ensure_ascii=False).encode("utf-8")

app = FastAPI(
    title="AI ClearCut - HF UltraFast BG Remover",
    default_response_class=PrettyJSONResponse
)

# --------------------------
# CORS
# --------------------------
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# --------------------------
# Root
# --------------------------
@app.get("/")
async def root():
    return {
        "message": "HF UltraFast Background Remover Running 🚀",
        "endpoints": {
            "POST /upload": "Upload image",
            "GET /json?url=": "Image from URL",
            "GET /file/{id}": "Get processed image"
        },
        "creator": "Jerrycoder"
    }

# --------------------------
# Serve temp file
# --------------------------
@app.get("/file/{file_id:path}")
async def get_file(file_id: str):
    file_data = TEMP_FILES.get(file_id)
    if not file_data:
        return {"status": "error", "message": "File expired"}

    return Response(content=file_data, media_type="image/png")

# --------------------------
# Process image
# --------------------------
def process_image(img_bytes: bytes):
    try:
        # Remove background
        output_bytes = remove(img_bytes)

        # Generate unique ID
        file_ext = "png"  # default
        file_id = f"{uuid.uuid4()}.{file_ext}"
        file_url = f"https://jerrycoder-rembg-as.hf.space/file/{file_id}"

        # Store temporarily
        TEMP_FILES[file_id] = output_bytes

        return {
            "status": "success",
            "url": file_url,
            "full_url": file_url,
            "creator": "Jerrycoder"
        }

    except Exception as e:
        return {"status": "error", "message": str(e)}

# --------------------------
# POST upload
# --------------------------
@app.post("/upload")
async def upload_file(file: UploadFile = File(...)):
    img_bytes = await file.read()
    return process_image(img_bytes)

# --------------------------
# GET from URL
# --------------------------
@app.get("/json")
async def from_url(url: str = Query(...)):
    try:
        resp = requests.get(url, timeout=10)
        resp.raise_for_status()
        return process_image(resp.content)
    except Exception as e:
        return {"status": "error", "message": str(e)}

# --------------------------
# Run
# --------------------------
if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=7860)