from dataclasses import dataclass
import io
import os
import re
import time
import random
import base64
import textwrap
import torch
import transformers
import gradio as gr
from PIL import Image
from modules import scripts_manager, shared, devices, errors, processing, sd_models, sd_modules, timer, ui_symbols
from modules import ui_control_helpers
from modules.sd_offload_aux import register_aux, deregister_aux, move_aux_to_gpu, offload_aux
from modules.logger import log
from modules.caption.logits import LogitsParser
from modules.caption import helpers


debug_enabled = os.environ.get('SD_LLM_DEBUG', None) is not None
debug_log = log.trace if debug_enabled else lambda *args, **kwargs: None


def b64(image):
    if image is None:
        return ''
    if isinstance(image, gr.Image): # should not happen
        return None
    with io.BytesIO() as stream:
        image.convert('RGB').save(stream, 'JPEG')
        values = stream.getvalue()
        encoded = base64.b64encode(values).decode()
        return encoded


def is_cloud_model(model_name: str) -> bool:
    if not model_name:
        return False
    return model_name in Options.cloud


def is_vision_model(model_name: str) -> bool:
    """Check if model supports vision/image input."""
    if not model_name:
        return False
    return model_name in Options.img2img or model_name in Options.cloud


def is_thinking_model(model_name: str) -> bool:
    """Check if model supports thinking/reasoning mode."""
    if not model_name:
        return False
    model_lower = model_name.lower()
    # Match VQA's detection patterns for consistency
    thinking_indicators = [
        'thinking',      # Qwen3-VL-*-Thinking models
        'reasoning',     # Ministral-3-*-Reasoning models
        'moondream3',    # Moondream 3 supports thinking
        'moondream 3',
        'moondream2',    # Moondream 2 supports reasoning mode
        'moondream 2',
        'mimo',          # XiaomiMiMo models
        'qwen3.5',      # Qwen3.5 native thinking (repo names)
        'qwen 3.5',     # Qwen3.5 native thinking (display names)
    ]
    return any(indicator in model_lower for indicator in thinking_indicators)


def get_model_display_name(model_repo: str) -> str:
    """Generate display name with vision/reasoning symbols."""
    symbols = []
    if model_repo in Options.img2img:
        symbols.append(ui_symbols.vision)
    if model_repo in Options.cloud:
        symbols.append(ui_symbols.cloud)
    if is_thinking_model(model_repo):
        symbols.append(ui_symbols.reasoning)
    return f"{model_repo} {' '.join(symbols)}" if symbols else model_repo


def get_model_repo_from_display(display_name: str) -> str:
    """Strip symbols from display name to get repo."""
    if not display_name:
        return display_name
    result = display_name
    for symbol in [ui_symbols.vision, ui_symbols.reasoning, ui_symbols.cloud]:
        result = result.replace(symbol, '')
    return result.strip()


def keep_think_block_open(text_prompt: str) -> str:
    """Remove closing </think> so model can continue reasoning with prefill."""
    think_open = "<think>"
    think_close = "</think>"
    last_open = text_prompt.rfind(think_open)
    if last_open == -1:
        return text_prompt
    close_index = text_prompt.find(think_close, last_open)
    if close_index == -1:
        return text_prompt
    end_close = close_index + len(think_close)
    while end_close < len(text_prompt) and text_prompt[end_close] in ' \t\r\n':
        end_close += 1
    return text_prompt[:close_index] + text_prompt[end_close:]


