__all__ = [
    "Process_Image_LMMH"
]



import os
import torch
import warnings
import logging
from PIL import Image
from transformers import AutoModelForCausalLM, TextIteratorStreamer
from janus.models import MultiModalityCausalLM, VLChatProcessor
from janus.utils.io import load_pil_images
from datetime import datetime
from threading import Thread

from .File_Text_utils import Detect_Language_LMMH


# 🧹 Silence most logs and warnings
os.environ["TRANSFORMERS_VERBOSITY"] = "error"
os.environ["HF_HUB_DISABLE_PROGRESS_BARS"] = "1"
os.environ["TOKENIZERS_PARALLELISM"] = "false"

logging.getLogger("transformers").setLevel(logging.ERROR)
logging.getLogger("torch").setLevel(logging.ERROR)
warnings.filterwarnings("ignore")


# 🔇 Suppress stdout temporarily during model load
import sys
import contextlib
from io import StringIO

@contextlib.contextmanager
def suppress_output():
    """Temporarily suppress stdout and stderr."""
    with open(os.devnull, "w") as devnull:
        old_stdout, old_stderr = sys.stdout, sys.stderr
        sys.stdout, sys.stderr = devnull, devnull
        try:
            yield
        finally:
            sys.stdout, sys.stderr = old_stdout, old_stderr


# 🧠 Load model silently
model_pathJanus = "deepseek-ai/Janus-Pro-7B"

with suppress_output():
    vl_chat_processor: VLChatProcessor = VLChatProcessor.from_pretrained(model_pathJanus)
    tokenizer = vl_chat_processor.tokenizer

    deviceIMG = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    vl_gpt: MultiModalityCausalLM = AutoModelForCausalLM.from_pretrained(
        model_pathJanus, trust_remote_code=True
    )

# 🔧 Multi-GPU support
if torch.cuda.device_count() > 1:
    vl_gpt = torch.nn.DataParallel(vl_gpt)

vl_gpt = vl_gpt.to(torch.bfloat16).to(deviceIMG).eval()


def Process_Image_LMMH(value, question, SAVE_DIR):
    if value is None:
        raise ValueError("No image file provided")

    original_filename = os.path.splitext(value)[0]
    timestamp = datetime.now().strftime("%Y%m%d")
    file_name = os.path.splitext(os.path.basename(value))[0]
    txt_filename = f"{file_name}_{timestamp}.txt"
    directory_path = os.path.dirname(value)
    txt_path = os.path.join(directory_path, txt_filename)

    conversation = [
        {"role": "User", "content": f"<image_placeholder>\n{question}", "images": [value]},
        {"role": "Assistant", "content": ""},
    ]

    pil_images = load_pil_images(conversation)
    prepare_inputs = vl_chat_processor(
        conversations=conversation, images=pil_images, force_batchify=True
    ).to(deviceIMG)

    # Prepare model and input
    if torch.cuda.device_count() > 1:
        inputs_embeds = vl_gpt.module.prepare_inputs_embeds(**prepare_inputs)
        model_ref = vl_gpt.module.language_model
    else:
        inputs_embeds = vl_gpt.prepare_inputs_embeds(**prepare_inputs)
        model_ref = vl_gpt.language_model

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

    generate_kwargs = {
        "inputs_embeds": inputs_embeds,
        "attention_mask": prepare_inputs.attention_mask,
        "pad_token_id": tokenizer.eos_token_id,
        "bos_token_id": tokenizer.bos_token_id,
        "eos_token_id": tokenizer.eos_token_id,
        "max_new_tokens": 512,
        "do_sample": False,
        "use_cache": True,
        "streamer": streamer,
    }

    # Background generation
    response = model_ref.generate(**generate_kwargs)

    description_buffer = ""
    for token in streamer:
        description_buffer += token
        yield token

    # Save output
    timestamp_str = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
    with open(txt_path, 'a' if os.path.exists(txt_path) else 'w', encoding='utf-8') as txt_file:
        txt_file.write(f'\n\n{timestamp_str} >> {question}\n{timestamp_str} >> {description_buffer}')

