from fastapi import FastAPI, File, UploadFile, Form, HTTPException
from fastapi.responses import JSONResponse
from gradio_client import Client, handle_file
import os
import tempfile
import shutil
from fastapi.middleware.cors import CORSMiddleware
import logging

# Initialize logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

app = FastAPI()

client = Client("Ashrafb/moondream_captioning")

# Create the "uploads" directory if it doesn't exist
os.makedirs("uploads", exist_ok=True)

# Define a function to save uploaded file to a temporary file
async def save_upload_file(upload_file: UploadFile) -> str:
    try:
        # Create a temporary directory if it doesn't exist
        os.makedirs("temp_uploads", exist_ok=True)
        # Create a temporary file path
        temp_file_path = os.path.join("temp_uploads", tempfile.NamedTemporaryFile(delete=False).name)
        # Save the uploaded file to the temporary file
        with open(temp_file_path, "wb") as buffer:
            shutil.copyfileobj(upload_file.file, buffer)
        return temp_file_path
    except Exception as e:
        logger.error(f"Error saving upload file: {e}")
        raise HTTPException(status_code=500, detail=f"Error saving upload file: {e}")

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # Adjust as needed, '*' allows requests from any origin
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

@app.post("/get_caption")
async def get_caption(image: UploadFile = File(...), context: str = Form(...)):
    try:
        # Save the uploaded image to a temporary file
        temp_file_path = await save_upload_file(image)
        
        # Debugging: Print the value of additional_context
        logger.info(f"Additional Context: {context}")
        
        # Check if additional context is provided and not None
        if context is not None:
            context = context.strip()
        
        # Log the parameters being passed to the Gradio client
        logger.info(f"Calling client.predict with image={temp_file_path} and context={context}")
        
        # Use handle_file to handle the file upload correctly
        result = client.predict(
            image=handle_file(temp_file_path),
            question=context,
            api_name="/get_caption"
        )
        
        return {"caption": result}
    except Exception as e:
        logger.error(f"Error in get_caption: {e}")
        raise HTTPException(status_code=500, detail=f"Error in get_caption: {e}")

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