import gradio as gr
from langchain_mistralai import ChatMistralAI
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import StrOutputParser
import os

# Set up API keys from environment variables
mistral_api_key = os.getenv("mistralapikey", "HQ2u8NYTq0Bpc8zpuEQ9rxLd9p7PgHXA")
langchain_api_key = os.getenv("langchainapikey", "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")

# Environment setup
os.environ["mistralapikey"] = mistral_api_key
if langchain_api_key:
    os.environ["langchaintracingv2"] = "true"
    os.environ["langchainapikey"] = langchain_api_key

# Prompt Template
prompt = ChatPromptTemplate.from_messages([
    ("system", "You are a helpful AI assistant powered by Codestral. Please provide clear, accurate, and helpful responses to user queries."),
    ("user", "Question: {question}")
])

# Initialize components
def initialize_chain(temperature=0.7, max_tokens=500):
    """Initialize the Codestral AI chain with given parameters"""
    llm = ChatMistralAI(
        model="codestral-latest",
        mistral_api_key=mistral_api_key,
        temperature=temperature,
        max_tokens=max_tokens
    )
    output_parser = StrOutputParser()
    return prompt | llm | output_parser

def chat_with_codestral(message, history, temperature, max_tokens):
    """Chat function for Gradio interface"""
    try:
        # Initialize chain with current parameters
        chain = initialize_chain(temperature, max_tokens)
        
        # Get response from Codestral
        response = chain.invoke({"question": message})
        
        # Add to history
        history.append([message, response])
        
        return history, ""
    
    except Exception as e:
        error_message = f"❌ Error: {str(e)}"
        history.append([message, error_message])
        return history, ""

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

# Responsive CSS with mobile-first approach
css = """
/* Base styles - Mobile first */
.gradio-container {
    font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif;
    max-width: 100%;
    margin: 0 auto;
    padding: 10px;
}

/* Large screens - center content with max width */
@media (min-width: 1200px) {
    .gradio-container {
        max-width: 1400px;
        padding: 20px;
    }
}

/* Chat message styling */
.chat-message {
    padding: 10px;
    margin: 5px 0;
    border-radius: 10px;
}

.user-message {
    background-color: #e3f2fd;
    margin-left: 10%;
}

.bot-message {
    background-color: #f5f5f5;
    margin-right: 10%;
}

/* Responsive layout adjustments */
@media (max-width: 768px) {
    .user-message {
        margin-left: 5%;
    }
    .bot-message {
        margin-right: 5%;
    }
    
    /* Stack settings and info boxes on mobile */
    .settings-row,
    .info-row {
        flex-direction: column;
    }
}

/* Chatbot container responsive height */
.chatbot-container {
    height: 400px;
    margin-bottom: 10px !important;
}

@media (min-width: 768px) {
    .chatbot-container {
        height: 500px;
        margin-bottom: 15px !important;
    }
}

@media (min-width: 1024px) {
    .chatbot-container {
        height: 600px;
        margin-bottom: 15px !important;
    }
}

/* Reduce gap between chatbot and input */
.input-row {
    display: flex;
    gap: 10px;
    align-items: flex-end;
    flex-wrap: nowrap;
    margin-top: 0 !important;
    padding-top: 0 !important;
}

@media (max-width: 480px) {
    .input-row {
        flex-direction: column;
        gap: 8px;
    }
    
    .input-row .textbox {
        width: 100% !important;
    }
    
    .input-row .button {
        width: 100% !important;
        min-width: unset !important;
    }
}

/* Button responsive sizing */
.send-button {
    min-width: 80px;
    white-space: nowrap;
}

@media (max-width: 480px) {
    .send-button {
        min-width: unset;
        padding: 8px 16px;
    }
}

/* Settings panel responsive behavior */
.settings-panel {
    min-width: 200px;
}

@media (max-width: 1024px) {
    .settings-panel {
        min-width: 180px;
    }
}

@media (max-width: 768px) {
    .settings-panel {
        min-width: 160px;
    }
}

/* Info boxes responsive */
.info-box {
    margin-top: 20px;
    padding: 15px;
    border-radius: 10px;
    font-size: 14px;
}

@media (max-width: 768px) {
    .info-box {
        margin-top: 15px;
        padding: 12px;
        font-size: 13px;
    }
}

@media (max-width: 480px) {
    .info-box {
        margin-top: 10px;
        padding: 10px;
        font-size: 12px;
    }
}

/* Header responsive */
.header {
    text-align: center;
    padding: 20px;
}

@media (max-width: 768px) {
    .header {
        padding: 15px 10px;
    }
    
    .header h1 {
        font-size: 24px;
    }
    
    .header p {
        font-size: 16px;
    }
}

@media (max-width: 480px) {
    .header {
        padding: 10px 5px;
    }
    
    .header h1 {
        font-size: 20px;
    }
    
    .header p {
        font-size: 14px;
    }
}

/* Footer responsive */
.footer {
    text-align: center;
    margin-top: 20px;
    padding: 15px;
    font-size: 14px;
    color: #666;
}

@media (max-width: 768px) {
    .footer {
        font-size: 12px;
        padding: 10px;
    }
}
"""