@dataclass
class Options:
    img2img = [
        # Gemma
        'google/gemma-3-4b-it',
        'google/gemma-3n-E2B-it',
        'google/gemma-3n-E4B-it',
        'google/gemma-4-E2B-it',
        'google/gemma-4-E4B-it',
        'google/gemma-4-12B-it-qat-w4a16-ct',
        # Qwen3.5
        'Qwen/Qwen3.5-2B',
        'Qwen/Qwen3.5-4B',
        'Qwen/Qwen3.5-9B',
        # Qwen3-VL
        'Qwen/Qwen3-VL-2B-Instruct',
        'Qwen/Qwen3-VL-2B-Thinking',
        'Qwen/Qwen3-VL-4B-Instruct',
        'Qwen/Qwen3-VL-4B-Thinking',
        'Qwen/Qwen3-VL-8B-Instruct',
        'Qwen/Qwen3-VL-8B-Thinking',
        # Qwen2.5-VL
        'Qwen/Qwen2.5-VL-3B-Instruct',
        # Mistral
        'mistralai/Ministral-3-3B-Instruct-2512-BF16',
        'mistralai/Ministral-3-8B-Instruct-2512-BF16',
        'mistralai/Ministral-3-3B-Reasoning-2512',
        'mistralai/Ministral-3-8B-Reasoning-2512',
        # Finetunes
        'trohrbaugh/gemma-4-E4B-it-heretic-ara',
        'trohrbaugh/Qwen3.5-9B-heretic-v2',
    ]
    cloud = [
        'google/gemini-3.5-flash',
        'google/gemini-3.1-pro-preview',
        'google/gemini-3.1-flash-lite',
        'google/gemini-3.1-flash-lite-preview',
        'google/gemini-2.5-flash',
        'google/gemini-2.5-flash-lite',
        'google/gemini-2.5-pro',
    ]
    models = {
        # Gemma
        'google/gemma-3-1b-it': {},
        'google/gemma-3-4b-it': {},
        'google/gemma-3n-E2B-it': {},
        'google/gemma-3n-E4B-it': {},
        'google/gemma-4-E2B-it': {},
        'google/gemma-4-E4B-it': {},
        'google/gemma-4-12B-it-qat-w4a16-ct': {}, # compressed-tensor model
        # Qwen3.5
        'Qwen/Qwen3.5-0.8B': {},
        'Qwen/Qwen3.5-2B': {},
        'Qwen/Qwen3.5-4B': {},
        'Qwen/Qwen3.5-9B': {},
        # Qwen3
        'Qwen/Qwen3-0.6B': {},
        'Qwen/Qwen3-1.7B': {},
        'Qwen/Qwen3-4B': {},
        'Qwen/Qwen3-4B-Instruct-2507': {},
        # Qwen3-VL
        'Qwen/Qwen3-VL-2B-Instruct': {},
        'Qwen/Qwen3-VL-2B-Thinking': {},
        'Qwen/Qwen3-VL-4B-Instruct': {},
        'Qwen/Qwen3-VL-4B-Thinking': {},
        'Qwen/Qwen3-VL-8B-Instruct': {},
        'Qwen/Qwen3-VL-8B-Thinking': {},
        # Qwen2.5
        'Qwen/Qwen2.5-0.5B-Instruct': {},
        'Qwen/Qwen2.5-1.5B-Instruct': {},
        'Qwen/Qwen2.5-3B-Instruct': {},
        # Qwen2.5-VL
        'Qwen/Qwen2.5-VL-3B-Instruct': {},
        # Llama
        'meta-llama/Llama-3.2-1B-Instruct': {},
        'meta-llama/Llama-3.2-3B-Instruct': {},
        'meta-llama/Llama-3.2-8B-Instruct': {},
        'cognitivecomputations/Dolphin3.0-Llama3.2-1B': {},
        'cognitivecomputations/Dolphin3.0-Llama3.2-3B': {},
        # Gemini
        'google/gemini-3.5-flash': {},
        'google/gemini-3.1-pro-preview': {},
        'google/gemini-3.1-flash-lite': {},
        'google/gemini-3.1-flash-lite-preview': {},
        'google/gemini-2.5-flash': {},
        'google/gemini-2.5-flash-lite': {},
        'google/gemini-2.5-pro': {},
        # SmolLM
        'HuggingFaceTB/SmolLM2-135M-Instruct': {},
        'HuggingFaceTB/SmolLM2-360M-Instruct': {},
        'HuggingFaceTB/SmolLM2-1.7B-Instruct': {},
        'HuggingFaceTB/SmolLM3-3B': {},
        # Phi
        'microsoft/Phi-4-mini-instruct': {},
        # Mistral
        'mistralai/Ministral-3-3B-Instruct-2512-BF16': {},
        'mistralai/Ministral-3-8B-Instruct-2512-BF16': {},
        'mistralai/Ministral-3-3B-Reasoning-2512': {},
        'mistralai/Ministral-3-8B-Reasoning-2512': {},
        # Finetunes
        'p-e-w/gemma-4-E2B-it-heretic-ara': {},
        'trohrbaugh/gemma-4-E4B-it-heretic-ara': {},
        'trohrbaugh/Qwen3.5-9B-heretic-v2': {},
        # GGUF
        'mradermacher/Llama-3.2-1B-Instruct-Uncensored-i1-GGUF': { # kept primarily as an example how to add gguf model
            'repo': 'meta-llama/Llama-3.2-1B-Instruct', # original repo so we can load missing components
            'type': 'llama', # required so gguf loader knows what to do
            'gguf': 'mradermacher/Llama-3.2-1B-Instruct-Uncensored-i1-GGUF', # gguf repo
            'file': 'Llama-3.2-1B-Instruct-Uncensored.i1-Q4_0.gguf', # gguf file inside repo
        },
    }
    models_cls = {
        'qwen3_5': 'Qwen3_5ForConditionalGeneration',
        'qwen3_5_moe': 'Qwen3_5MoeForConditionalGeneration',
        'qwen3_vl': 'Qwen3VLForConditionalGeneration',
        'qwen2_5_vl': 'Qwen2_5_VLForConditionalGeneration',
        'qwen2_vl': 'Qwen2VLForConditionalGeneration',
        'mistral3': 'Mistral3ForConditionalGeneration',
        'gemma4': 'Gemma4ForConditionalGeneration',
    }

    # default = list(models)[1] # gemma-3-4b-it
    default = 'google/gemma-3-4b-it'
    supported = list(transformers.integrations.ggml.GGUF_CONFIG_MAPPING)
    t2i_prompt: str = textwrap.dedent('''\
        You are an image prompt expert.
        You will receive a user prompt for image generation.
        Improve the prompt by adding relevant visual specificity for composition, lighting, color, texture, and atmosphere.
        Keep the result faithful to the original prompt and the intended image.
        Do not add unrelated concepts, non-visual commentary, or fluff.
        ''')
    i2i_prompt: str = textwrap.dedent('''\
        You are an image prompt expert.
        You will receive an image and a user prompt for editing or refinement.
        Improve the prompt with concrete visual detail that remains faithful to the image and edit intent.
        Keep the result grounded in image-generation language.
        Do not invent unrelated objects, actions, or concepts.
        ''')
    i2i_noprompt: str = textwrap.dedent('''\
        You are an image prompt expert.
        You will receive an image and must create an editing prompt from the visible content.
        Base the prompt on what is present in the image, its style, and its atmosphere.
        Use concrete, image-generation language and avoid unrelated fluff or commentary.
        ''')
    nsfw_ok: str = textwrap.dedent('''\
        NSFW, nudity and explicit content and language are allowed.
        ''')
    nsfw_no: str = textwrap.dedent('''\
        NSFW, nudity or explicit content and language are not allowed; if present, remove those elements from the enhanced prompt.
        ''')
    details_prompt: str = textwrap.dedent('''\
        Enhance the prompt with relevant image details.
        Identify the main subject(s), physical attributes, clothing if non-nude, and distinctive visual features.
        Describe the scene, environment, key objects, and overall mood or atmosphere.
        Be precise and efficient; avoid redundancy, abstract commentary, unrelated fluff, or instructions.
        Do not invent any objects, settings, or themes not implied by the input.
        Do not add era, background props, or atmosphere unless explicitly present in the prompt.
        ''')
    details_format: str = textwrap.dedent('''\
        Output exactly one enhanced prompt string.
        Do not add greetings, comments, explanations, follow-up questions, labels, formatting, or numbering.
        Do not include any extra prose or analysis.
        Start immediately with the prompt content.
        No stray tokens!
        ''')

    censored = ["i cannot", "i can't", "i am sorry", "against my programming", "i am not able", "i am unable", 'i am not allowed']

    max_delim_index: int = 60
    min_tokens: int = 0
    max_tokens: int = 256
    do_sample: bool = True
    temperature: float = 0.6
    repetition_penalty: float = 1.2
    top_k: int = 0
    top_p: float = 0.0
    thinking_mode: bool = False

    @staticmethod
    def get_model_choices():
        """Return list of display names for dropdown."""
        return [get_model_display_name(repo) for repo in Options.models.keys()]

    @staticmethod
    def get_default_display():
        """Return display name for default model."""
        return get_model_display_name(Options.default)


