#!/usr/bin/env python3
"""
Step3-VL-10B Gradio Interface
"""

import os
import sys
import time
import re
import io
import base64
from typing import Optional, List, Generator

from PIL import Image
import gradio as gr
import spaces
from transformers import AutoProcessor, AutoModelForCausalLM

def set_pythonpath(path: str, allow_missing: bool = False):
    """Prepend path to PYTHONPATH and sys.path."""
    if not path:
        return
    if not allow_missing and not os.path.isdir(path):
        return
    if path not in sys.path:
        sys.path.insert(0, path)
    os.environ["PYTHONPATH"] = (
        path
        + os.pathsep
        + os.environ.get("PYTHONPATH", "")
    )

def resolve_remote_snapshot_path(model_name: str, revision: Optional[str]):
    """Resolve HF cache snapshot path for remote model."""
    if not model_name or not revision:
        return None
    hub_cache = os.environ.get("HUGGINGFACE_HUB_CACHE")
    if not hub_cache:
        return None
    model_dir = f"models--{model_name.replace('/', '--')}"
    return os.path.join(hub_cache, model_dir, "snapshots", revision)

def ensure_model_repo_on_pythonpath(model_name: str, revision: Optional[str] = None):
    """Add local/remote model repo path to PYTHONPATH and sys.path."""
    if not model_name:
        return
    if os.path.isdir(model_name):
        set_pythonpath(model_name)
        return
    snapshot_path = resolve_remote_snapshot_path(model_name, revision)
    if snapshot_path:
        # Path may not exist until after first download.
        set_pythonpath(snapshot_path, allow_missing=True)

MODEL_NAME = os.getenv("MODEL_NAME", "stepfun-ai/Step3-VL-10B")
ensure_model_repo_on_pythonpath(MODEL_NAME)

_MODEL = None
_PROCESSOR = None
_LOADED_MODEL_NAME = None



@spaces.GPU
def get_model(model_name: Optional[str] = None):
    """Load and cache the Step3-VL-10B model and processor."""
    global _MODEL, _PROCESSOR, _LOADED_MODEL_NAME

    if model_name is None:
        model_name = MODEL_NAME

    if _MODEL is not None and _PROCESSOR is not None and _LOADED_MODEL_NAME == model_name:
        return _MODEL, _PROCESSOR

    key_mapping = {
        "^vision_model": "model.vision_model",
        r"^model(?!\.(language_model|vision_model))": "model.language_model",
        "vit_large_projector": "model.vit_large_projector",
    }

    print(f"加载模型: {model_name}")
    revision = None
    if os.path.isdir(model_name):
        # Ensure local custom modules (e.g., configuration_step_vl.py) are importable.
        ensure_model_repo_on_pythonpath(model_name, revision)
    _PROCESSOR = AutoProcessor.from_pretrained(
        model_name,
        trust_remote_code=True,
        fix_mistral_regex=True,
        revision=revision
    )
    _MODEL = AutoModelForCausalLM.from_pretrained(
        model_name,
        trust_remote_code=True,
        device_map="auto",
        torch_dtype="auto",
        key_mapping=key_mapping,
        revision=revision
    ).eval()
    _LOADED_MODEL_NAME = model_name
    return _MODEL, _PROCESSOR


@spaces.GPU
def warmup_model(model_name: Optional[str] = None):
    """Warmup the model once at app startup."""
    get_model(model_name)
    return True


