import torch
from transformers import AutoProcessor, AutoModelForCausalLM
from config import DEVICE
import gc

processor_florence = None
model_florence = None

def load_florence():
    global processor_florence, model_florence
    if model_florence is not None:
        return
    try:
        processor_florence = AutoProcessor.from_pretrained("microsoft/Florence-2-base-ft", trust_remote_code=True)
        model_florence = AutoModelForCausalLM.from_pretrained("microsoft/Florence-2-base-ft", trust_remote_code=True).to(DEVICE)
        model_florence.eval()
    except Exception as e:
        print("Florence-2 failed to load:", e)

def florence_answer(image, question, lang="en"):
    load_florence()
    if model_florence is None:
        return "Server Error", "Florence-2 Not Loaded", "The model ran out of memory or failed to load on the Hugging Face space."

    if image.mode != "RGB":
        image = image.convert("RGB")

    def run_task(task_prompt, text_input=None):
        prompt = task_prompt if text_input is None else task_prompt + text_input
        inputs = processor_florence(text=prompt, images=image, return_tensors="pt").to(DEVICE)
        with torch.no_grad():
            generated_ids = model_florence.generate(
                input_ids=inputs["input_ids"],
                pixel_values=inputs["pixel_values"],
                max_new_tokens=100,
                num_beams=3
            )
        generated_text = processor_florence.batch_decode(generated_ids, skip_special_tokens=False)[0]
        parsed = processor_florence.post_process_generation(generated_text, task=task_prompt, image_size=(image.width, image.height))
        return parsed[task_prompt]

    try:
        # Step 1: Detailed Captioning
        caption = run_task("<MORE_DETAILED_CAPTION>")
        
        # Step 2: VQA Answering
        answer = run_task("<VQA>", question)
        
        # Step 3: Explanation Logic
        # Florence-2 dense captions act as a perfect contextual reasoning base
        explanation = f"Florence-2 performed a massive deep visual scan, identifying: '{caption}'. Because of these precise extracted contextual elements, the model deduced the answer is {answer}."
        
        return caption, answer, explanation
    except Exception as e:
        return "Florence API Error", "Execution Crashed", repr(e)