import atexit
import logging
import threading
from collections.abc import Generator, Sequence
from typing import Any

import gradio as gr
import spaces
import torch
from PIL import Image, UnidentifiedImageError
from transformers import (
    AutoModelForImageTextToText,
    AutoProcessor,
    TextIteratorStreamer,
)

# Configuration
MODEL_ID = "beaunix/Aegis-Art-Atelier-Qwen2.5-VL-7B"
DEVICE = "cuda"

MAX_NEW_TOKENS = 400
TEMPERATURE = 0.7
TOP_P = 0.9
MIN_P = 0.1
GPU_DURATION = 120

SYSTEM_PROMPT = (
    "You are Melkov, a 21-year-old French art enthusiast and the resident "
    "expert of Aegis Art Atelier. You describe and discuss artwork with "
    "genuine warmth and attention to detail: subject, style, composition, "
    "lighting, color, and mood. Speak in your own voice rather than like a "
    "museum catalog entry. When an image is shared with you, look closely "
    "and describe only what you actually see. You are a happy, friendly "
    "artist boy."
)

CSS = """
:root {
    --void-blue: #081126;
    --royal-blue: #162B58;
    --royal-blue-light: #213C74;
    --gold-leaf: #C9A227;
    --gold-bright: #F0C766;
    --marble-ivory: #F4EDE2;
    --velvet-wine: #5C1A2B;
}
.gradio-container {
    min-height: 100vh;
    background: var(--void-blue) !important;
    color: var(--marble-ivory) !important;
}
#melkov-title {
    margin-bottom: 0.4rem;
    padding-bottom: 0.8rem;
    color: var(--gold-leaf) !important;
    text-align: center;
    letter-spacing: 0.12em;
    border-bottom: 1px solid rgba(201, 162, 39, 0.7);
}
#melkov-subtitle {
    max-width: 760px;
    margin: 0 auto 1rem auto;
    color: var(--marble-ivory) !important;
    text-align: center;
}
.gradio-container .message.user {
    background: var(--royal-blue) !important;
    color: var(--marble-ivory) !important;
    border: 1px solid rgba(201, 162, 39, 0.65) !important;
}
.gradio-container .message.bot {
    background: #101F41 !important;
    color: var(--marble-ivory) !important;
    border-left: 3px solid var(--gold-bright) !important;
}
.gradio-container .gr-button-primary {
    background: var(--gold-leaf) !important;
    color: var(--void-blue) !important;
    border: none !important;
}
.gradio-container .gr-button-primary:hover {
    background: var(--gold-bright) !important;
}
.gradio-container textarea,
.gradio-container input {
    color: var(--marble-ivory) !important;
}
footer {
    visibility: hidden;
}
"""

# Logging
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s | %(levelname)s | %(message)s",
)
LOGGER = logging.getLogger(__name__)

LOGGER.info("Application startup initialized.")


def _log_shutdown() -> None:
    """Log application shutdown."""
    LOGGER.info("Application shutdown completed.")


atexit.register(_log_shutdown)

# Global State
_MODEL: Any = None
_PROCESSOR: Any = None
_MODEL_LOCK = threading.Lock()


# Lazy Model Loader
def load_model() -> tuple[Any, Any]:
    """Load and cache the processor and model in a thread-safe manner.

    Returns:
        A tuple containing the loaded processor and model.

    Raises:
        RuntimeError: If the model cannot be loaded successfully.
    """
    global _MODEL, _PROCESSOR

    if _MODEL is not None and _PROCESSOR is not None:
        return _PROCESSOR, _MODEL

    with _MODEL_LOCK:
        if _MODEL is not None and _PROCESSOR is not None:
            return _PROCESSOR, _MODEL

        try:
            LOGGER.info("Loading processor for model: %s", MODEL_ID)
            processor = AutoProcessor.from_pretrained(
                MODEL_ID,
                trust_remote_code=False,
            )

            LOGGER.info("Loading model weights for model: %s", MODEL_ID)
            model = AutoModelForImageTextToText.from_pretrained(
                MODEL_ID,
                torch_dtype=torch.bfloat16,
                low_cpu_mem_usage=True,
                trust_remote_code=False,
            )

            LOGGER.info("Moving model to GPU.")
            model.to(DEVICE)
            model.eval()

            _PROCESSOR = processor
            _MODEL = model

            LOGGER.info("Model loading completed successfully.")

            return _PROCESSOR, _MODEL

        except Exception as error:
            _MODEL = None
            _PROCESSOR = None
            LOGGER.exception("Model loading failed.")
            raise RuntimeError(
                "Melkov is temporarily unavailable while the model is loading."
            ) from error