def resolve_generation_tokens(processor):
    """Resolve eos/pad token ids and stop strings for safer stopping."""
    tokenizer = getattr(processor, "tokenizer", processor)
    eos_ids = []
    eos_token_id = getattr(tokenizer, "eos_token_id", None)
    if eos_token_id is not None:
        if isinstance(eos_token_id, (list, tuple)):
            eos_ids.extend(list(eos_token_id))
        else:
            eos_ids.append(eos_token_id)

    def _try_add_token(token_str: str):
        try:
            token_id = tokenizer.convert_tokens_to_ids(token_str)
            unk_id = getattr(tokenizer, "unk_token_id", None)
            if token_id is not None and token_id != unk_id:
                eos_ids.append(token_id)
        except Exception:
            return

    # Common end-of-turn markers across chat templates
    for token_str in [
        "<|eot_id|>",
        "<|end_of_turn|>",
        "<|end_of_text|>",
        "<|im_end|>",
        "</s>",
    ]:
        _try_add_token(token_str)

    # De-duplicate while preserving order
    seen = set()
    eos_ids = [i for i in eos_ids if not (i in seen or seen.add(i))]

    pad_token_id = getattr(tokenizer, "pad_token_id", None)
    if pad_token_id is None and eos_ids:
        pad_token_id = eos_ids[0]

    stop_strings = [
        "\nassistant",
        "\nAssistant",
        "<|eot_id|>",
        "<|end_of_turn|>",
        "<|end_of_text|>",
        "<|im_end|>",
        "</s>",
    ]
    return eos_ids, pad_token_id, stop_strings


def truncate_on_stop_strings(text: str, stop_strings):
    if not text:
        return text
    earliest = None
    for s in stop_strings:
        idx = text.find(s)
        if idx != -1 and (earliest is None or idx < earliest):
            earliest = idx
    return text if earliest is None else text[:earliest].rstrip()

def escape_html(text):
    """Escape HTML special characters to prevent XSS"""
    if not isinstance(text, str):
        return text
    return (text
        .replace("&", "&amp;")
        .replace("<", "&lt;")
        .replace(">", "&gt;")
        .replace('"', "&quot;")
        .replace("'", "&#x27;"))

def load_image(image_path: Optional[str]) -> Optional[Image.Image]:
    """Load image from filepath for model input."""
    if not image_path:
        return None
    if not os.path.exists(image_path):
        return None
    try:
        image = Image.open(image_path)
        image.load()  # Ensure file is fully read before temp cleanup.
        if image.mode != "RGB":
            image = image.convert("RGB")
        return image
    except Exception as e:
        print(f"[ERROR] Failed to load image: {e}")
        return None


def image_to_data_url(image: Image.Image) -> Optional[str]:
    """Convert PIL image to data URL for stable chat rendering."""
    if image is None:
        return None
    try:
        buffer = io.BytesIO()
        image.save(buffer, format="PNG")
        encoded = base64.b64encode(buffer.getvalue()).decode("ascii")
        return f"data:image/png;base64,{encoded}"
    except Exception as e:
        print(f"[ERROR] Failed to encode image: {e}")
        return None

def _strip_data_urls(text: str) -> str:
    if not text:
        return text
    # Remove base64 data URLs to avoid massive token counts.
    return re.sub(r"data:image\/[a-zA-Z]+;base64,[A-Za-z0-9+/=\s]+", "", text)


def _strip_html(text: str) -> str:
    if not text:
        return text
    return re.sub(r"<[^>]+>", "", text)


def _extract_text_content(content) -> Optional[str]:
    if content is None:
        return None
    if isinstance(content, str):
        cleaned = _strip_data_urls(content)
        cleaned = _strip_html(cleaned)
        return cleaned.strip() if cleaned.strip() else None
    if isinstance(content, list):
        parts = []
        for item in content:
            if isinstance(item, dict) and item.get("type") == "text":
                parts.append(item.get("text", ""))
        text = "\n".join(p for p in parts if p)
        return text.strip() if text.strip() else None
    return None


