import gradio as gr
import requests
import os
import time

# Get the token from environment variables
API_TOKEN = os.environ.get("HF_TOKEN", "")
MODEL_NAME = "deepseek-ai/Janus-Pro-7B"
API_URL = f"https://api-inference.huggingface.co/models/{MODEL_NAME}"
headers = {"Authorization": f"Bearer {API_TOKEN}"} if API_TOKEN else {}

def check_token_setup():
    """Check if token is properly set up"""
    if not API_TOKEN:
        return "❌ HF_TOKEN not found. Please add your Hugging Face token in Space Settings → Secrets → HF_TOKEN"
    elif not API_TOKEN.startswith("hf_"):
        return "❌ Invalid token format. Token should start with 'hf_'"
    else:
        return "✅ Token is properly configured!"

def query_janus_model(payload):
    """Send request to Hugging Face Inference API"""
    if not API_TOKEN:
        return {"error": "No API token configured"}
    
    try:
        response = requests.post(API_URL, headers=headers, json=payload, timeout=60)
        
        if response.status_code == 503:
            return {"error": "🔄 Model is loading, please try again in 30-60 seconds..."}
        elif response.status_code == 401:
            return {"error": "🔐 Invalid API token. Please check your HF_TOKEN secret."}
        elif response.status_code == 404:
            return {"error": "❌ Model not found. The model might be temporarily unavailable."}
        elif response.status_code != 200:
            return {"error": f"API Error {response.status_code}: {response.text}"}
            
        return response.json()
    except requests.exceptions.Timeout:
        return {"error": "⏰ Request timeout - model might be loading"}
    except Exception as e:
        return {"error": f"🔌 Connection error: {str(e)}"}

def chat_with_janus(message, history):
    """Chat function for Janus-Pro model"""
    # First check token setup
    token_status = check_token_setup()
    if "❌" in token_status:
        yield token_status
        return
    
    # Show token status
    yield "✅ Token configured! Processing your request..."
    time.sleep(1)
    
    # Prepare the payload
    payload = {
        "inputs": message,
        "parameters": {
            "max_new_tokens": 350,
            "temperature": 0.7,
            "top_p": 0.9,
            "do_sample": True,
            "return_full_text": False
        },
        "options": {
            "wait_for_model": True
        }
    }
    
    # Show loading message
    yield "🔄 Sending request to Janus-Pro-7B... (First time may take 30-60 seconds)"
    
    # Query the model
    result = query_janus_model(payload)
    
    # Process the response
    if "error" in result:
        yield f"{result['error']}"
    elif isinstance(result, list) and len(result) > 0:
        if 'generated_text' in result[0]:
            generated_text = result[0]['generated_text']
            # Clean up the response
            if generated_text.startswith(message):
                generated_text = generated_text[len(message):].strip()
            yield generated_text
        else:
            yield f"📄 Response: {str(result[0])}"
    elif isinstance(result, dict) and 'generated_text' in result:
        yield result['generated_text']
    else:
        yield f"📄 Response: {str(result)}"

def clear_chat():
    """Clear chat history"""
    return [], []

# Create the chat interface
with gr.Blocks(theme=gr.themes.Soft(), title="Janus-Pro-7B API") as demo:
    gr.Markdown(
        """
        # 🚀 Janus-Pro-7B Chat API
        **Free API for DeepSeek's Multimodal AI Model**
        
        *Understanding & Generation Capabilities*
        """
    )
    
    # Token status display
    token_status = gr.Markdown(check_token_setup())
    
    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("### 📋 Instructions")
            gr.Markdown("""
            1. **First message may take 30-60 seconds** (model loading)
            2. **Subsequent messages are faster**
            3. **Be specific** for better responses
            4. Model supports **technical explanations** and **creative writing**
            """)
            
            gr.Markdown("### ℹ️ Model Info")
            gr.Markdown("""
            - **Name**: Janus-Pro-7B
            - **Developer**: DeepSeek AI
            - **Type**: Multimodal (Text + Image)
            - **License**: MIT
            """)
            
        with gr.Column(scale=2):
            chatbot = gr.Chatbot(
                label="💬 Chat with Janus-Pro",
                height=400,
                show_copy_button=True,
                placeholder="Your conversation will appear here..."
            )
            
            with gr.Row():
                msg = gr.Textbox(
                    label="Type your message",
                    placeholder="Ask me anything... (Press Enter to send)",
                    scale=4,
                    container=False,
                    autofocus=True
                )
                clear_btn = gr.Button("🗑️ Clear Chat", scale=1)
            
            with gr.Row():
                gr.Examples(
                    examples=[
                        "Explain quantum computing like I'm 10 years old",
                        "Write a short story about a robot learning to paint",
                        "What are the main differences between AI and human intelligence?",
                        "How do transformers work in machine learning?"
                    ],
                    inputs=msg,
                    label="💡 Try these examples:"
                )
    
    # Event handlers
    def respond_and_clear(message, history):
        for response in chat_with_janus(message, history):
            yield response
    
    msg.submit(
        fn=respond_and_clear,
        inputs=[msg, chatbot],
        outputs=chatbot
    ).then(
        lambda: "",  # Clear input
        outputs=msg
    )
    
    clear_btn.click(
        fn=clear_chat,
        outputs=[chatbot, msg]
    )
    
    gr.Markdown(
        """
        ---
        **Note**: This is a free API using Hugging Face's Inference API. There might be rate limits for high usage.
        
        **Troubleshooting**:
        - If you get token errors, check your HF_TOKEN secret in Space Settings
        - If model is loading, wait 30-60 seconds and try again
        - For persistent issues, check the Space logs
        """
    )

if __name__ == "__main__":
    demo.launch(debug=True)