import gradio as gr
import requests
import os
import json
from typing import List, Dict
from datetime import datetime

# Groq API Configuration
API_URL = "https://api.groq.com/openai/v1/chat/completions"
API_KEY = os.getenv("GROQ_API_KEY")

# In-memory chat history storage
chat_history: List[Dict[str, str]] = []

def groq_with_memory(message: str, topic: str = "general") -> tuple:
    """Groq API call with chat history"""
    
    if not API_KEY:
        return "❌ No API Key found", ""
    
    if not message.strip():
        return "❌ Empty message", ""
    
    try:
        headers = {
            "Authorization": f"Bearer {API_KEY}",
            "Content-Type": "application/json"
        }
        
        # Add current message to history
        chat_history.append({
            "role": "user",
            "content": message.strip(),
            "topic": topic,
            "timestamp": datetime.now().strftime("%Y-%m-%d %H:%M")
        })
        
        # Prepare messages for API call (last 50 messages for longer context)
        recent_history = chat_history[-50:]
        messages = [{"role": msg["role"], "content": msg["content"]} for msg in recent_history]
        
        payload = {
            "model": "gemma2-9b-it",
            "messages": messages,
            "max_tokens": 3000,  # Increased for longer responses
            "temperature": 0.7
        }
        
        response = requests.post(API_URL, headers=headers, json=payload, timeout=30)
        
        if response.status_code == 200:
            result = response.json()
            if "choices" in result and result["choices"]:
                response_content = result["choices"][0]["message"]["content"]
                chat_history.append({
                    "role": "assistant",
                    "content": response_content,
                    "topic": topic,
                    "timestamp": datetime.now().strftime("%Y-%m-%d %H:%M")
                })
                return response_content, ""  # Clear input after send
            return f"❌ No response: {result}", ""
        return f"❌ HTTP {response.status_code}: {response.text}", ""
            
    except Exception as e:
        return f"❌ Error: {str(e)}", ""

def get_chat_summary(topic_filter: str = None) -> str:
    """Generate conversation summary"""
    if not chat_history:
        return "❌ No chat history to summarize"
    
    filtered = [msg for msg in chat_history if not topic_filter or msg.get("topic") == topic_filter]
    
    if not filtered:
        return f"❌ No messages found for topic: {topic_filter}"
    
    # Group by topic
    topics = {}
    for msg in filtered:
        topic = msg.get("topic", "general")
        if topic not in topics:
            topics[topic] = []
        topics[topic].append(msg)
    
    summary = "📋 **Chat Summary**\n\n"
    for topic, messages in topics.items():
        user_msgs = [m for m in messages if m["role"] == "user"]
        ai_msgs = [m for m in messages if m["role"] == "assistant"]
        
        summary += f"**🏷️ Topic: {topic}**\n"
        summary += f"- Messages: {len(user_msgs)} user, {len(ai_msgs)} AI\n"
        summary += f"- Time span: {messages[0]['timestamp']} - {messages[-1]['timestamp']}\n"
        
        if user_msgs:
            summary += f"- Key topics discussed: {', '.join(msg['content'][:50] + '...' for msg in user_msgs[:3])}\n"
        summary += "\n"
    
    return summary

def get_full_history(topic_filter: str = None) -> str:
    """Get full chat history with optional topic filter"""
    if not chat_history:
        return "❌ No chat history available"
    
    filtered = [msg for msg in chat_history if not topic_filter or msg.get("topic") == topic_filter]
    
    if not filtered:
        return f"❌ No messages found for topic: {topic_filter}"
    
    history_text = f"📚 **Chat History** ({len(filtered)} messages)\n\n"
    
    current_topic = None
    for msg in filtered:
        # Add topic header when topic changes
        if msg.get("topic") != current_topic:
            current_topic = msg.get("topic")
            history_text += f"\n**🏷️ Topic: {current_topic}**\n"
            history_text += "---\n"
        
        role_icon = "👤" if msg["role"] == "user" else "🤖"
        history_text += f"{role_icon} **{msg['timestamp']}**\n"
        history_text += f"{msg['content']}\n\n"
    
    return history_text

def clear_all_history():
    """Clear all chat history"""
    global chat_history
    chat_history.clear()
    return "✅ All chat history cleared", "", ""

def get_topics_list() -> List[str]:
    """Get list of unique topics"""
    topics = list(set(msg.get("topic", "general") for msg in chat_history))
    return ["All Topics"] + sorted(topics)

# Custom CSS for better UI
custom_css = """
.gradio-container {
    max-width: 100% !important;
    padding: 0 !important;
    margin: 0 !important;
}

.main {
    max-width: 100% !important;
    padding: 10px !important;
}

/* Better font for readability */
* {
    font-family: 'Inter', 'Segoe UI', system-ui, -apple-system, sans-serif !important;
}

/* AI Response text spacing */
.response-area textarea {
    line-height: 1.7 !important;
    padding: 20px !important;
    font-size: 15px !important;
    font-family: 'Inter', 'Segoe UI', system-ui, sans-serif !important;
}

/* History text spacing */
.history-display textarea {
    line-height: 1.7 !important;
    padding: 20px !important;
    font-size: 15px !important;
    font-family: 'Inter', 'Segoe UI', system-ui, sans-serif !important;
}

/* Input text styling */
.input-area textarea {
    font-family: 'Inter', 'Segoe UI', system-ui, sans-serif !important;
    font-size: 15px !important;
    line-height: 1.6 !important;
}

/* Better text readability everywhere */
textarea, input {
    line-height: 1.6 !important;
    font-family: 'Inter', 'Segoe UI', system-ui, sans-serif !important;
}
"""

