# app.py
import gradio as gr
from openai import OpenAI
import os
import time

print(f"Gradio version: {gr.__version__}")

# ==================== Configuration ====================
DEFAULT_SYSTEM_PROMPT = "You are DeepSeek-V4, an advanced AI assistant with strong reasoning capabilities. Provide accurate, helpful, and well-reasoned responses."

REASONING_EFFORT_MAP = {
    "Non-think": "minimal",
    "Think High": "high", 
    "Think Max": "maximum"
}

THINKING_TYPE_MAP = {
    "Non-think": "disabled",
    "Think High": "enabled",
    "Think Max": "enabled"
}

# ==================== API Client ====================
def get_client():
    api_key = os.environ.get('DEEPSEEK_API_KEY')
    if not api_key:
        raise ValueError(
            "⚠️ DEEPSEEK_API_KEY not found!\n\n"
            "Get your key: https://platform.deepseek.com/api_keys\n"
            "Then set: DEEPSEEK_API_KEY=your-key-here"
        )
    return OpenAI(api_key=api_key, base_url="https://api.deepseek.com")

# ==================== Response Generation ====================
def generate_response(
    message: str,
    history: list,
    thinking_mode: str,
    max_tokens: int,
    temperature: float,
    top_p: float,
    system_prompt: str,
    show_thinking: bool
):
    if not message.strip():
        yield history, "", "Please enter a message."
        return
    
    client = get_client()
    
    # Build messages from history
    api_messages = [{"role": "system", "content": system_prompt}]
    
    if history:
        for item in history:
            if isinstance(item, dict):
                api_messages.append({"role": item["role"], "content": item["content"]})
            elif isinstance(item, (list, tuple)):
                # Handle [(user, assistant), ...] format
                api_messages.append({"role": "user", "content": item[0]})
                if item[1]:
                    api_messages.append({"role": "assistant", "content": item[1]})
    
    api_messages.append({"role": "user", "content": message})
    
    reasoning_effort = REASONING_EFFORT_MAP.get(thinking_mode, "high")
    thinking_type = THINKING_TYPE_MAP.get(thinking_mode, "enabled")
    
    try:
        start_time = time.time()
        
        stream = client.chat.completions.create(
            model="deepseek-v4-pro",
            messages=api_messages,
            stream=True,
            max_tokens=max_tokens,
            temperature=temperature,
            top_p=top_p,
            reasoning_effort=reasoning_effort,
            extra_body={"thinking": {"type": thinking_type}}
        )
        
        content_chunks = []
        thinking_chunks = []
        
        for chunk in stream:
            delta = chunk.choices[0].delta
            
            if delta.content:
                content_chunks.append(delta.content)
            
            if hasattr(delta, 'reasoning_content') and delta.reasoning_content:
                thinking_chunks.append(delta.reasoning_content)
            
            current_content = ''.join(content_chunks)
            current_thinking = ''.join(thinking_chunks)
            
            # Build response
            if show_thinking and current_thinking:
                response_text = f"**[Thinking]**\n{current_thinking}\n\n**[Response]**\n{current_content}"
            else:
                response_text = current_content
            
            elapsed = time.time() - start_time
            
            # Try dict format, fallback to tuple format
            try:
                new_history = history.copy() if history else []
                new_history.append({"role": "user", "content": message})
                new_history.append({"role": "assistant", "content": response_text})
            except:
                new_history = list(history) if history else []
                new_history.append((message, response_text))
            
            yield new_history, current_thinking, f"🔄 {elapsed:.1f}s"
        
        # Final
        final_content = ''.join(content_chunks)
        final_thinking = ''.join(thinking_chunks)
        
        if show_thinking and final_thinking:
            final_response = f"**[Thinking]**\n{final_thinking}\n\n**[Response]**\n{final_content}"
        else:
            final_response = final_content
        
        elapsed = time.time() - start_time
        
        try:
            final_history = history.copy() if history else []
            final_history.append({"role": "user", "content": message})
            final_history.append({"role": "assistant", "content": final_response})
        except:
            final_history = list(history) if history else []
            final_history.append((message, final_response))
        
        yield final_history, final_thinking, f"✅ {elapsed:.1f}s"
        
    except Exception as e:
        error_msg = f"❌ {str(e)}"
        try:
            new_history = history.copy() if history else []
            new_history.append({"role": "user", "content": message})
            new_history.append({"role": "assistant", "content": error_msg})
        except:
            new_history = list(history) if history else []
            new_history.append((message, error_msg))
        yield new_history, "", error_msg

