import gradio as gr
import requests
import os
import json

# These will be set as Hugging Face Spaces secrets
API_KEY = os.environ.get("FIREWORKS_API_KEY", "")
SYSTEM_PROMPT = os.environ.get("SYSTEM_PROMPT")

# API endpoint
API_URL = "https://api.fireworks.ai/inference/v1/chat/completions"

def chat_with_model(message, history, temperature, max_tokens, top_p, top_k):
    """
    Stream responses from the API
    """
    # Build conversation history
    messages = [{"role": "system", "content": SYSTEM_PROMPT}]
    
    # Add conversation history
    for human, assistant in history:
        messages.append({"role": "user", "content": human})
        messages.append({"role": "assistant", "content": assistant})
    
    # Add current message
    messages.append({"role": "user", "content": message})
    
    # Prepare the request
    headers = {
        "Accept": "text/event-stream",
        "Content-Type": "application/json",
        "Authorization": f"Bearer {API_KEY}"
    }
    
    payload = {
        "model": "accounts/fireworks/models/kimi-k2-instruct-0905",
        "max_tokens": max_tokens,
        "top_p": top_p,
        "top_k": top_k,
        "presence_penalty": 0,
        "frequency_penalty": 0,
        "temperature": temperature,
        "messages": messages,
        "stream": True
    }
    
    try:
        response = requests.post(API_URL, headers=headers, json=payload, timeout=120, stream=True)
        response.raise_for_status()
        
        assistant_message = ""
        
        for line in response.iter_lines():
            if line:
                line = line.decode('utf-8')
                if line.startswith('data: '):
                    data = line[6:]
                    if data == '[DONE]':
                        break
                    try:
                        json_data = json.loads(data)
                        if 'choices' in json_data and len(json_data['choices']) > 0:
                            delta = json_data['choices'][0].get('delta', {})
                            content = delta.get('content', '')
                            if content:
                                assistant_message += content
                                yield assistant_message
                    except json.JSONDecodeError:
                        continue
        
        if not assistant_message:
            yield "❌ No response received from the model."
        
    except requests.exceptions.RequestException as e:
        yield f"❌ Error: {str(e)}\n\nPlease check your API key in Hugging Face Spaces secrets."
    except Exception as e:
        yield f"❌ Unexpected error: {str(e)}"

# Custom CSS for a modern look
custom_css = """
.gradio-container {
    font-family: 'Inter', sans-serif;
    max-width: 1400px !important;
}
#title {
    text-align: center;
    background: linear-gradient(135deg, #667eea 0%, #764ba2 50%, #f093fb 100%);
    -webkit-background-clip: text;
    -webkit-text-fill-color: transparent;
    font-size: 3em;
    font-weight: 800;
    margin-bottom: 0.3em;
    letter-spacing: -0.02em;
}
#subtitle {
    text-align: center;
    font-size: 1.2em;
    color: #888;
    margin-bottom: 2em;
    font-weight: 300;
}
.message-wrap {
    border-radius: 16px !important;
}
footer {
    display: none !important;
}
"""

