import torch
from transformers import BlipProcessor, BlipForQuestionAnswering, BlipForConditionalGeneration
from config import DEVICE

processor_vqa = BlipProcessor.from_pretrained("Salesforce/blip-vqa-base")
model_vqa = BlipForQuestionAnswering.from_pretrained("Salesforce/blip-vqa-base").to(DEVICE)
model_vqa.eval()

processor_cap = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
model_cap = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base").to(DEVICE)
model_cap.eval()

def blip_answer(image, question, lang="en"):
    # Fix numpy tensor casting issues in HF deployment
    if image.mode != "RGB":
        image = image.convert("RGB")

    # 1. Generate Caption
    cap_inputs = processor_cap(images=image, return_tensors="pt").to(DEVICE)
    with torch.no_grad():
        cap_output = model_cap.generate(**cap_inputs, max_new_tokens=20)
    caption = processor_cap.decode(cap_output[0], skip_special_tokens=True)

    # 2. Generate Base Answer
    vqa_inputs = processor_vqa(images=image, text=question, return_tensors="pt").to(DEVICE)
    with torch.no_grad():
        vqa_output = model_vqa.generate(**vqa_inputs, max_new_tokens=10)
    answer = processor_vqa.decode(vqa_output[0], skip_special_tokens=True)

    return caption, answer, "Direct vision reasoning processed smoothly by local BLIP model."