import gradio as gr
import os
from huggingface_hub import InferenceClient
from utils.retriever import KnowledgeRetriever
import json

class AccessibilityChatbot:
    def __init__(self):
        # Initialize DeepSeek-R1 client
        self.client = InferenceClient(
            model="deepseek-ai/DeepSeek-R1",
            token=os.getenv("HF_TOKEN")
        )
        
        # Initialize knowledge retriever
        self.retriever = KnowledgeRetriever()
        
        # System prompt for accessibility education
        self.system_prompt = """You are an expert web accessibility instructor helping university students learn about web accessibility. 

Your knowledge comes from WebAIM resources, which are authoritative sources for web accessibility information.

Guidelines for responses:
1. Provide clear, student-friendly explanations
2. Use the provided WebAIM context to answer questions accurately
3. Always cite your sources by mentioning the WebAIM document and page number
4. Include practical examples and code snippets when relevant
5. Break down complex concepts into digestible parts
6. Encourage best practices and standards compliance
7. If asked about assignments, provide actionable guidance

Remember: You're teaching students, so be encouraging and educational while maintaining accuracy."""

    def generate_response(self, message, history):
        """Generate response using DeepSeek-R1 with WebAIM context"""
        
        # Retrieve relevant content from WebAIM PDFs
        relevant_content = self.retriever.retrieve_relevant_content(message)
        context = self.retriever.format_context_for_llm(relevant_content)
        
        # Prepare messages for the LLM
        messages = [
            {"role": "system", "content": f"{self.system_prompt}\n\nContext from WebAIM resources:\n{context}"}
        ]
        
        # 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})
        
        try:
            response = self.client.chat_completion(
                messages=messages,
                max_tokens=1500,
                temperature=0.7,
                top_p=0.9
            )
            
            assistant_response = response.choices[0].message.content
            
            # Add source information
            if relevant_content and assistant_response:
                sources = self.format_sources(relevant_content)
                assistant_response += f"\n\n**Sources:**\n{sources}"
            
            return assistant_response or "I apologize, but I couldn't generate a response. Please try again."
            
        except Exception as e:
            return f"I apologize, but I'm experiencing technical difficulties. Please try again. Error: {str(e)}"
    
    def format_sources(self, content_list):
        """Format source citations for display"""
        sources = []
        seen_sources = set()
        
        for item in content_list:
            source_key = f"{item['source_file']}_{item['page_number']}"
            if source_key not in seen_sources:
                sources.append(f"• {item['source_file']} (Page {item['page_number']})")
                seen_sources.add(source_key)
        
        return "\n".join(sources)

# Initialize chatbot
chatbot = AccessibilityChatbot()

# Create Gradio interface
def create_interface():
    # Custom CSS for improved styling
    custom_css = """
    .gradio-container {
        max-width: 800px !important;
        margin: 0 auto !important;
    }
    
    .chat-container {
        background: white;
        border-radius: 15px;
        padding: 2rem;
        box-shadow: 0 8px 32px rgba(0,0,0,0.1);
        border: 1px solid #e1e5e9;
    }
    
    .input-container {
        background: #f8f9fa;
        border-radius: 12px;
        padding: 1.5rem;
        margin-top: 1rem;
    }
    
    .gradio-button {
        background: linear-gradient(135deg, #667eea 0%, #764ba2 100%) !important;
        border: none !important;
        border-radius: 8px !important;
        color: white !important;
        font-weight: 600 !important;
        padding: 12px 24px !important;
        transition: all 0.3s ease !important;
    }
    
    .gradio-button:hover {
        transform: translateY(-2px) !important;
        box-shadow: 0 8px 25px rgba(102, 126, 234, 0.4) !important;
    }
    
    .gradio-textbox {
        border-radius: 12px !important;
        border: 2px solid #e1e5e9 !important;
        transition: border-color 0.3s ease !important;
    }
    
    .gradio-textbox:focus-within {
        border-color: #667eea !important;
        box-shadow: 0 0 0 3px rgba(102, 126, 234, 0.1) !important;
    }
    
    .chatbot-container {
        border-radius: 12px !important;
        border: 1px solid #e1e5e9 !important;
        background: white !important;
    }
    """
    
    with gr.Blocks(
        title="Web Accessibility Learning Assistant",
        css=custom_css
    ) as demo:
        
        # Main chat interface
        with gr.Row():
            with gr.Column(scale=1):
                gr.HTML('<div class="chat-container">')
                
                chatbot_interface = gr.Chatbot(
                    height=600,
                    placeholder="👋 Ask me anything about web accessibility! I'm here to help you learn.",
                    show_label=False,
                    container=True,
                    bubble_full_width=False,
                    elem_classes=["chatbot-container"]
                )
                
                gr.HTML('</div>')
        
        # Input section
        with gr.Row():
            with gr.Column(scale=1):
                gr.HTML('<div class="input-container">')
                
                msg = gr.Textbox(
                    placeholder="Type your question here... (e.g., 'How do I write good alt text?' or 'What are the WCAG contrast requirements?')",
                    label="Your Question",
                    lines=3,
                    max_lines=6,
                    elem_classes=["gradio-textbox"]
                )
                
                with gr.Row():
                    clear_btn = gr.Button("🗑️ Clear Chat", variant="secondary", size="sm")
                    submit_btn = gr.Button("🚀 Ask Question", variant="primary", size="lg")
                
                gr.HTML('</div>')
        
        # Handle message submission
        def respond(message, history):
            if not message.strip():
                return history, ""
            
            response = chatbot.generate_response(message, history)
            history.append((message, response))
            return history, ""
        
        def clear_chat():
            return [], ""
        
        # Event handlers
        msg.submit(respond, [msg, chatbot_interface], [chatbot_interface, msg])
        submit_btn.click(respond, [msg, chatbot_interface], [chatbot_interface, msg])
        clear_btn.click(clear_chat, outputs=[chatbot_interface, msg])
    
    return demo

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