def format_messages(history, user_text, image: Optional[Image.Image] = None):
    """Format message list for Step3-VL-10B."""
    messages: List[dict] = []

    if not history:
        history = []

    # 处理历史记录（仅保留文本）
    for item in history:
        role = item.get("role") if isinstance(item, dict) else getattr(item, "role", None)
        content = item.get("content") if isinstance(item, dict) else getattr(item, "content", None)
        raw_text = item.get("raw_text") if isinstance(item, dict) else None

        if not role or content is None:
            continue

        # 如果内容包含思考块，提取最终回答
        if role == "assistant" and isinstance(content, str) and '<div class="thinking-block">' in content:
            pattern = r'<div class="thinking-block">.*?</div>\s*</div>\s*'
            remaining_content = re.sub(pattern, '', content, flags=re.DOTALL).strip()
            if remaining_content and not remaining_content.startswith('<'):
                content = remaining_content
            else:
                continue

        if role == "user" and raw_text is not None:
            text_content = raw_text
        else:
            text_content = _extract_text_content(content)
        if isinstance(text_content, str) and text_content.strip():
            messages.append({
                "role": role,
                "content": [{"type": "text", "text": text_content}]
            })

    # 添加当前用户消息（支持图像 + 文本）
    content_list = []
    if image is not None:
        content_list.append({"type": "image", "image": image})
    if user_text:
        content_list.append({"type": "text", "text": user_text})

    if content_list:
        if len(content_list) == 1 and content_list[0]["type"] == "text":
            messages.append({
                "role": "user",
                "content": [{"type": "text", "text": content_list[0]["text"]}]
            })
        else:
            messages.append({"role": "user", "content": content_list})

    return messages

def build_user_display(image_url: Optional[str], user_text: Optional[str]) -> str:
    parts = []
    if image_url:
        parts.append(
            f'<div class="uploaded-image">'
            f'<img src="{image_url}" alt="uploaded image" />'
            f'</div>'
        )
    if user_text:
        parts.append(user_text)
    return "\n\n".join(parts).strip()


@spaces.GPU
def chat(user_text, image_file, history, max_tokens, temperature, top_p, show_thinking=True, model_name=None):
    """Chat function for Step3-VL-10B."""
    if model_name is None:
        model_name = MODEL_NAME

    if not user_text and not image_file:
        yield history or [], "", None
        return

    # Ensure history is a list and formatted correctly
    history = history or []
    clean_history = []
    for item in history:
        if isinstance(item, dict) and 'role' in item and 'content' in item:
            clean_history.append(item)
        elif hasattr(item, "role") and hasattr(item, "content"):
            clean_history.append(item)
    history = clean_history

    # Load image if provided
    image = load_image(image_file)
    image_url = image_to_data_url(image) if image is not None else None

    messages = format_messages(history, user_text, image)
    if not messages:
        yield history or [], "", None
        return

    # Update history with user message
    user_display = build_user_display(image_url, user_text)
    history.append({
        "role": "user",
        "content": user_display,
        "raw_text": user_text or ""
    })

    # Add thinking placeholder
    if show_thinking:
        history.append({
            "role": "assistant",
            "content": (
                '<div class="thinking-block">\n'
                '<div class="thinking-header">💭 Thinking...</div>\n'
                '<div class="thinking-content">Processing your request...</div>\n'
                '</div>'
            )
        })
    else:
        history.append({
            "role": "assistant",
            "content": "⏳ Generating response..."
        })
    yield history, "", None

    start_time = time.time()
    try:
        # 直接复用已加载的全局模型，必要时再触发加载
        model = _MODEL
        processor = _PROCESSOR
        if model is None or processor is None or _LOADED_MODEL_NAME != model_name:
            model, processor = get_model(model_name)

        inputs = processor.apply_chat_template(
            messages, add_generation_prompt=True, tokenize=True,
            return_dict=True, return_tensors="pt"
        ).to(model.device)

        gen_kwargs = {
            "max_new_tokens": max_tokens,
            "do_sample": temperature > 0
        }
        if temperature > 0:
            gen_kwargs["temperature"] = temperature
            gen_kwargs["top_p"] = top_p

        eos_ids, pad_token_id, stop_strings = resolve_generation_tokens(processor)
        if eos_ids:
            gen_kwargs["eos_token_id"] = eos_ids if len(eos_ids) > 1 else eos_ids[0]
        if pad_token_id is not None:
            gen_kwargs["pad_token_id"] = pad_token_id

        generate_ids = model.generate(**inputs, **gen_kwargs)
        decoded = processor.decode(
            generate_ids[0, inputs["input_ids"].shape[-1]:],
            skip_special_tokens=True
        )
        decoded = truncate_on_stop_strings(decoded, stop_strings).strip()

        if "</think>" in decoded:
            think_content, response_content = decoded.split("</think>", 1)
            if think_content.startswith("<think>"):
                think_content = think_content[len("<think>"):].strip()
            response_content = response_content.strip()
            if show_thinking:
                escaped_think = escape_html(think_content)
                formatted_content = (
                    f'<div class="thinking-block">\n'
                    f'<div class="thinking-header">💭 Thinking Process</div>\n'
                    f'<div class="thinking-content">{escaped_think}</div>\n'
                    f'</div>\n\n'
                    f'{response_content}'
                )
                history[-1]["content"] = formatted_content
            else:
                history[-1]["content"] = response_content
        else:
            history[-1]["content"] = decoded

        elapsed_time = time.time() - start_time
        print(f"[MODEL] ✅ SUCCESS - Time: {elapsed_time:.2f}s")
        yield history, "", None

    except Exception as e:
        elapsed_time = time.time() - start_time
        import traceback
        print(f"[MODEL] ❌ FAILED - Error: {traceback.format_exc()}, Time: {elapsed_time:.2f}s")
        history[-1]["content"] = f"❌ Error: {str(e)}"
        yield history, "", None