# Create Gradio interface
with gr.Blocks(css=css, title="🤖 Codestral AI Chat", theme=gr.themes.Soft()) as demo:
    
    # Header
    gr.HTML("""
        <div class="header">
            <h1>🤖 Codestral AI Chat with LangChain</h1>
            <p style="color: #666;">
                Powered by <strong>Codestral AI</strong> from Mistral AI | Built with <strong>LangChain</strong> & <strong>Gradio</strong>
            </p>
        </div>
    """)
    
    # Main chat area - centered layout
    with gr.Column():
        chatbot = gr.Chatbot(
            label="💬 Chat with Codestral AI",
            height=500,
            show_label=True,
            container=True,
            bubble_full_width=False,
            elem_classes=["chatbot-container"],
            show_copy_button=True
        )
        
        # Input row with responsive design - reduced gap
        with gr.Row(elem_classes=["input-row"], variant="compact"):
            msg = gr.Textbox(
                label="Your Message",
                placeholder="Type your question here...",
                lines=2,
                max_lines=4,
                scale=4,
                container=True,
                elem_classes=["textbox"]
            )
            send_btn = gr.Button(
                "Send 🚀", 
                variant="primary", 
                scale=1,
                elem_classes=["button", "send-button"]
            )
        
        # Clear button row
        with gr.Row():
            clear_btn = gr.Button("🗑️ Clear Chat", variant="secondary")
    
    # Model Settings and Info - Centered below chat
    gr.HTML("<h3 style='text-align: center; margin: 30px 0 20px 0;'>⚙️ Model Settings</h3>")
    
    # Settings in a responsive row
    with gr.Row(elem_classes=["settings-row"]):
        with gr.Column(scale=1):
            temperature = gr.Slider(
                minimum=0.0,
                maximum=1.0,
                value=0.7,
                step=0.1,
                label="🌡️ Temperature",
                info="Controls randomness (0=focused, 1=creative)"
            )
        
        with gr.Column(scale=1):
            max_tokens = gr.Slider(
                minimum=50,
                maximum=2000,
                value=500,
                step=50,
                label="📝 Max Tokens",
                info="Maximum response length"
            )
    
    # Tips and Features in responsive columns
    with gr.Row(elem_classes=["info-row"]):
        with gr.Column(scale=1):
            # Tips box
            gr.HTML("""
                <div class="info-box" style="background-color: #f8f9fa;">
                    <h4 style="margin-top: 0;">💡 Tips:</h4>
                    <ul style="margin: 10px 0; padding-left: 20px;">
                        <li>Ask coding questions</li>
                        <li>Request explanations</li>
                        <li>Get creative writing help</li>
                        <li>Solve problems step-by-step</li>
                    </ul>
                </div>
            """)
        
        with gr.Column(scale=1):
            # Features box
            gr.HTML("""
                <div class="info-box" style="background-color: #e8f5e8;">
                    <h4 style="margin-top: 0;">🔧 Features:</h4>
                    <ul style="margin: 10px 0; padding-left: 20px;">
                        <li>✅ Codestral AI Model</li>
                        <li>✅ LangChain Integration</li>
                        <li>✅ Real-time Responses</li>
                        <li>✅ Adjustable Parameters</li>
                        <li>✅ Chat History</li>
                    </ul>
                </div>
            """)

    # Event handlers
    def submit_message(message, history, temp, tokens):
        return chat_with_codestral(message, history, temp, tokens)
    
    # Button and enter key events
    send_btn.click(
        submit_message,
        inputs=[msg, chatbot, temperature, max_tokens],
        outputs=[chatbot, msg]
    )
    
    msg.submit(
        submit_message,
        inputs=[msg, chatbot, temperature, max_tokens],
        outputs=[chatbot, msg]
    )
    
    clear_btn.click(
        clear_chat,
        outputs=[chatbot, msg]
    )
    
    # Footer
    gr.HTML("""
        <div class="footer">
            <p>🚀 Built with <strong>Codestral AI</strong>, <strong>LangChain</strong>, and <strong>Gradio</strong></p>
            <p>Deploy on <a href="https://huggingface.co/spaces" target="_blank">🤗 Hugging Face Spaces</a></p>
        </div>
    """)

# Launch the app
if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=True,
        show_error=True
    )