from fastapi import FastAPI, HTTPException
from pydantic import BaseModel, Field
import uvicorn
import sys
import os

# Add liblibai_controlnet to path so we can reuse its OSS infrastructure
liblib_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), "../liblibai_controlnet"))
if liblib_dir not in sys.path:
    sys.path.insert(0, liblib_dir)

try:
    from liblibai_client import LibLibAIClient
except ImportError:
    LibLibAIClient = None

app = FastAPI(title="LibLib Image Uploader Proxy")

class UploadRequest(BaseModel):
    image_base64: str = Field(..., description="Base64 encoded image data (e.g. data:image/png;base64,iVBORw0...)")
    filename: str = Field("image.png", description="Optional filename for extension detection")

@app.post("/upload")
async def upload_image(req: UploadRequest):
    if not LibLibAIClient:
        raise HTTPException(status_code=500, detail="LibLibAIClient module not found")
        
    try:
        client = LibLibAIClient()
        # Ensure it has basic prefix logic if omitted, though upload_base64_image handles 'data:image...' natively
        raw_b64 = req.image_base64
        url = client.upload_base64_image(raw_b64)
        return {"url": url, "status": "success"}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))

@app.get("/health")
async def health():
    return {"status": "ok"}

if __name__ == "__main__":
    port = int(sys.argv[sys.argv.index("--port") + 1]) if "--port" in sys.argv else 8004
    uvicorn.run(app, host="0.0.0.0", port=port)