from typing import List
from diffusers.modular_pipelines import (
    PipelineState,
    ModularPipelineBlocks,
    InputParam,
    OutputParam,
)
from transformers import AutoModelForCausalLM, AutoTokenizer
import os

SYSTEM_PROMPT = (
    "You are an expert image generation assistant. "
    "Take the user's short description and expand it into a vivid, detailed, and clear image generation prompt. "
    "Ensure rich colors, depth, realistic lighting, and an imaginative composition. "
    "Avoid vague terms — be specific about style, perspective, and mood. "
    "Try to keep the output under 512 tokens. "
    "Please don't return any prefix or suffix tokens, just the expanded user description."
)

class QwenPromptExpander(ModularPipelineBlocks):
    def __init__(self, model_id="Qwen/Qwen2.5-3B-Instruct", system_prompt=SYSTEM_PROMPT):
        super().__init__()

        self.system_prompt = system_prompt
        self.model = AutoModelForCausalLM.from_pretrained(
            model_id, torch_dtype="auto"
        ).to("cuda")
        self.tokenizer = AutoTokenizer.from_pretrained(model_id)

    @property
    def expected_components(self):
        return []

    @property
    def inputs(self) -> List[InputParam]:
        return [
            InputParam(
                "prompt",
                type_hint=str,
                required=True,
                description="Prompt to use",
            )
        ]
    
    @property
    def intermediate_inputs(self) -> List[InputParam]:
        return []

    @property
    def intermediate_outputs(self) -> List[OutputParam]:
        return [
            OutputParam(
                "prompt",
                type_hint=str,
                description="Expanded prompt by the LLM",
            ),
            OutputParam(
                "old_prompt",
                type_hint=str,
                description="Old prompt provided by the user",
            )
        ]


    def __call__(self, components, state: PipelineState) -> PipelineState:
        block_state = self.get_block_state(state)

        old_prompt = block_state.prompt
        print(f"Actual prompt: {old_prompt}")
        
        messages = [
            {"role": "system", "content": self.system_prompt},
            {"role": "user", "content": old_prompt}
        ]
        text = self.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
        model_inputs = self.tokenizer([text], return_tensors="pt").to(self.model.device)
        generated_ids = self.model.generate(**model_inputs,max_new_tokens=512)
        generated_ids = [
            output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
        ]

        block_state.prompt = self.tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
        block_state.old_prompt = old_prompt
        print(f"{block_state.prompt=}")
        self.set_block_state(state, block_state)

        return components, state