# Custom CSS for better UI
custom_css = """
/* 全局样式 */
.gradio-container {
    max-width: 100% !important;
    font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", Roboto, "Helvetica Neue", Arial, sans-serif;
}
/* 标题样式 */
.app-header {
    text-align: center;
    padding: 2.5rem 1.5rem;
    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
    position: relative;
    overflow: hidden;
    border-radius: 16px;
    margin-bottom: 1.5rem;
    box-shadow: 0 8px 24px rgba(102, 126, 234, 0.35);
}
/* 标题背景装饰 */
.app-header::before {
    content: '';
    position: absolute;
    top: -50%;
    right: -50%;
    width: 200%;
    height: 200%;
    background: radial-gradient(circle, rgba(255, 255, 255, 0.1) 0%, transparent 70%);
    animation: rotate 20s linear infinite;
}
@keyframes rotate {
    from { transform: rotate(0deg); }
    to { transform: rotate(360deg); }
}
.app-header h1 {
    margin: 0;
    font-size: 2.8rem;
    font-weight: 700;
    color: white !important;
    text-shadow: 0 3px 6px rgba(0, 0, 0, 0.25);
    letter-spacing: 1px;
    position: relative;
    z-index: 1;
}
.app-header p {
    color: rgba(255, 255, 255, 0.95) !important;
    text-shadow: 0 2px 4px rgba(0, 0, 0, 0.2);
    position: relative;
    z-index: 1;
    line-height: 1.5;
}
/* 聊天框样式 */
.chatbot-container {
    border-radius: 12px;
    box-shadow: 0 2px 8px rgba(0, 0, 0, 0.08);
    overflow: hidden;
}
/* 思考过程样式 - 模仿Claude/ChatGPT的风格 */
.thinking-block {
    background: linear-gradient(135deg, #f5f7fa 0%, #eef2f7 100%);
    border-left: 4px solid #667eea;
    padding: 16px 20px;
    margin: 12px 0;
    border-radius: 8px;
    box-shadow: 0 2px 4px rgba(0, 0, 0, 0.05);
}
.thinking-header {
    display: flex;
    align-items: center;
    font-weight: 600;
    color: #667eea;
    margin-bottom: 10px;
    font-size: 0.95rem;
}
.thinking-content {
    background: #ffffff;
    padding: 12px 16px;
    border-radius: 6px;
    font-family: 'SF Mono', Monaco, 'Cascadia Code', 'Roboto Mono', Consolas, 'Courier New', monospace;
    font-size: 0.9rem;
    line-height: 1.6;
    color: #374151;
    white-space: pre-wrap;
    word-wrap: break-word;
    border: 1px solid #e5e7eb;
}
/* 回复分隔线 */
.response-divider {
    border: none;
    height: 2px;
    background: linear-gradient(to right, transparent, #e5e7eb, transparent);
    margin: 20px 0;
}
/* 按钮样式 */
.primary-btn {
    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%) !important;
    border: none !important;
    transition: all 0.3s ease !important;
}
.primary-btn:hover {
    transform: translateY(-2px);
    box-shadow: 0 4px 12px rgba(102, 126, 234, 0.4) !important;
}
/* 左侧面板样式 */
.left-panel {
    background: #f9fafb;
    border-radius: 12px;
    padding: 1rem;
    height: 100%;
}
/* 输入框样式 */
.input-box textarea {
    border-radius: 8px !important;
    border: 2px solid #e5e7eb !important;
    transition: border-color 0.3s ease !important;
}
.input-box textarea:focus {
    border-color: #667eea !important;
    box-shadow: 0 0 0 3px rgba(102, 126, 234, 0.1) !important;
}
/* 聊天中展示上传图片 */
.uploaded-image img {
    max-width: 100%;
    border-radius: 8px;
    border: 1px solid #e5e7eb;
    display: block;
}
/* 输入区域标题 */
h3 {
    color: #374151;
    font-size: 1.1rem;
    margin: 1rem 0 0.5rem 0;
}
/* 聊天消息样式优化 */
.message-wrap {
    padding: 1rem !important;
}
.message {
    padding: 1rem !important;
    border-radius: 12px !important;
    line-height: 1.6 !important;
}
/* 用户消息 */
.message.user {
    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%) !important;
    color: white !important;
}
/* 助手消息 */
.message.bot {
    background: #f9fafb !important;
    border: 1px solid #e5e7eb !important;
}
/* 左侧面板整体样式 */
.left-column {
    background: linear-gradient(to bottom, #ffffff 0%, #f9fafb 100%);
    border-radius: 12px;
    padding: 1rem;
    box-shadow: 0 2px 8px rgba(0, 0, 0, 0.05);
}
/* 按钮容器样式 */
.button-row {
    margin-top: 1rem;
    gap: 0.5rem;
}
/* Dark Mode Support */
.dark .message.bot {
    background: #1f2937 !important;
    border: 1px solid #374151 !important;
    color: #e5e7eb !important;
}
.dark .thinking-block {
    background: linear-gradient(135deg, #1f2937 0%, #111827 100%);
    border-left: 4px solid #4f46e5;
}
.dark .thinking-content {
    background: #111827;
    color: #e5e7eb;
    border: 1px solid #374151;
}
.dark .thinking-header {
    color: #818cf8;
}
.dark .left-panel {
    background: #111827;
}
.dark .left-column {
    background: linear-gradient(to bottom, #1f2937 0%, #111827 100%);
}
.dark .input-box textarea {
    background-color: #1f2937;
    border-color: #374151 !important;
    color: #e5e7eb;
}
.dark h3 {
    color: #e5e7eb;
}
/* 滚动条美化 */
::-webkit-scrollbar {
    width: 8px;
    height: 8px;
}
::-webkit-scrollbar-track {
    background: #f1f1f1;
    border-radius: 4px;
}
::-webkit-scrollbar-thumb {
    background: #888;
    border-radius: 4px;
}
::-webkit-scrollbar-thumb:hover {
    background: #555;
}
"""