# ==================== Gradio Interface ====================
def create_demo():
    # Try with type parameter, fallback without
    try:
        # Try newest format
        chatbot = gr.Chatbot(
            label="💬 Chat with DeepSeek-V4 Pro",
            height=500,
            type="messages"
        )
    except:
        try:
            # Try without type
            chatbot = gr.Chatbot(
                label="💬 Chat with DeepSeek-V4 Pro",
                height=500
            )
        except:
            # Minimal
            chatbot = gr.Chatbot(label="💬 Chat with DeepSeek-V4 Pro")
    
    with gr.Blocks(title="DeepSeek-V4 Pro Demo") as demo:
        
        # Header
        gr.Markdown("""
        # 🚀 DeepSeek-V4 Pro
        
        **Million-Token Context Intelligence**
        1.6T Parameters | 49B Activated | 1M Context
        """)
        
        with gr.Row():
            # Left sidebar
            with gr.Column(scale=1, min_width=300):
                gr.Markdown("### ⚙️ Settings")
                
                thinking_mode = gr.Radio(
                    choices=["Non-think", "Think High", "Think Max"],
                    value="Think High",
                    label="🧠 Reasoning Mode"
                )
                
                show_thinking = gr.Checkbox(
                    value=True,
                    label="📝 Show thinking process"
                )
                
                system_prompt = gr.Textbox(
                    label="📋 System Prompt",
                    value=DEFAULT_SYSTEM_PROMPT,
                    lines=3
                )
                
                with gr.Accordion("🔧 Advanced", open=False):
                    max_tokens = gr.Slider(64, 32768, value=4096, step=64, label="Max Tokens")
                    temperature = gr.Slider(0.0, 2.0, value=0.7, step=0.05, label="Temperature")
                    top_p = gr.Slider(0.0, 1.0, value=1.0, step=0.05, label="Top P")
            
            # Right - Chat
            with gr.Column(scale=2):
                chatbot_component = gr.Chatbot(
                    label="💬 Chat with DeepSeek-V4 Pro",
                    height=500
                )
                
                with gr.Accordion("🧠 Thinking Process", open=True):
                    thinking_display = gr.Textbox(
                        label="Reasoning",
                        value="*Waiting...*",
                        lines=5,
                        interactive=False
                    )
                
                with gr.Row():
                    message_input = gr.Textbox(
                        label="Message",
                        placeholder="Type your message...",
                        lines=2,
                        scale=9
                    )
                    send_btn = gr.Button("🚀 Send", variant="primary", scale=1)
                
                with gr.Row():
                    clear_btn = gr.Button("🗑️ Clear", size="sm")
                
                status_display = gr.Textbox(
                    label="Status",
                    value="✅ Ready",
                    interactive=False
                )
        
        # Footer
        gr.Markdown("""
        ---
        🔑 [Get API Key](https://platform.deepseek.com/api_keys) | 
        📦 [Model](https://huggingface.co/deepseek-ai/DeepSeek-V4-Pro)
        
        ### 🔥 Trending AI Tools — Free to Try
        - 🖼️ [Free AI Image Upscaler 4K/8K — Enhance Any Photo Instantly](https://dlss5nvidia.com)
        - 🏠 [Free AI House Design Generator — Floor Plans & 3D Render in Seconds](https://housepluspus.com)
        - 🎨 [Free GPT-4o Image Generator — Text to Stunning Visuals](https://gpst-image2.com)
        - ⚡ [DLSS 5 AI Rendering Boost — 2x FPS for Free](https://dlss5.app)
        """)
        
        # ==================== Events ====================
        
        def on_send(msg, hist, mode, show, sys_prompt, max_tok, temp, top_p_val):
            if not msg.strip():
                yield hist, "*Waiting...*", "Enter a message."
                return
            
            if not os.environ.get('DEEPSEEK_API_KEY'):
                new_hist = list(hist) if hist else []
                new_hist.append((msg, "⚠️ Set DEEPSEEK_API_KEY"))
                yield new_hist, "*Error*", "❌ API Key missing"
                return
            
            hist = list(hist) if hist else []
            
            for h, t, s in generate_response(msg, hist, mode, max_tok, temp, top_p_val, sys_prompt, show):
                yield h, t if t else "*No reasoning*", s
        
        send_btn.click(
            fn=on_send,
            inputs=[message_input, chatbot_component, thinking_mode, show_thinking,
                    system_prompt, max_tokens, temperature, top_p],
            outputs=[chatbot_component, thinking_display, status_display]
        ).then(
            fn=lambda: "",
            outputs=[message_input]
        )
        
        message_input.submit(
            fn=on_send,
            inputs=[message_input, chatbot_component, thinking_mode, show_thinking,
                    system_prompt, max_tokens, temperature, top_p],
            outputs=[chatbot_component, thinking_display, status_display]
        ).then(
            fn=lambda: "",
            outputs=[message_input]
        )
        
        clear_btn.click(
            fn=lambda: ([], "*Cleared*", "✅ Cleared"),
            outputs=[chatbot_component, thinking_display, status_display]
        )
    
    return demo

# ==================== Main ====================
if __name__ == "__main__":
    try:
        from dotenv import load_dotenv
        load_dotenv()
    except:
        pass
    
    if not os.environ.get('DEEPSEEK_API_KEY'):
        print("\n⚠️  Set DEEPSEEK_API_KEY environment variable!")
        print("   export DEEPSEEK_API_KEY='your-key'\n")
    
    demo = create_demo()
    demo.queue(max_size=50).launch(
        server_name="0.0.0.0",
        server_port=7860
    )