#!/usr/bin/env python3
import os
import requests
import base64
from pathlib import Path
from backend.config import HF_API_KEY, HF_IMAGE_MODEL_ID
from backend.database.connection import get_connection
from backend.ingestion.image_loader import get_pending_image_jobs, mark_job_completed, mark_job_failed

def get_hf_api_url():
    model_id = HF_IMAGE_MODEL_ID or "Salesforce/blip-image-captioning-large"
    return f"https://api-inference.huggingface.co/models/{model_id}"

def caption_single_image_hf(image_path: str) -> str:
    """
    Generates a caption for a single image using Hugging Face Inference API.
    Zero local VRAM usage.
    """
    if not HF_API_KEY:
        print("[Captioner] Error: HF_API_KEY not set in backend/.env")
        return "Caption generation failed: No API Key."

    headers = {"Authorization": f"Bearer {HF_API_KEY}"}
    api_url = get_hf_api_url()

    try:
        with open(image_path, "rb") as f:
            data = f.read()

        print(f"[Captioner] Calling HF API ({api_url}) for {Path(image_path).name}...")
        response = requests.post(api_url, headers=headers, data=data)
        
        if response.status_code == 200:
            result = response.json()
            if isinstance(result, list) and len(result) > 0 and "generated_text" in result[0]:
                caption = result[0]["generated_text"].strip()
                return caption
            elif "error" in result:
                print(f"[Captioner] HF API Error: {result['error']}")
                return "Caption generation failed."
        else:
            print(f"[Captioner] HF API failed with status {response.status_code}: {response.text}")
            return "Caption generation failed."
            
    except Exception as e:
        print(f"[Captioner] Error processing {image_path}: {e}")
        return "Caption generation failed."
    
    return "Caption generation failed."

def run_caption_pipeline():
    """
    Main pipeline to process pending image jobs via Hugging Face API.
    """
    print("[Captioner] Starting HF Image Captioning pipeline")
    
    jobs = get_pending_image_jobs()
    print(f"[Captioner] Found {len(jobs)} pending jobs")
    
    if not jobs:
        print("No jobs. Exiting.")
        return

    for job in jobs:
        print(f"[Captioner] Processing: {job['image_path']}")
        caption = caption_single_image_hf(job['image_path'])
        
        if caption and not caption.startswith("Caption generation failed"):
            mark_job_completed(job['id'], caption)
            print(f"[Captioner] Job {job['id']} completed: {caption}")
        else:
            mark_job_failed(job['id'], caption)
            print(f"[Captioner] Job {job['id']} failed.")
            
    print("[Captioner] Pipeline finished.")

if __name__ == "__main__":
    run_caption_pipeline()