# Gradio Interface
with gr.Blocks(title="Step3-VL-10B", css=custom_css, theme=gr.themes.Soft()) as demo:
    model_ready = gr.State(False)
    # Header
    gr.HTML("""
        <div class="app-header">
            <h1 style="color: white;">🧠 Step3-VL-10B</h1>
            <p style="color: white; margin: 0.8rem 0 0 0; opacity: 0.95; font-size: 1.15rem; font-weight: 500;">
                Lightweight 10B Multimodal Foundation Model
            </p>
        </div>
    """)

    with gr.Row():
        # Left Panel - Input Area
        with gr.Column(scale=1, min_width=350):
            # Configuration
            with gr.Accordion("⚙️ Configuration", open=False):
                max_tokens = gr.Slider(
                    1, 56000, 
                    value=16384, 
                    label="Max Tokens",
                    info="Maximum tokens to generate"
                )
                temperature = gr.Slider(
                    0.0, 2.0, 
                    value=1.0, 
                    label="Temperature",
                    info="Higher = more random"
                )
                top_p = gr.Slider(
                    0.0, 1.0, 
                    value=1.0, 
                    label="Top P",
                    info="Nucleus sampling"
                )
                show_thinking = gr.Checkbox(
                    label="💭 Show Thinking Process", 
                    value=True,
                    info="Display reasoning steps"
                )
            
            # Input Area
            gr.Markdown("### 📝 Your Input")
            user_text = gr.Textbox(
                label="Text Message",
                lines=4,
                placeholder="Type your message here...",
                elem_classes=["input-box"],
                show_label=False
            )
            
            image_file = gr.Image(
                label="🖼️ Image Input",
                type="filepath",
                elem_classes=["image-upload"]
            )

            # Buttons
            with gr.Row():
                clear_btn = gr.Button("🗑️ Clear", scale=1, size="lg")
                submit_btn = gr.Button(
                    "🚀 Send",
                    variant="primary",
                    scale=2,
                    size="lg",
                    elem_classes=["primary-btn"]
                )
            
            # Usage Guide at bottom
            with gr.Accordion("📖 Quick Guide", open=False):
                gr.Markdown("""
                **Usage:**
                - Enter text, upload an image, or do both
                - Supports multi-turn conversation
                - Toggle the thinking process display
                
                **Tips:**
                - Use an image for better visual Q&A
                - Thinking steps show in a blue block
                - Adjust settings in Configuration
                """)

        # Right Panel - Conversation Area
        with gr.Column(scale=2):
            chatbot = gr.Chatbot(
                label="💬 Conversation",
                height=700,
                type="messages",
                elem_classes=["chatbot-container"],
                show_label=True,
                avatar_images=(None, None),
                bubble_full_width=False,
                sanitize_html=False
            )

    submit_btn.click(
        fn=chat,
        inputs=[user_text, image_file, chatbot, max_tokens, temperature, top_p, show_thinking],
        outputs=[chatbot, user_text, image_file]
    )

    clear_btn.click(
        fn=lambda: ([], "", None),
        outputs=[chatbot, user_text, image_file]
    )

    demo.load(
        fn=warmup_model,
        inputs=None,
        outputs=model_ready
    )

if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument("--host", default="0.0.0.0")
    parser.add_argument("--port", type=int, default=7860)
    parser.add_argument("--model", default=MODEL_NAME)
    parser.add_argument("--share", action="store_true", help="启用 Gradio 共享链接")
    args = parser.parse_args()

    # 更新全局模型名称
    if args.model:
        MODEL_NAME = args.model

    print(f"启动Gradio: http://{args.host}:{args.port}")
    print(f"模型: {MODEL_NAME}")
    try:
        demo.launch(server_name=args.host, server_port=args.port, share=args.share)
    except ValueError as e:
        if "localhost is not accessible" in str(e) and not args.share:
            print("检测到 localhost 不可访问，自动启用 share=True 重试。")
            demo.launch(server_name=args.host, server_port=args.port, share=True)
        else:
            raise