import requests
from fastapi import FastAPI, UploadFile, File
import shutil
import os

# Initialize FastAPI
app = FastAPI()

# Create an upload directory
UPLOAD_DIR = "/app/uploads"

if not os.path.exists(UPLOAD_DIR):
    os.makedirs(UPLOAD_DIR, exist_ok=True)
    os.chmod(UPLOAD_DIR, 0o777)

# Hugging Face API Key (Replace with your actual API key)
HUGGINGFACE_API_KEY = "HF_API_KEY"


# Hugging Face Model URL
MODEL_URL = "https://api-inference.huggingface.co/models/Salesforce/blip-image-captioning-base"

def get_caption(image_path):
    """Sends image to Hugging Face API for captioning"""
    headers = {
        "Authorization": f"Bearer {HUGGINGFACE_API_KEY}",
        "Content-Type": "application/octet-stream"
    }
    
    with open(image_path, "rb") as image_file:
        image_data = image_file.read()

    response = requests.post(MODEL_URL, headers=headers, data=image_data)
    
    if response.status_code == 200:
        try:
            return response.json()[0]['generated_text']
        except (IndexError, KeyError):
            return f"Error: Unexpected response format: {response.text}"
    else:
        return f"Error: {response.text}"

@app.post("/upload-image/")
async def upload_image(file: UploadFile = File(...)):
    """Uploads image & gets caption"""
    file_path = f"{UPLOAD_DIR}/{file.filename}"
    
    # Save the uploaded file
    with open(file_path, "wb") as buffer:
        shutil.copyfileobj(file.file, buffer)

    # Get caption
    caption = get_caption(file_path)

    return {"filename": file.filename, "caption": caption}

@app.get("/")
async def root():
    return {"message": "Image-to-Text API using Hugging Face Inference API is running."}