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

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

# সোর্স স্পেস যেখানে আসল মডেলটি আছে
SOURCE_SPACE = "selfit-camera/omni-image-editor"

def get_client():
    """চেষ্টা করবে মূল AI স্পেসের সাথে কানেক্ট করার"""
    max_retries = 3
    for i in range(max_retries):
        try:
            print(f"Connecting to AI Source (Attempt {i+1})...")
            client = Client(SOURCE_SPACE)
            return client
        except Exception as e:
            print(f"Connection failed: {e}")
            if i < max_retries - 1:
                time.sleep(5) # ৫ সেকেন্ড অপেক্ষা করে আবার চেষ্টা করবে
    return None

class EditRequest(BaseModel):
    imageUrl: str
    prompt: str = "enhance quality"

@app.get("/")
def home():
    return {
        "status": "Running",
        "service": "Omni-Editor-API",
        "author": "Xalman"
    }

@app.post("/predict")
async def predict(data: EditRequest):
    client = get_client()
    
    if client is None:
        raise HTTPException(
            status_code=503, 
            detail="AI Source Space is currently unavailable or sleeping. Please try again in a minute."
        )
    
    try:
        # handle_file ব্যবহার করে প্রসেস করা
        result = client.predict(
            image=handle_file(data.imageUrl),
            edit_command=data.prompt,
            api_name="/predict"
        )
        
        return {
            "status": "success",
            "result_url": result
        }
    except Exception as e:
        print(f"Prediction Error: {e}")
        raise HTTPException(status_code=500, detail=str(e))

if __name__ == "__main__":
    port = int(os.environ.get("PORT", 7860))
    uvicorn.run(app, host="0.0.0.0", port=port)