# Create Gradio interface
with gr.Blocks(css=custom_css, theme=gr.themes.Soft(primary_hue="purple", secondary_hue="pink")) as demo:
    gr.HTML("<h1 id='title'>Palmyra-Sec Playground</h1>")
    gr.HTML("<p id='subtitle'>Intelligent conversations powered by advanced language models</p>")
    
    with gr.Row():
        with gr.Column(scale=3):
            chatbot = gr.Chatbot(
                height=550,
                bubble_full_width=False,
                avatar_images=(
                    "https://api.dicebear.com/7.x/avataaars/svg?seed=user",
                    "https://api.dicebear.com/7.x/bottts-neutral/svg?seed=ai&backgroundColor=b6e3f4"
                ),
                show_copy_button=True
            )
            
            with gr.Row():
                msg = gr.Textbox(
                    placeholder="Ask me anything...",
                    show_label=False,
                    scale=4,
                    container=False,
                    lines=2
                )
                submit_btn = gr.Button("Send 💬", scale=1, variant="primary", size="lg")
            
            with gr.Row():
                clear_btn = gr.ClearButton([msg, chatbot], value="🗑️ Clear", size="sm")
                retry_btn = gr.Button("🔄 Retry", size="sm", variant="secondary")
        
        with gr.Column(scale=1):
            gr.Markdown("### ⚙️ Model Settings")
            
            temperature = gr.Slider(
                minimum=0,
                maximum=2,
                value=0.6,
                step=0.1,
                label="🌡️ Temperature",
                info="Creativity level"
            )
            
            max_tokens = gr.Slider(
                minimum=100,
                maximum=4000,
                value=2000,
                step=100,
                label="📏 Max Tokens",
                info="Response length"
            )
            
            top_p = gr.Slider(
                minimum=0,
                maximum=1,
                value=1,
                step=0.05,
                label="🎯 Top P",
                info="Diversity control"
            )
            
            top_k = gr.Slider(
                minimum=1,
                maximum=100,
                value=40,
                step=1,
                label="🔢 Top K",
                info="Token selection"
            )
            
            with gr.Accordion("📚 Quick Presets", open=False):
                gr.Markdown("""
                **Creative** → Temp: 0.9, Top P: 0.95  
                **Balanced** → Temp: 0.6, Top P: 1.0  
                **Precise** → Temp: 0.3, Top P: 0.9
                """)
                
                with gr.Row():
                    creative_btn = gr.Button("🎨 Creative", size="sm")
                    balanced_btn = gr.Button("⚖️ Balanced", size="sm")
                    precise_btn = gr.Button("🎯 Precise", size="sm")
            
        
    
    # Handle message submission with streaming
    def respond(message, chat_history, temp, max_tok, top_p_val, top_k_val):
        if not message.strip():
            return chat_history, ""
            
        if not API_KEY:
            chat_history.append((message, "⚠️ Please configure API key in Space secrets!"))
            return chat_history, ""
        
        chat_history.append((message, ""))
        
        for partial_response in chat_with_model(message, chat_history[:-1], temp, max_tok, top_p_val, top_k_val):
            chat_history[-1] = (message, partial_response)
            yield chat_history, ""
    
    msg.submit(
        respond,
        [msg, chatbot, temperature, max_tokens, top_p, top_k],
        [chatbot, msg]
    )
    
    submit_btn.click(
        respond,
        [msg, chatbot, temperature, max_tokens, top_p, top_k],
        [chatbot, msg]
    )
    
    # Retry last message
    def retry_last(chat_history, temp, max_tok, top_p_val, top_k_val):
        if not chat_history:
            return chat_history
        
        last_message = chat_history[-1][0]
        chat_history = chat_history[:-1]
        
        for updated_history, _ in respond(last_message, chat_history, temp, max_tok, top_p_val, top_k_val):
            yield updated_history
    
    retry_btn.click(
        retry_last,
        [chatbot, temperature, max_tokens, top_p, top_k],
        [chatbot]
    )
    
    # Preset buttons
    def set_creative():
        return 0.9, 0.95
    
    def set_balanced():
        return 0.6, 1.0
    
    def set_precise():
        return 0.3, 0.9
    
    creative_btn.click(set_creative, None, [temperature, top_p])
    balanced_btn.click(set_balanced, None, [temperature, top_p])
    precise_btn.click(set_precise, None, [temperature, top_p])
    
    # Add examples
    gr.Examples(
        examples=[
            ["Give me defensive actions. Also map the mitigations to NIST and ISO 27001 controls."],
            ["A financial services company detected unusual outbound traffic from an internal SQL server at 2am. Provide an executive summary of what might be happening, a technical analysis, likely MITRE ATT&CK techniques, and immediate incident response steps the SOC should take."],
            ["We are preparing for a SOC-2 audit. Draft a short security policy on how our engineering team should handle secrets (API keys, SSH keys, environment variables) in source code and CI/CD pipelines. Include recommended tools and preventive controls."],
            ["What are the best practices for Python code?"],
        ],
        inputs=msg,
        label="💡 Example Prompts"
    )
    
    gr.Markdown("""
    ---
    <div style='text-align: center; color: #888; font-size: 0.9em;'>
    🔒 <b>Privacy First</b> • All credentials stored securely • Conversations are ephemeral
    </div>
    """)

# Launch the app
if __name__ == "__main__":
    demo.launch()