from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from gradio_client import Client, handle_file
import uvicorn

app = FastAPI(title="Omni Editor API")

# আপনার স্ক্রিনশট অনুযায়ী সোর্স স্পেস
SOURCE_SPACE = "selfit-camera/omni-image-editor"

try:
    # ক্লায়েন্ট ইনিশিয়ালাইজ করা
    client = Client(SOURCE_SPACE)
except Exception as e:
    print(f"Connection Failed: {e}")

class ImageEditRequest(BaseModel):
    imageUrl: str # এটি আপনার JSON বডির সাথে মিল থাকতে হবে
    prompt: str = "enhance"

@app.get("/")
def health_check():
    return {"status": "Live", "author": "Xalman"}

@app.post("/predict")
async def predict(request: ImageEditRequest):
    try:
        # Gradio Client-এ সরাসরি URL পাঠালে অনেক সময় ৫০০০ এরর দেয়
        # তাই handle_file ব্যবহার করা বাধ্যতামূলক।
        result = client.predict(
            image=handle_file(request.imageUrl), 
            edit_command=request.prompt,        
            api_name="/predict"
        )
        
        # সফল হলে রেজাল্ট ইমেজের পাথ রিটার্ন করবে
        return {
            "status": "success",
            "result_url": result 
        }
    except Exception as e:
        print(f"Detailed Server Error: {e}")
        raise HTTPException(status_code=500, detail=str(e))

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