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:
                sources = self.format_sources(relevant_content)
                assistant_response += f"\n\n**Sources:**\n{sources}"
            
            return assistant_response
            
        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():
    with gr.Blocks(
        title="Web Accessibility Learning Assistant",
        theme=gr.themes.Soft(),
        css="""
        .gradio-container {
            max-width: 1000px !important;
        }
        """
    ) as demo:
        
        gr.Markdown("""
        # 🌐 Web Accessibility Learning Assistant
        
        Welcome to your personal web accessibility tutor! I'm here to help you learn about web accessibility using **WebAIM resources** - the gold standard for accessibility guidance.
        
        ## What I can help you with:
        - **WCAG Guidelines**: Understanding success criteria and implementation
        - **Screen Reader Testing**: How to test with assistive technologies  
        - **Code Examples**: Accessible HTML, CSS, and JavaScript patterns
        - **Best Practices**: Real-world accessibility solutions
        - **Assignment Help**: Guidance for your accessibility projects
        
        *All answers are based on authoritative WebAIM documentation with proper citations.*
        """)
        
        # Main chat interface
        chatbot_interface = gr.Chatbot(
            height=500,
            placeholder="👋 Ask me anything about web accessibility!",
            show_label=False,
            container=True,
            bubble_full_width=False
        )
        
        msg = gr.Textbox(
            placeholder="Type your question here... (e.g., 'How do I write good alt text?')",
            label="Your Question",
            lines=2,
            max_lines=4
        )
        
        # Quick start examples
        gr.Markdown("### 🚀 Quick Start Examples:")
        gr.Examples(
            examples=[
                "What are the WCAG 2.1 AA requirements for color contrast?",
                "How do I make forms accessible to screen readers?",
                "What's the difference between aria-label and aria-labelledby?",
                "How can I test my website with a screen reader?",
                "What are the most common accessibility mistakes students make?",
                "How do I write effective alt text for complex images?",
                "What ARIA roles should I use for a navigation menu?",
                "How do I make data tables accessible?",
                "What are the keyboard navigation requirements?",
                "How do I ensure my site works without JavaScript?"
            ],
            inputs=msg,
            examples_per_page=5
        )
        
        # Additional resources
        with gr.Accordion("📚 Additional Learning Resources", open=False):
            gr.Markdown("""
            ### Recommended Tools:
            - **WAVE**: Web accessibility evaluation tool
            - **axe DevTools**: Browser extension for accessibility testing
            - **Lighthouse**: Built-in accessibility audit in Chrome
            - **NVDA/JAWS**: Screen readers for testing
            
            ### Key Standards:
            - **WCAG 2.1**: Current web accessibility guidelines
            - **Section 508**: US federal accessibility requirements
            - **ADA**: Americans with Disabilities Act considerations
            """)
        
        # 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, ""
        
        msg.submit(respond, [msg, chatbot_interface], [chatbot_interface, msg])
        
        # Footer
        gr.Markdown("""
        ---
        *This chatbot uses WebAIM resources and is powered by DeepSeek-R1. 
        For the most up-to-date information, always refer to the original WebAIM documentation.*
        """)
    
    return demo

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