class PromptEnhanceScript(scripts_manager.Script):
    prompt: gr.Textbox = None
    image: gr.Image = None
    model: str = None
    llm: transformers.AutoModelForCausalLM = None
    processor: transformers.AutoProcessor = None
    tokenizer: transformers.AutoTokenizer = None
    busy: bool = False
    server = None
    options = Options()

    def title(self):
        return 'Prompt enhance'

    def show(self, _is_img2img):
        return scripts_manager.AlwaysVisible

    def compile(self):
        if self.llm is None or 'LLM' not in shared.opts.cuda_compile:
            return
        from modules.sd_models_compile import compile_torch
        self.llm = compile_torch(self.llm, apply_to_components=False, op="LLM")

    def load(self, name:str | None=None, use_openai:bool=False, model_repo:str | None=None, model_gguf:str | None=None, model_type:str | None=None, model_file:str | None=None):
        # Strip symbols from display name if present
        name = get_model_repo_from_display(name) if name else self.options.default
        if self.busy:
            log.debug('Prompt enhance: busy')
            return model_repo
        if is_cloud_model(name):
            return model_repo
        if (self.model is not None) and (self.model == name):
            return model_repo

        model_repo = sd_models.path_to_repo(model_repo) if model_repo else None

        self.busy = True
        from modules import modelloader, model_quant, ggml
        modelloader.hf_login()
        model_repo = model_repo or self.options.models.get(name, {}).get('repo', None) or name
        model_gguf = model_gguf or self.options.models.get(name, {}).get('gguf', None) or model_repo
        model_type = model_type or self.options.models.get(name, {}).get('type', None)
        model_file = model_file or self.options.models.get(name, {}).get('file', None)
        model_subfolder = self.options.models.get(name, {}).get('subfolder', None)
        model_tokenizer = self.options.models.get(name, {}).get('tokenizer', None)

        gguf_args = {}
        if model_type is not None and model_file is not None and len(model_type) > 2 and len(model_file) > 2:
            debug_log(f'Prompt enhance: gguf supported={self.options.supported}')
            if model_type not in self.options.supported:
                log.error(f'Prompt enhance: name="{name}" repo="{model_repo}" fn="{model_file}" type={model_type} gguf not supported')
                log.trace(f'Prompt enhance: gguf supported={self.options.supported}')
                self.busy = False
                return model_repo
            ggml.install_gguf()
            gguf_args['model_type'] = model_type
            gguf_args['gguf_file'] = model_file

        quant_args = model_quant.create_config(module='LLM') if not gguf_args else {}

        try:
            t0 = time.time()
            if self.llm is not None:
                deregister_aux('prompt_enhance')
                sd_models.move_model(self.llm, devices.cpu, force=True)
                self.llm = None
                self.tokenizer = None
                self.processor = None
                devices.torch_gc(force=True, reason='prompt-enhance:load')
                log.debug(f'Prompt enhance: unload="{self.model}"')
            self.model = None
            load_args = { 'pretrained_model_name_or_path': model_repo if not gguf_args else model_gguf }
            if model_subfolder:
                load_args['subfolder'] = model_subfolder # Comma was incorrect here

            model_config = transformers.AutoConfig.from_pretrained(load_args['pretrained_model_name_or_path'], trust_remote_code=True, cache_dir=shared.opts.hfcache_dir)
            model_type = getattr(model_config, 'model_type', '')
            cls_name = transformers.AutoModelForCausalLM
            custom_cls_name = self.options.models_cls.get(model_type, None)
            if custom_cls_name:
                custom_cls = getattr(transformers, custom_cls_name, None)
                if custom_cls:
                    cls_name = custom_cls

            log.info(f'Prompt enhance load: name="{name}" repo="{model_repo}" cls={cls_name.__name__}')

            if '-ct' in model_repo.lower():
                from installer import install
                install('compressed-tensors')
                quant_args = {}

            sd_models.set_caption_load_options()
            try:
                self.llm = cls_name.from_pretrained(
                    **load_args,
                    trust_remote_code=True,
                    torch_dtype=devices.dtype,
                    low_cpu_mem_usage=True,
                    cache_dir=shared.opts.hfcache_dir,
                    # _attn_implementation="eager",
                    **gguf_args,
                    **quant_args,
                )
            finally:
                sd_models.set_huggingface_options(quiet=True)
            self.llm.eval()
            register_aux('prompt_enhance', self.llm)
            tokenizer_args = { 'pretrained_model_name_or_path': model_repo }
            if model_tokenizer:
                tokenizer_args['subfolder'] = model_tokenizer
            self.tokenizer = transformers.AutoTokenizer.from_pretrained(**tokenizer_args, cache_dir=shared.opts.hfcache_dir)
            if model_repo in self.options.img2img:
                self.processor = transformers.AutoProcessor.from_pretrained(**tokenizer_args, cache_dir=shared.opts.hfcache_dir)

            if debug_enabled:
                modules = sd_modules.get_model_stats(self.llm) + sd_modules.get_model_stats(self.tokenizer)
                for m in modules:
                    debug_log(f'Prompt enhance: {m}')
            self.model = name
            t1 = time.time()
            log.debug(f'Prompt enhance: cls={self.llm.__class__.__name__} name="{name}" repo="{model_repo}" fn="{model_file}" processor="{self.processor.__class__.__name__ if self.processor else None}" tokenizer="{self.tokenizer.__class__.__name__ if self.tokenizer else None}" time={t1-t0:.2f} loaded')
            self.compile()
        except Exception as e:
            log.error(f'Prompt enhance: load {e}')
            if debug_enabled:
                errors.display(e, 'Prompt enhance')

        devices.torch_gc()
        self.set_openai(enable=use_openai)
        self.busy = False
        return model_repo

    def censored(self, response):
        text = response.lower().replace("i'm", "i am")
        return any(c.lower() in text for c in self.options.censored)

    def unload(self):
        if self.llm is not None:
            model_name = self.model
            self.set_openai(enable=False)
            log.debug(f'Prompt enhance: unloading model="{model_name}"')
            deregister_aux('prompt_enhance')
            sd_models.move_model(self.llm, devices.cpu, force=True)
            self.model = None
            self.llm = None
            self.tokenizer = None
            self.processor = None
            devices.torch_gc(force=True, reason='prompt-enhance:unload')
            log.debug(f'Prompt enhance: model="{model_name}" unloaded')
        else:
            log.debug('Prompt enhance: no model loaded')

    def set_openai(self, enable: bool):
        from modules.openai.serve import OpenAIServer
        if enable and self.llm is not None and self.tokenizer is not None:
            self.server = OpenAIServer(
                model=self.llm,
                tokenizer=self.tokenizer,
                host="127.0.0.1",
                port=8000,
                server=shared.api.app,
            )
            self.server.start()
        elif self.server is not None:
            self.server.stop()
            self.server = None

    def clean(self, response, keep_thinking=False, prefill_text='', keep_prefill=False):
        # Handle thinking tags FIRST (before generic tag removal)
        if '<think>' in response or '</think>' in response:
            if keep_thinking:
                # Format: handle partial tags (</think> without <think> means thinking was in prompt)
                if '</think>' in response and '<think>' not in response:
                    response = 'Reasoning:\n' + response.replace('</think>', '\n\nAnswer:\n')
                else:
                    response = response.replace('<think>', 'Reasoning:\n').replace('</think>', '\n\nAnswer:\n')
            else:
                # Strip all thinking content
                response = re.sub(r'<think>.*?</think>', '', response, flags=re.DOTALL)
                response = response.replace('</think>', '')  # Handle orphaned closing tags

        # remove special characters
        response = response.replace('"', '').replace("'", "").replace('"', '').replace('"', '').replace('**', '')
        # remove repeating characters and short repeated tokens from model collapse
        response = response.replace('\n\n', '\n').replace('  ', ' ').replace('...', '.')
        response = re.sub(r'\b([A-Za-z]{1,3})(?:\s+\1){1,}\b', r'\1', response, flags=re.IGNORECASE)

        # remove comments between brackets (but not Reasoning:/Answer: which we may have added)
        response = re.sub(r'<.*?>', '', response)
        response = re.sub(r'\[.*?\]', '', response)
        response = re.sub(r'\/.*?\/', '', response)

        # remove llm commentary
        removed = ''
        if response.startswith('Prompt'):
            removed, response = response.split('Prompt', maxsplit=1)
        if 0 <= response.find(':') < self.options.max_delim_index:
            # Don't split on "Reasoning:" or "Answer:" if we're keeping thinking
            colon_pos = response.find(':')
            prefix_text = response[:colon_pos].strip()
            if not keep_thinking or (prefix_text not in ['Reasoning', 'Answer']):
                removed, response = response.split(':', maxsplit=1)
        if 0 <= response.find('---') < self.options.max_delim_index:
            response, removed = response.split('---', maxsplit=1)
        if len(removed) > 0:
            debug_log(f'Prompt enhance: max={self.options.max_delim_index} removed="{removed}"')

        # remove bullets and lists
        lines = [re.sub(r'^(\s*[-*]|\s*\d+)\s+', '', line).strip() for line in response.splitlines()]
        response = '\n'.join(lines)

        response = response.strip()

        # Remove leading conversational filler that some LLMs prepend
        response = re.sub(
            r'^(?:\s*(?:wait|okay|ok|sure|alright|yes|yep|hello|hi|thanks|thank you|no problem|of course|got it|i love|i like|i appreciate|great|excellent|right)[^.!?]*[.!?]\s*)+',
            '',
            response,
            flags=re.IGNORECASE,
        )

        # Handle prefill retention/removal
        prefill_text = (prefill_text or '').strip()
        if prefill_text:
            if keep_prefill:
                # Add prefill if it's missing from the cleaned response
                if not response.startswith(prefill_text):
                    sep = '' if (not response or response[0] in '.,!?;:') else ' '
                    response = f'{prefill_text}{sep}{response}'
            else:
                # Remove prefill if it's present in the cleaned response
                if response.startswith(prefill_text):
                    response = response[len(prefill_text):].strip()

        return response

    def post(self, response, prefix, suffix, networks):
        response = response.strip()
        prefix = prefix.strip()
        suffix = suffix.strip()
        if len(prefix) > 0:
            response = f'{prefix} {response}'
        if len(suffix) > 0:
            response = f'{response} {suffix}'
        if len(networks) > 0:
            response = f'{response} {" ".join(networks)}'
        return response

    def extract(self, prompt):
        pattern = r'(<.*?>)'
        matches = re.findall(pattern, prompt)
        filtered = re.sub(pattern, '', prompt)
        return filtered, matches

    def get_image(self, image):
        current_image = None
        try:
            if image is not None and isinstance(image, gr.Image):
                current_image = image.value
            elif image is not None and isinstance(image, Image.Image): # if image is already a PIL image
                current_image = image
            if current_image is not None and (current_image.width <= 64 or current_image.height <= 64):
                current_image = None
            # Fallback to Kanvas/Control input if no image from Gradio component (e.g., when Kanvas is active)
            if current_image is None and ui_control_helpers.input_source is not None:
                if isinstance(ui_control_helpers.input_source, list) and len(ui_control_helpers.input_source) > 0:
                    current_image = ui_control_helpers.input_source[0]
                elif isinstance(ui_control_helpers.input_source, Image.Image):
                    current_image = ui_control_helpers.input_source
        except Exception:
            current_image = None
        return current_image

    def enhance(self,
                model: str | None=None,
                prompt:str | None=None,
                system:str | None=None,
                prefix:str | None=None,
                suffix:str | None=None,
                sample:bool | None=None,
                min_tokens:int | None=None,
                max_tokens:int | None=None,
                temperature:float | None=None,
                penalty:float | None=None,
                top_k:int | None=None,
                top_p:float | None=None,
                thinking:bool=False,
                seed:int=-1,
                image=None,
                nsfw:bool | None=None,
                use_vision:bool=True,
                prefill:str='',
                keep_prefill:bool=False,
                keep_thinking:bool=False,
                custom_args:str | None=None,
                process_words:str='',
                semantic_threshold:float=0.0,
                embedding_similarity:float=0.0,
                use_openai:bool=False,
               ):
        # Strip symbols from model name if present
        model = get_model_repo_from_display(model) if model else self.options.default
        prompt = prompt or (self.prompt.value if self.prompt else "") # Check if self.prompt is None
        image = None
        if use_vision and is_vision_model(model): # handle vision toggle
            image = image or self.image
        if image is None:
            use_vision = False
        prefix = prefix or ''
        suffix = suffix or ''
        min_tokens = min_tokens or self.options.min_tokens
        max_tokens = max_tokens or self.options.max_tokens
        penalty = penalty or self.options.repetition_penalty
        temperature = temperature or self.options.temperature
        top_k = top_k if top_k is not None else self.options.top_k
        top_p = top_p if top_p is not None else self.options.top_p
        thinking = thinking or self.options.thinking_mode
        sample = sample if sample is not None else self.options.do_sample
        nsfw = nsfw if nsfw is not None else True # Default nsfw to True if not provided
        debug_log(f'Prompt enhance: model="{model}" model_class="{self.llm.__class__.__name__ if self.llm is not None else "not loaded"}" nsfw={nsfw} thinking={thinking} prefill="{prefill[:30] if prefill else ""}" use_vision={use_vision} image={image is not None}')

        while self.busy:
            time.sleep(0.1)

        if not is_cloud_model(model):
            self.load(model, use_openai=use_openai)

        if seed is None or seed == -1:
            random.seed()
            seed = int(random.randrange(4294967294))
        torch.manual_seed(seed)
        if (self.llm is None) and (not is_cloud_model(model)):
            log.error('Prompt enhance: model not loaded')
            return prompt
        prompt_text, networks = self.extract(prompt) # Use prompt_text after extraction
        debug_log(f'Prompt enhance: networks={networks}')

        current_image = None
        # Only process images if vision is enabled and model supports it
        if use_vision and is_vision_model(model):
            current_image = self.get_image(image)
        debug_log(f'Prompt enhance: image={current_image}')

        # Check if vision was requested but no image is available
        if use_vision and is_vision_model(model) and current_image is None:
            log.error(f'Prompt enhance: model="{model}" error="No input image provided"')
            return 'Error: No input image provided. Please upload or select an image.'

        # Resize large images to match VQA performance (Qwen3-VL performance is sensitive to resolution)
        # Create a copy to avoid modifying the original image used by img2img
        if current_image is not None and isinstance(current_image, Image.Image):
            original_size = (current_image.width, current_image.height)
            needs_resize = current_image.width > 768 or current_image.height > 768
            needs_rgb = current_image.mode != 'RGB'

            if needs_resize or needs_rgb:
                # Copy the image before any modifications to preserve the original
                current_image = current_image.copy()

                if needs_resize:
                    current_image.thumbnail((768, 768), Image.Resampling.LANCZOS)
                    debug_log(f'Prompt enhance: Resized image from {original_size} to {(current_image.width, current_image.height)}')

                if needs_rgb:
                    current_image = current_image.convert('RGB')
                    debug_log('Prompt enhance: Converted image to RGB mode')

        has_system = system is not None and len(system) > 4

        if current_image is not None and isinstance(current_image, Image.Image):
            if is_cloud_model(model):
                pass
            elif self.processor is None:
                log.error('Prompt enhance: image not supported by model')
                return prompt_text # Return original text part if image cannot be processed
            if prompt_text is not None and len(prompt_text) > 0:
                if not has_system:
                    system = self.options.i2i_prompt
                    system += self.options.nsfw_ok if nsfw else self.options.nsfw_no
                    system += self.options.details_prompt
                    system += self.options.details_format
                chat_template = [
                    { "role": "system", "content": [
                        {"type": "text", "text": system }
                    ] },
                    { "role": "user",   "content": [
                        {"type": "text", "text": prompt_text},
                        {"type": "image", "image": b64(current_image)}
                    ] },
                ]
            else:
                if not has_system:
                    system = self.options.i2i_noprompt
                    system += self.options.nsfw_ok if nsfw else self.options.nsfw_no
                    system += self.options.details_prompt
                    system += self.options.details_format
                chat_template = [
                    { "role": "system", "content": [
                        {"type": "text", "text": system }
                    ] },
                    { "role": "user",   "content": [
                        {"type": "image", "image": b64(current_image)}
                    ] },
                ]
        else:
            if not has_system:
                system = self.options.t2i_prompt
                system += self.options.nsfw_ok if nsfw else self.options.nsfw_no
                system += self.options.details_prompt
                system += self.options.details_format
            if self.processor is None:
                chat_template = [
                    { "role": "system", "content": system },
                    { "role": "user",   "content": prompt_text },
                ]
            else:
                chat_template = [
                    { "role": "system", "content": [
                        {"type": "text", "text": system }
                    ] },
                    { "role": "user",   "content": [
                        {"type": "text", "text": prompt_text},
                    ] },
                ]

        # Prepare prefill (VQA approach: string concatenation, not assistant message)
        prefill_text = (prefill or '').strip()
        use_prefill = len(prefill_text) > 0
        is_thinking = is_thinking_model(model)

        debug_log(f'Prompt enhance: system="{system}"')
        debug_log(f'Prompt enhance: prompt="{prompt_text}"')
        debug_log(f'Prompt template: roles={[msg["role"] for msg in chat_template]} thinking={is_thinking}:{thinking} prefill={use_prefill}')
        t0 = time.time()
        self.busy = True

        if is_cloud_model(model):
            if 'gemini' in model:
                from modules.caption import gemini
                kwargs = {
                    'temperature': temperature,
                    'min_output_tokens': min_tokens,
                    'max_output_tokens': max_tokens,
                }
                model_name = model.replace('google/', '')
                response = gemini.predict(prompt_text, current_image, model_name, system, model, prefill_text, thinking, kwargs)
                t1 = time.time()
                log.info(f'Prompt enhance: model="{model}" nsfw={nsfw} time={t1-t0:.2f} prefill="{prefill_text[:20] if prefill_text else None}" response={len(response)}')
                debug_log(f'Prompt enhance: response="{response}"')
                self.busy = False
                return response

            else:
                return 'Model not recognized'

        try:
            # Qwen3.5 uses native enable_thinking parameter in the chat template
            is_qwen35 = 'qwen3.5' in model.lower()
            template_kwargs = {'enable_thinking': thinking} if is_qwen35 else {}

            # Generate text prompt using template
            apply_fn = self.processor if self.processor is not None else self.tokenizer
            try:
                text_prompt = apply_fn.apply_chat_template(
                    chat_template,
                    add_generation_prompt=True,
                    tokenize=False,
                    **template_kwargs,
                )
            except TypeError:
                text_prompt = apply_fn.apply_chat_template(
                    chat_template,
                    tokenize=False,
                )

            # Manual think handling - skip for Qwen3.5 (template handles it natively)
            if is_thinking and not is_qwen35:
                if not thinking:
                    # User wants to SKIP thinking
                    # Template opened the block with <think>, close it immediately
                    text_prompt += "</think>\n"
                    if use_prefill:
                        text_prompt += prefill_text
                    debug_log('Prompt enhance: forced thinking off, appended </think>')
                else:
                    # User wants thinking - prefill becomes part of thought process
                    if use_prefill:
                        text_prompt += prefill_text
                    debug_log('Prompt enhance: thinking enabled, prefill inside think block')
            else:
                # Standard model or Qwen3.5 (no manual <think> manipulation needed)
                if use_prefill:
                    text_prompt += prefill_text

            # debug_log(f'Prompt enhance: template="{text_prompt}"')

            # Tokenize the final prompt
            # For VL models with images, pass the image to the processor (like VQA does)
            if self.processor is not None and current_image is not None:
                inputs = self.processor(text=[text_prompt], images=[current_image], padding=True, return_tensors="pt")
            elif self.processor is not None:
                # VL processor without image - must use explicit text= parameter
                inputs = self.processor(text=[text_prompt], images=None, padding=True, return_tensors="pt")
            else:
                inputs = self.tokenizer(text_prompt, return_tensors="pt")
            inputs = inputs.to(devices.device).to(devices.dtype)
            input_len = inputs['input_ids'].shape[1]
        except Exception as e:
            log.error(f'Prompt enhance tokenize: {e}')
            if debug_enabled:
                errors.display(e, 'Prompt enhance')
            self.busy = False
            return prompt_text # Return original text part on error

        try:
            with devices.inference_context():
                move_aux_to_gpu('prompt_enhance')
                gen_kwargs = {
                    'do_sample': sample,
                    'temperature': float(temperature),
                    'max_new_tokens': int(max_tokens),
                    'repetition_penalty': float(penalty),
                }
                if min_tokens > 0:
                    gen_kwargs['min_new_tokens'] = int(min_tokens)
                if top_k > 0:
                    gen_kwargs['top_k'] = int(top_k)
                if top_p > 0:
                    gen_kwargs['top_p'] = float(top_p)

                logits_processor = None
                if process_words is not None and len(process_words.strip()) > 0 and self.tokenizer is not None:
                    logits_processor = LogitsParser(self.tokenizer, process_words, semantic_threshold=semantic_threshold, embedding_similarity=embedding_similarity)
                    gen_kwargs['logits_processor'] = [logits_processor]

                custom = helpers.get_custom_args(self.llm, custom_args)
                for k, v in custom.items():
                    gen_kwargs[k] = v

                log.debug(f'Prompt enhance: cls={self.llm.__class__.__name__} model="{model}" tokens={input_len} args={gen_kwargs} custom={custom}')
                defaults = {k: v for k, v in helpers.get_default_args(self.llm).items() if k not in gen_kwargs}
                log.debug(f'Prompt enhance: defaults={defaults}')

                outputs = self.llm.generate(**inputs, **gen_kwargs)

                if logits_processor is not None:
                    log.debug(f'Prompt enhance: process={logits_processor.get_replacements()}')

            outputs_cropped = outputs[:, input_len:]
            decode_fn = self.processor if self.processor is not None else self.tokenizer
            response = decode_fn.batch_decode(
                outputs_cropped,
                skip_special_tokens=True,
                clean_up_tokenization_spaces=True,
            )
            if debug_enabled:
                response_before_clean = response[0] if isinstance(response, list) else response
                debug_log(f'Prompt enhance: response_before_clean="{response_before_clean}"')
        except Exception as e:
            outputs = None
            log.error(f'Prompt enhance generate: {e}')
            if debug_enabled:
                errors.display(e, 'Prompt enhance')
            self.busy = False
            response = f'Error: {str(e)}'
        finally:
            offload_aux('prompt_enhance')
            devices.torch_gc(force=False, reason='prompt-enhance')
        t1 = time.time()

        if isinstance(response, list):
            response = response[0]
        is_censored =  self.censored(response)
        if not is_censored:
            response = self.clean(response, keep_thinking=keep_thinking, prefill_text=prefill_text, keep_prefill=keep_prefill)
            response = self.post(response, prefix, suffix, networks)
        log.info(f'Prompt enhance: model="{model}" nsfw={nsfw} time={t1-t0:.2f} seed={seed} thinking={thinking} keep={keep_thinking}:{keep_prefill} prefill="{prefill_text[:20] if prefill_text else None}" inputs={input_len} outputs={outputs.shape[-1] if isinstance(outputs, torch.Tensor) else 0} prompt={len(prompt_text)} response={len(response)}')
        debug_log(f'Prompt enhance: prompt="{prompt_text}"')
        debug_log(f'Prompt enhance: response_after_clean="{response}"')
        self.busy = False
        if is_censored:
            log.warning(f'Prompt enhance: censored response="{response}"')
            return prompt # Return original full prompt on censorship
        return response

    def apply(self, prompt, image, apply_prompt, llm_model, prompt_system, prompt_prefix, prompt_suffix, min_tokens, max_tokens, do_sample, temperature, repetition_penalty, top_k, top_p, thinking_mode, nsfw_mode, use_vision, prefill_text, keep_prefill, keep_thinking, custom_args, process_words, semantic_threshold, embedding_similarity, use_openai):
        response = self.enhance(
            prompt=prompt,
            image=image,
            prefix=prompt_prefix,
            suffix=prompt_suffix,
            model=llm_model,
            system=prompt_system,
            sample=do_sample,
            min_tokens=min_tokens,
            max_tokens=max_tokens,
            temperature=temperature,
            penalty=repetition_penalty,
            top_k=top_k,
            top_p=top_p,
            thinking=thinking_mode,
            nsfw=nsfw_mode,
            use_vision=use_vision,
            prefill=prefill_text,
            keep_prefill=keep_prefill,
            keep_thinking=keep_thinking,
            custom_args=custom_args,
            process_words=process_words,
            semantic_threshold=semantic_threshold,
            embedding_similarity=embedding_similarity,
            use_openai=use_openai,
        )
        if apply_prompt:
            return [response, response]
        return [response, gr.update()]

    def get_custom(self, name):
        # Strip symbols from display name to get repo
        repo_name = get_model_repo_from_display(name)
        model_repo = self.options.models.get(repo_name, {}).get('repo', None) or repo_name
        model_gguf = self.options.models.get(repo_name, {}).get('gguf', None)
        model_type = self.options.models.get(repo_name, {}).get('type', None)
        model_file = self.options.models.get(repo_name, {}).get('file', None)
        return [model_repo, model_gguf, model_type, model_file]

    def update_vision_toggle(self, model_name):
        """Update vision toggle interactivity and value based on model selection."""
        repo_name = get_model_repo_from_display(model_name)
        is_vl = is_vision_model(repo_name)
        if not is_vl:
            return gr.update(interactive=False, value=False)
        return gr.update(interactive=is_vl)

    def ui(self, _is_img2img):
        with gr.Accordion('Prompt enhance', open=False, elem_id='prompt_enhance'):
            gr.HTML('<style>#prompt_enhance_use_vision:has(input:disabled) { opacity: 0.5; }</style>')
            with gr.Row():
                apply_btn = gr.Button(value='Enhance now', elem_id='prompt_enhance_apply', variant='primary')
            with gr.Row():
                apply_prompt = gr.Checkbox(label='Apply to prompt', value=False)
                apply_auto = gr.Checkbox(label='Auto enhance', value=False)
            with gr.Row():
                # Set initial state based on whether default model supports vision
                default_is_vl = is_vision_model(Options.default)
                use_vision = gr.Checkbox(label='Use vision', value=False, interactive=default_is_vl, elem_id='prompt_enhance_use_vision')
                use_openai = gr.Checkbox(label='OpenAI interface', value=False, elem_id='prompt_enhance_openai')
            gr.HTML('<br>')
            with gr.Group():
                with gr.Row():
                    llm_model = gr.Dropdown(label='LLM model', choices=Options.get_model_choices(), value=Options.get_default_display(), interactive=True, allow_custom_value=True, elem_id='prompt_enhance_model')
                with gr.Row():
                    load_btn = gr.Button(value='Load model', elem_id='prompt_enhance_load', variant='secondary')
                    load_btn.click(fn=self.load, inputs=[llm_model, use_openai], outputs=[])
                    unload_btn = gr.Button(value='Unload model', elem_id='prompt_enhance_unload', variant='secondary')
                    unload_btn.click(fn=self.unload, inputs=[], outputs=[])
                with gr.Accordion('Custom model', open=False, elem_id='prompt_enhance_custom'):
                    with gr.Row():
                        model_repo = gr.Textbox(label='Model repo', value=None, interactive=True, elem_id='prompt_enhance_model_repo', placeholder='Original model repo on huggingface')
                    with gr.Row():
                        model_gguf = gr.Textbox(label='Model gguf', value=None, interactive=True, elem_id='prompt_enhance_model_gguf', placeholder='Optional GGUF model repo on huggingface')
                    with gr.Row():
                        model_type = gr.Textbox(label='Model type', value=None, interactive=True, elem_id='prompt_enhance_model_type', placeholder='Optional GGUF model type')
                    with gr.Row():
                        model_file = gr.Textbox(label='Model file', value=None, interactive=True, elem_id='prompt_enhance_model_file', placeholder='Optional GGUF model file inside GGUF model repo')
                    with gr.Row():
                        custom_btn = gr.Button(value='Load custom model', elem_id='prompt_enhance_custom_load', variant='secondary')
                        custom_btn.click(fn=self.load, inputs=[model_repo, use_openai, model_repo, model_gguf, model_type, model_file], outputs=[llm_model])
                        llm_model.change(fn=self.get_custom, inputs=[llm_model], outputs=[model_repo, model_gguf, model_type, model_file])
                        gr.HTML('<br>')
                with gr.Accordion('Options', open=False, elem_id='prompt_enhance_options'):
                    with gr.Row():
                        min_tokens = gr.Slider(label='Min tokens', value=self.options.min_tokens, minimum=0, maximum=4096, step=1, interactive=True)
                        max_tokens = gr.Slider(label='Max tokens', value=self.options.max_tokens, minimum=10, maximum=4096, step=1, interactive=True)
                        do_sample = gr.Checkbox(label='Use samplers', value=self.options.do_sample, interactive=True)
                    with gr.Row():
                        temperature = gr.Slider(label='Temperature', value=self.options.temperature, minimum=0.0, maximum=1.0, step=0.01, interactive=True)
                        repetition_penalty = gr.Slider(label='Repetition penalty', value=self.options.repetition_penalty, minimum=0.0, maximum=2.0, step=0.01, interactive=True)
                    with gr.Row():
                        top_k = gr.Slider(label='Top-K', value=self.options.top_k, minimum=0, maximum=100, step=1, interactive=True)
                        top_p = gr.Slider(label='Top-P', value=self.options.top_p, minimum=0.0, maximum=1.0, step=0.01, interactive=True)
                    with gr.Row():
                        nsfw_mode = gr.Checkbox(label='NSFW allowed', value=True, interactive=True)
                        thinking_mode = gr.Checkbox(label='Thinking mode', value=False, interactive=True)
                    with gr.Row():
                        keep_thinking = gr.Checkbox(label='Keep Thinking Trace', value=False, interactive=True)
                        keep_prefill = gr.Checkbox(label='Keep Prefill', value=False, interactive=True)
                    with gr.Row():
                        custom_args = gr.Textbox(label='Custom arguments', value='', placeholder='Optional: custom arguments for the model as k=v, semicolon delimited', interactive=True, lines=1)
                    with gr.Row():
                        prefill_text = gr.Textbox(label='Prefill text', value='', placeholder='Optional: pre-fill start of model response', interactive=True, lines=1)
                    gr.HTML('<br>')
                with gr.Accordion('Input', open=False, elem_id='prompt_enhance_system_prompt'): # Corrected elem_id reference
                    with gr.Row():
                        prompt_prefix = gr.Textbox(label='Prompt prefix', value='', placeholder='Text prepended to the enhanced result', interactive=True, lines=2, elem_id='prompt_enhance_prefix')
                    with gr.Row():
                        prompt_suffix = gr.Textbox(label='Prompt suffix', value='', placeholder='Text appended to the enhanced result', interactive=True, lines=2, elem_id='prompt_enhance_suffix')
                    with gr.Row():
                        prompt_system = gr.Textbox(label='System prompt', value='', placeholder='Leave empty to use built-in enhancement instructions', interactive=True, lines=4, elem_id='prompt_enhance_system')
                with gr.Accordion('Process', open=False, elem_id='prompt_enhance_logits'): # Corrected elem_id reference
                    with gr.Row():
                        process_words = gr.Textbox(label='Words to process', value='', placeholder='list of words with optional substitutions', interactive=True, lines=3, elem_id='prompt_enhance_process_words')
                    with gr.Row():
                        semantic_threshold = gr.Slider(label='Semantic threshold', value=0.0, minimum=0.0, maximum=1.0, step=0.01, interactive=True, elem_id='prompt_enhance_semantic_threshold')
                        embedding_similarity = gr.Slider(label='Embedding similarity', value=0.0, minimum=0.0, maximum=1.0, step=0.01, interactive=True, elem_id='prompt_enhance_embedding_similarity')
                with gr.Accordion('Output', open=True, elem_id='prompt_enhance_output'): # Corrected elem_id reference
                    with gr.Row():
                        prompt_output = gr.Textbox(label='Enhanced prompt', value='', placeholder='Enhanced prompt will appear here', interactive=True, lines=4, max_lines=12, elem_id='prompt_enhance_result')
                    with gr.Row():
                        clear_btn = gr.Button(value='Clear', elem_id='prompt_enhance_clear', variant='secondary')
                        clear_btn.click(fn=lambda: '', inputs=[], outputs=[prompt_output])
                        copy_btn = gr.Button(value='Set prompt', elem_id='prompt_enhance_copy', variant='secondary')
                        if self.prompt: # not registered for api script runner
                            copy_btn.click(fn=lambda x: x, inputs=[prompt_output], outputs=[self.prompt])
            if self.image is None:
                self.image = gr.Image(type='pil', interactive=False, visible=False, width=64, height=64) # dummy image
            # Update vision toggle interactivity when model changes
            llm_model.change(fn=self.update_vision_toggle, inputs=[llm_model], outputs=[use_vision], show_progress=False)
            if self.prompt:
                apply_btn.click(fn=self.apply, inputs=[self.prompt, self.image, apply_prompt, llm_model, prompt_system, prompt_prefix, prompt_suffix, min_tokens, max_tokens, do_sample, temperature, repetition_penalty, top_k, top_p, thinking_mode, nsfw_mode, use_vision, prefill_text, keep_prefill, keep_thinking, custom_args, process_words, semantic_threshold, embedding_similarity, use_openai], outputs=[prompt_output, self.prompt])
        return [self.prompt, self.image, apply_auto, llm_model, prompt_system, prompt_prefix, prompt_suffix, min_tokens, max_tokens, do_sample, temperature, repetition_penalty, top_k, top_p, thinking_mode, nsfw_mode, use_vision, prefill_text, keep_prefill, keep_thinking, custom_args, process_words, semantic_threshold, embedding_similarity, use_openai]

    def after_component(self, component, **_kwargs): # searching for actual ui prompt components
        if getattr(component, 'elem_id', '') in ['txt2img_prompt', 'img2img_prompt', 'control_prompt', 'video_prompt']:
            self.prompt = component
            self.prompt.use_original = True
        if getattr(component, 'elem_id', '') in ['img2img_image', 'control_input_select']:
            self.image = component
            self.image.use_original = True

    def before_process(self, p: processing.StableDiffusionProcessing, *args, **kwargs): # pylint: disable=unused-argument
        _self_prompt, self_image, apply_auto, llm_model, prompt_system, prompt_prefix, prompt_suffix, min_tokens, max_tokens, do_sample, temperature, repetition_penalty, top_k, top_p, thinking_mode, nsfw_mode, use_vision, prefill_text, keep_prefill, keep_thinking, custom_args, process_words, semantic_threshold, embedding_similarity, use_openai = args
        if not apply_auto and not p.enhance_prompt:
            return
        if shared.state.skipped or shared.state.interrupted:
            return
        p.prompt = shared.prompt_styles.apply_styles_to_prompt(p.prompt, p.styles)
        p.negative_prompt = shared.prompt_styles.apply_negative_styles_to_prompt(p.negative_prompt, p.styles)
        shared.prompt_styles.apply_styles_to_extra(p)
        p.styles = []
        jobid = shared.state.begin('LLM')
        p.extra_generation_params['LLM'] = get_model_repo_from_display(llm_model)
        p.extra_generation_params['Original'] = p.prompt
        p.prompt = self.enhance(
            prompt=p.prompt,
            seed=p.seed,
            image=self_image,
            prefix=prompt_prefix,
            suffix=prompt_suffix,
            model=llm_model,
            system=prompt_system,
            sample=do_sample,
            min_tokens=min_tokens,
            max_tokens=max_tokens,
            temperature=temperature,
            penalty=repetition_penalty,
            top_k=top_k,
            top_p=top_p,
            thinking=thinking_mode,
            nsfw=nsfw_mode,
            use_vision=use_vision,
            prefill=prefill_text,
            keep_prefill=keep_prefill,
            keep_thinking=keep_thinking,
            custom_args=custom_args,
            process_words=process_words,
            semantic_threshold=semantic_threshold,
            embedding_similarity=embedding_similarity,
            use_openai=use_openai,
        )
        timer.process.record('prompt')
        shared.state.end(jobid)