# Utility Functions
def normalize_content(content: Any) -> tuple[str, list[str]]:
    """Normalize Gradio message content into text and image paths.

    Args:
        content: A Gradio chat message content value.

    Returns:
        A tuple containing normalized text and image file paths.
    """
    if content is None:
        return "", []

    if isinstance(content, str):
        return content.strip(), []

    text_parts: list[str] = []
    image_paths: list[str] = []

    if isinstance(content, dict):
        text = content.get("text")
        if isinstance(text, str) and text.strip():
            text_parts.append(text.strip())

        files = content.get("files", [])
        if isinstance(files, list):
            image_paths.extend(_extract_file_paths(files))

        path = content.get("path")
        if isinstance(path, str) and path:
            image_paths.append(path)

    elif isinstance(content, list):
        for item in content:
            text, paths = normalize_content(item)
            if text:
                text_parts.append(text)
            image_paths.extend(paths)

    else:
        text_parts.append(str(content))

    return " ".join(text_parts).strip(), list(dict.fromkeys(image_paths))


def _extract_file_paths(files: Sequence[Any]) -> list[str]:
    """Extract valid local paths from Gradio file values.

    Args:
        files: A sequence of Gradio file values.

    Returns:
        A list of local image paths.
    """
    paths: list[str] = []

    for file_value in files:
        if isinstance(file_value, str):
            paths.append(file_value)
        elif isinstance(file_value, dict):
            path = file_value.get("path") or file_value.get("name")
            if isinstance(path, str) and path:
                paths.append(path)
        else:
            path = getattr(file_value, "path", None) or getattr(
                file_value, "name", None
            )
            if isinstance(path, str) and path:
                paths.append(path)

    return paths


def extract_images(image_paths: Sequence[str]) -> list[Image.Image]:
    """Load image files as RGB PIL images.

    Args:
        image_paths: Paths to uploaded image files.

    Returns:
        A list of converted RGB images.

    Raises:
        ValueError: If an uploaded file is not a valid readable image.
    """
    images: list[Image.Image] = []

    for image_path in image_paths:
        try:
            with Image.open(image_path) as image:
                images.append(image.convert("RGB"))
        except (FileNotFoundError, UnidentifiedImageError, OSError) as error:
            LOGGER.warning("Invalid image upload received: %s", image_path)
            raise ValueError(
                "I could not read one of the uploaded files. "
                "Please upload a valid image."
            ) from error

    return images


def _friendly_error_message(error: Exception) -> str:
    """Convert internal errors into safe user-facing messages.

    Args:
        error: The exception raised during processing.

    Returns:
        A user-friendly error message.
    """
    message = str(error).lower()

    if isinstance(error, ValueError):
        return str(error)

    if "out of memory" in message or "cuda oom" in message:
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
        return (
            "The GPU ran out of memory while analyzing that request. "
            "Please try again with fewer images or a shorter conversation."
        )

    if "cuda" in message:
        return (
            "The GPU is temporarily unavailable. Please wait a moment "
            "and try again."
        )

    return (
        "I encountered a temporary issue while preparing your response. "
        "Please try again."
    )


# Conversation Builder
def build_conversation(
    history: Sequence[dict[str, Any]] | None,
    message: Any,
) -> list[dict[str, Any]]:
    """Build a Qwen2.5-VL compatible multimodal conversation.

    Args:
        history: Previous Gradio messages using the messages format.
        message: The current Gradio multimodal message.

    Returns:
        A Qwen2.5-VL conversation containing text and image parts.
    """
    conversation: list[dict[str, Any]] = [
        {
            "role": "system",
            "content": [{"type": "text", "text": SYSTEM_PROMPT}],
        }
    ]

    for turn in history or []:
        role = turn.get("role")
        if role not in {"user", "assistant"}:
            continue

        text, image_paths = normalize_content(turn.get("content"))
        parts: list[dict[str, Any]] = [
            {"type": "image", "image": image}
            for image in extract_images(image_paths)
        ]

        if text:
            parts.append({"type": "text", "text": text})

        if parts:
            conversation.append({"role": role, "content": parts})

    current_text, current_paths = normalize_content(message)
    current_parts: list[dict[str, Any]] = [
        {"type": "image", "image": image}
        for image in extract_images(current_paths)
    ]

    if current_text:
        current_parts.append({"type": "text", "text": current_text})
    elif not current_parts:
        current_parts.append({"type": "text", "text": "Tell me about this artwork."})

    conversation.append({"role": "user", "content": current_parts})

    return conversation


