import gradio as gr
import os
import requests
import json

def create_deepseek_interface():
    # Initialize chat history
    chat_history = []
    
    # Function to call the DeepSeek API
    def query_deepseek(message, history, api_key):
        if not api_key:
            return history, "Please provide your Fireworks AI API key"
        
        # Prepare the conversation history for the API
        messages = []
        for h in history:
            messages.append({"role": "user", "content": h[0]})
            messages.append({"role": "assistant", "content": h[1]})
        
        # Add the new user message
        messages.append({"role": "user", "content": message})
        
        # Prepare the API request
        url = "https://api.fireworks.ai/inference/v1/chat/completions"
        payload = {
            "model": "accounts/fireworks/models/deepseek-v3",
            "max_tokens": 16384,
            "top_p": 1,
            "top_k": 40,
            "presence_penalty": 0,
            "frequency_penalty": 0,
            "temperature": 0.6,
            "messages": messages
        }
        headers = {
            "Accept": "application/json",
            "Content-Type": "application/json",
            "Authorization": f"Bearer {api_key}"
        }
        
        try:
            # Make the API request
            response = requests.post(url, headers=headers, data=json.dumps(payload))
            response.raise_for_status()  # Raise exception for HTTP errors
            
            # Extract the response
            result = response.json()
            assistant_response = result.get("choices", [{}])[0].get("message", {}).get("content", "")
            
            # Update history with the new exchange
            history.append((message, assistant_response))
            return history, ""
        except requests.exceptions.RequestException as e:
            error_msg = f"API Error: {str(e)}"
            if response.status_code == 401:
                error_msg = "Authentication failed. Please check your API key."
            return history, error_msg
    
    # Create Gradio interface
    with gr.Blocks(theme="soft", fill_height=True) as demo:
        # Header Section
        gr.Markdown(
            """
            # 🤖 DeepSeek V3 Inference Interface
            ### Advanced AI Model Powered by Fireworks AI
            """
        )
        
        # State for error messages
        error_msg = gr.State("")
        
        # Main layout with two columns
        with gr.Row():
            # Sidebar with Model Information and API Key
            with gr.Column(scale=1):
                gr.Markdown(
                    """
                    ## 🔑 Access Control
                    ### Inference Provider
                    This interface connects to the DeepSeek-V3 model, served by the Fireworks AI API.
                    
                    #### Authentication
                    - Enter your Fireworks AI API key below
                    - Secure API access with end-to-end encryption
                    """
                )
                
                # API Key input
                api_key = gr.Textbox(
                    label="Fireworks AI API Key",
                    placeholder="Enter your API key...",
                    type="password"
                )
                
                # Model Details Section
                gr.Markdown(
                    """
                    ### 📊 Model Details
                    - **Model**: DeepSeek-V3
                    - **Provider**: Fireworks AI
                    - **Max Tokens**: 16,384
                    - **Temperature**: 0.6
                    - **Capabilities**: Advanced Language Understanding
                    """
                )
            
            # Main Content Area
            with gr.Column(scale=2):
                # Chat interface
                chatbot = gr.Chatbot(
                    height=500,
                    show_label=False,
                    container=True,
                    bubble=True,
                    avatar_images=("👤", "🤖")
                )
                
                # Error message display
                error_display = gr.Markdown(visible=False)
                
                # Input area
                with gr.Row():
                    msg = gr.Textbox(
                        label="Your Message",
                        placeholder="Type your prompt here...",
                        show_label=False,
                        scale=9
                    )
                    submit = gr.Button("Send", variant="primary", scale=1)
                
                # Buttons for clearing and examples
                with gr.Row():
                    clear = gr.ClearButton([msg, chatbot], value="🧹 Clear Conversation")
                    
                # Example queries
                gr.Examples(
                    examples=[
                        "Explain the differences between transformers and RNNs in deep learning.",
                        "Write a Python function to find prime numbers in a range.",
                        "Summarize the key concepts of reinforcement learning."
                    ],
                    inputs=msg
                )
        
        # Handle form submission
        def process_submission(message, history, api_key):
            if not message.strip():
                return history, error_display.update(visible=False)
            
            updated_history, error = query_deepseek(message, history.copy(), api_key)
            
            if error:
                return history, error_display.update(value=f"**Error:** {error}", visible=True)
            else:
                return updated_history, error_display.update(visible=False)
        
        # Connect the button to the function
        submit.click(
            process_submission,
            inputs=[msg, chatbot, api_key],
            outputs=[chatbot, error_display],
            postprocess=lambda _: ("", )  # Clear input box after sending
        )
        
        # Also allow pressing Enter to submit
        msg.submit(
            process_submission,
            inputs=[msg, chatbot, api_key],
            outputs=[chatbot, error_display],
            postprocess=lambda _: ("", )  # Clear input box after sending
        )
        
    return demo

# Launch the interface
if __name__ == "__main__":
    demo = create_deepseek_interface()
    demo.launch(debug=True)