# Main Gradio Interface
with gr.Blocks(
    title="🤖 AI Journal Chat", 
    theme=gr.themes.Soft(),
    css=custom_css
) as demo:
    
    gr.Markdown("# 📝 AI Journal Chat Interface")
    gr.Markdown("*Write, chat, and keep track of your thoughts with AI assistance*")
    
    with gr.Tabs() as tabs:
        # MAIN CHAT TAB
        with gr.Tab("💬 Chat"):
            # AI Response Area (Top)
            ai_response = gr.Textbox(
                label="🤖 AI Response",
                lines=12,
                max_lines=20,
                interactive=False,
                placeholder="AI responses will appear here...",
                show_copy_button=True,
                elem_classes="response-area"
            )
            
            # Input Area (Bottom)
            with gr.Group():
                with gr.Row():
                    user_input = gr.Textbox(
                        label="✍️ Your Message",
                        placeholder="Type your thoughts, questions, or journal entry here...",
                        lines=4,
                        max_lines=10,
                        scale=3
                    )
                    with gr.Column(scale=1):
                        topic_input = gr.Textbox(
                            label="🏷️ Topic",
                            value="journal",
                            placeholder="e.g., work, personal, ideas"
                        )
                        send_btn = gr.Button("📤 Send", variant="primary", size="lg")
                
                with gr.Row():
                    clear_response_btn = gr.Button("🗑️ Clear Response", variant="secondary")
                    show_context_btn = gr.Button("📋 Show Current Context", variant="secondary")
        
        # HISTORY TAB
        with gr.Tab("📚 Chat History"):
            with gr.Group():
                with gr.Row():
                    topic_filter = gr.Dropdown(
                        label="🔍 Filter by Topic",
                        choices=["All Topics"],
                        value="All Topics",
                        interactive=True
                    )
                    refresh_topics_btn = gr.Button("🔄 Refresh Topics", variant="secondary")
                
                with gr.Row():
                    show_history_btn = gr.Button("📖 Show Full History", variant="primary")
                    show_summary_btn = gr.Button("📋 Show Summary", variant="secondary")
                    clear_history_btn = gr.Button("🗑️ Clear All History", variant="stop")
                
                history_display = gr.Textbox(
                    label="📚 History & Summary",
                    lines=20,
                    max_lines=30,
                    interactive=False,
                    show_copy_button=True,
                    placeholder="Chat history and summaries will appear here...",
                    elem_classes="history-display"
                )
    
    # Event Handlers
    def send_message(message, topic):
        response, cleared_input = groq_with_memory(message, topic)
        return response, cleared_input
    
    def show_current_context():
        """Show current conversation context that AI can see"""
        if not chat_history:
            return "❌ No conversation context yet"
        
        recent_history = chat_history[-50:]  # Same as what AI sees
        context_text = f"🧠 **Current AI Context** ({len(recent_history)} messages)\n\n"
        
        for msg in recent_history:
            role_icon = "👤" if msg["role"] == "user" else "🤖"
            context_text += f"{role_icon} **{msg['timestamp']}** [{msg.get('topic', 'general')}]\n"
            context_text += f"{msg['content'][:100]}{'...' if len(msg['content']) > 100 else ''}\n\n"
        
        context_text += f"\n💡 *AI can remember these {len(recent_history)} messages in current conversation*"
        return context_text
    
    def refresh_topic_choices():
        return gr.Dropdown(choices=get_topics_list())
    
    def clear_only_response():
        return ""
    
    def filter_and_show_history(topic_filter):
        filter_topic = None if topic_filter == "All Topics" else topic_filter
        return get_full_history(filter_topic)
    
    def filter_and_show_summary(topic_filter):
        filter_topic = None if topic_filter == "All Topics" else topic_filter
        return get_chat_summary(filter_topic)
    
    # Button Events
    send_btn.click(
        send_message,
        inputs=[user_input, topic_input],
        outputs=[ai_response, user_input]
    )
    
    user_input.submit(
        send_message,
        inputs=[user_input, topic_input],
        outputs=[ai_response, user_input]
    )
    
    clear_response_btn.click(
        clear_only_response,
        outputs=[ai_response]
    )
    
    show_context_btn.click(
        show_current_context,
        outputs=[ai_response]
    )
    
    refresh_topics_btn.click(
        refresh_topic_choices,
        outputs=[topic_filter]
    )
    
    show_history_btn.click(
        filter_and_show_history,
        inputs=[topic_filter],
        outputs=[history_display]
    )
    
    show_summary_btn.click(
        filter_and_show_summary,
        inputs=[topic_filter],
        outputs=[history_display]
    )
    
    clear_history_btn.click(
        clear_all_history,
        outputs=[ai_response, user_input, history_display]
    )

# Launch configuration
print("🚀 Starting AI Journal Chat Interface...")
print(f"🌐 Access at: http://localhost:7860")
print(f"🔑 API Key: {'✅ Found' if API_KEY else '❌ Missing'}")

if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        show_error=True,
        share=False
    )