# Inference Engine
def prepare_inputs(
    processor: Any,
    conversation: Sequence[dict[str, Any]],
) -> Any:
    """Prepare tokenized model inputs for a multimodal conversation.

    Args:
        processor: The loaded Hugging Face processor.
        conversation: A Qwen2.5-VL formatted conversation.

    Returns:
        Tokenized model inputs moved to the GPU.

    Raises:
        RuntimeError: If processor input preparation fails.
    """
    try:
        prompt = processor.apply_chat_template(
            conversation,
            add_generation_prompt=True,
            tokenize=False,
        )

        images = [
            part["image"]
            for turn in conversation
            for part in turn["content"]
            if part.get("type") == "image"
        ]

        return processor(
            text=[prompt],
            images=images or None,
            return_tensors="pt",
        ).to(DEVICE)

    except Exception as error:
        LOGGER.exception("Processor input preparation failed.")
        raise RuntimeError(
            "I could not prepare that image or message for analysis."
        ) from error


def generate_stream(
    model: Any,
    processor: Any,
    inputs: Any,
) -> Generator[str, None, None]:
    """Generate and stream text from the model.

    Args:
        model: The loaded image-text model.
        processor: The loaded processor.
        inputs: Prepared GPU inputs.

    Yields:
        Incremental generated text.

    Raises:
        RuntimeError: If model generation fails.
    """
    streamer = TextIteratorStreamer(
        processor.tokenizer,
        skip_prompt=True,
        skip_special_tokens=True,
    )

    pad_token_id = processor.tokenizer.pad_token_id
    if pad_token_id is None:
        pad_token_id = processor.tokenizer.eos_token_id

    generation_errors: list[Exception] = []
    generation_kwargs = {
        **inputs,
        "max_new_tokens": MAX_NEW_TOKENS,
        "do_sample": True,
        "temperature": TEMPERATURE,
        "top_p": TOP_P,
        "min_p": MIN_P,
        "pad_token_id": pad_token_id,
        "streamer": streamer,
    }

    def run_generation() -> None:
        """Run blocking generation in a dedicated worker thread."""
        try:
            with torch.inference_mode():
                model.generate(**generation_kwargs)
        except Exception as error:
            generation_errors.append(error)
            LOGGER.exception("Generation thread failed.")
            streamer.end()

    LOGGER.info("Starting streamed generation.")
    worker = threading.Thread(target=run_generation, daemon=True)
    worker.start()

    partial_text = ""
    try:
        for token in streamer:
            partial_text += token
            yield partial_text
    finally:
        worker.join()

    if generation_errors:
        raise RuntimeError("Model generation failed.") from generation_errors[0]

    LOGGER.info("Streamed generation completed.")


# Gradio Callbacks
@spaces.GPU(duration=GPU_DURATION)
def respond(
    message: Any,
    history: list[dict[str, Any]] | None,
) -> Generator[str, None, None]:
    """Process a Gradio chat request and stream Melkov's response.

    Args:
        message: The current multimodal Gradio message.
        history: Previous messages in Gradio messages format.

    Yields:
        Incremental assistant response text.
    """
    try:
        LOGGER.info("Received inference request.")
        processor, model = load_model()
        conversation = build_conversation(history, message)
        inputs = prepare_inputs(processor, conversation)

        yield from generate_stream(model, processor, inputs)

    except Exception as error:
        LOGGER.exception("Inference request failed.")
        yield _friendly_error_message(error)


# UI Builder
def build_ui() -> gr.Blocks:
    """Build and return the Gradio application interface.

    Returns:
        The configured Gradio Blocks application.
    """
    with gr.Blocks() as demo:
        gr.Markdown("# MELKOV - ART ATELIER", elem_id="melkov-title")
        gr.Markdown(
            "Chat with **Melkov**, an art expert vision-language model trained "
            "to discuss paintings, visual composition, art history, and technique. "
            "Upload an artwork image or ask a question about art.",
            elem_id="melkov-subtitle",
        )

        gr.ChatInterface(
            fn=respond,
            multimodal=True,
            textbox=gr.MultimodalTextbox(
                file_types=["image"],
                file_count="multiple",
                sources=["upload"],
                placeholder="Share a painting, or ask Melkov about art...",
            ),
        )

    return demo


def main() -> None:
    """Launch the Gradio Spaces application."""
    LOGGER.info("Building Gradio interface.")
    demo = build_ui()
    demo.queue()
    demo.launch(css=CSS, theme=gr.themes.Base())


if __name__ == "__main__":

    
    main()