import torch
import spaces
from threading import Thread
from transformers import AutoProcessor, AutoModelForCausalLM, TextIteratorStreamer
from PIL import Image
from config import settings

MODEL_ID = settings.model
hf_token = settings.hf_token

print(f"Loading {MODEL_ID}...")
processor = AutoProcessor.from_pretrained(MODEL_ID, token=hf_token)
model = AutoModelForCausalLM.from_pretrained(
    MODEL_ID, torch_dtype=torch.float16, token=hf_token
)


@spaces.GPU
def generate_vision_pass(image_path: str, prompt: str) -> str:
    """Run the multimodal vision pass locally on ZeroGPU."""
    model.to("cuda")
    image = Image.open(image_path).convert("RGB")

    messages = [
        {
            "role": "user",
            "content": [{"type": "image"}, {"type": "text", "text": prompt}],
        }
    ]

    text = processor.apply_chat_template(messages, add_generation_prompt=True)
    inputs = processor(text=text, images=image, return_tensors="pt").to("cuda")

    with torch.no_grad():
        output = model.generate(**inputs, max_new_tokens=512)

    generated_ids = output[0, inputs["input_ids"].shape[1] :]
    return processor.decode(generated_ids, skip_special_tokens=True)


@spaces.GPU
def stream_text_generation(messages: list):
    """Stream text generation locally on ZeroGPU."""
    model.to("cuda")

    # Text-only input for the LLM during debate
    # Format to string first, then tokenize
    text = processor.apply_chat_template(
        messages, add_generation_prompt=True, tokenize=False
    )

    tokenizer = processor.tokenizer if hasattr(processor, "tokenizer") else processor
    inputs = tokenizer(text, return_tensors="pt").to("cuda")

    streamer = TextIteratorStreamer(
        tokenizer, skip_prompt=True, skip_special_tokens=True
    )

    generation_kwargs = dict(**inputs, streamer=streamer, max_new_tokens=512)

    thread = Thread(target=model.generate, kwargs=generation_kwargs)
    thread.start()

    for text_chunk in streamer:
        yield text_chunk