app.py

import os
from typing import List, Dict, Any, Optional, Tuple
import chromadb
from chromadb.utils import embedding_functions
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain_community.chat_message_histories import ChatMessageHistory
from langchain_core.messages import HumanMessage, AIMessage, SystemMessage
from langchain_community.utilities import GoogleSerperAPIWrapper
import gradio as gr
import re

class NigerianLegalAidChatbot:
    """
    A comprehensive legal aid chatbot for Nigerian law using RAG + Web Search fallback.
    """

    def __init__(
        self,
        google_api_key: str,
        serper_api_key: str,
        chroma_persist_directory: str = "./legal_chroma_db",
        model_name: str = "gemini-2.5-flash"
    ):
        """
        Initialize the Nigerian Legal Aid Chatbot.
        """
        os.environ["GOOGLE_API_KEY"] = google_api_key
        os.environ["SERPER_API_KEY"] = serper_api_key

        self.chroma_client = chromadb.PersistentClient(path=chroma_persist_directory)
        self.embedding_function = embedding_functions.SentenceTransformerEmbeddingFunction(
            model_name="all-MiniLM-L6-v2"
        )

        try:
            self.collection = self.chroma_client.get_collection(
                name="nigerian_legal_docs",
                embedding_function=self.embedding_function
            )
            print(f"✓ Connected to ChromaDB with {self.collection.count()} chunks")
        except Exception as e:
            raise Exception(f"Failed to load ChromaDB collection: {e}")

        self.llm = ChatGoogleGenerativeAI(
            model=model_name,
            temperature=0.2,
            convert_system_message_to_human=True
        )

        self.chat_history = ChatMessageHistory()
        self.search = GoogleSerperAPIWrapper()

        self.legal_system_prompt = """You are a Nigerian Legal Aid Assistant, an expert AI lawyer specializing in Nigerian law, particularly:
- Federal Competition and Consumer Protection Act (FCCPC) 2018
- Lagos State Tenancy Law 2011
- Nigerian Labour Act

YOUR CORE RESPONSIBILITIES:
1. Provide accurate legal guidance based on Nigerian laws
2. Always cite specific sections and acts when giving legal advice
3. Be precise with legal terminology and explain complex terms
4. Remain objective and professional in all responses
5. Acknowledge limitations - you provide information, not legal representation

RESPONSE FORMAT RULES:
✓ When using context from legal documents:
  - Quote the specific Act and Section
  - Include relevant excerpts in quotes
  
✓ When context is insufficient:
  - Use web search to find additional Nigerian legal resources
  - Clearly distinguish between document-based answers and web-sourced information
  
✓ Structure your answers:
  1. Direct answer to the question
  2. Legal basis (Act, Section, specific provision)
  3. Practical implications or examples
  4. Any important caveats or exceptions
  5. Recommendation to consult a lawyer if needed

Remember: You are providing legal information and guidance, NOT acting as a licensed attorney."""

        self.casual_system_prompt = """You are a friendly and helpful AI assistant with knowledge about Nigerian legal matters, but currently in casual conversation mode.

When users ask non-legal questions:
- Be warm, helpful, and conversational
- Provide accurate and helpful information
- Stay professional but friendly

If users want legal advice, politely let them know you're ready to help with Nigerian legal matters."""

    def is_legal_query(self, query: str) -> Tuple[bool, str, str]:
        """Determine if a query requires legal knowledge or is casual conversation."""
        query_lower = query.lower()

        legal_keywords = ['law', 'legal', 'right', 'tenant', 'landlord', 'rent',
                         'evict', 'contract', 'consumer', 'employment', 'wage',
                         'section', 'act', 'fccpc', 'tenancy', 'labour', 'notice',
                         'termination', 'refund', 'complaint', 'court', 'sue',
                         'eviction', 'lease', 'employer', 'employee', 'salary',
                         'compensation', 'dismiss', 'terminate', 'agreement']

        casual_keywords = ['hi', 'hello', 'hey', 'thanks', 'thank you', 'bye',
                          'good morning', 'good afternoon', 'good evening',
                          'how are you', "what's your name", 'who are you']

        if any(keyword in query_lower for keyword in legal_keywords):
            return True, "HIGH", "Legal keywords detected"

        if len(query.split()) <= 5:
            if any(casual in query_lower for casual in casual_keywords):
                return False, "HIGH", "Casual greeting detected"

        return True, "MEDIUM", "Defaulting to legal query"

    def retrieve_legal_context(self, query: str, n_results: int = 5) -> Dict[str, Any]:
        """Retrieve relevant legal context from ChromaDB."""
        results = self.collection.query(
            query_texts=[query],
            n_results=n_results
        )

        formatted_context = []
        sources = []

        for doc, meta, dist in zip(
            results['documents'][0],
            results['metadatas'][0],
            results['distances'][0]
        ):
            citation = f"{meta['document']} - {meta['section']}"
            context_block = f"""
[Source: {citation}]
{doc}
---"""
            formatted_context.append(context_block)
            sources.append({
                'document': meta['document'],
                'section': meta['section'],
                'part': meta.get('part', 'N/A'),
                'distance': dist
            })

        return {
            'context': '\n'.join(formatted_context),
            'sources': sources,
            'has_relevant_context': len(formatted_context) > 0
        }

    def web_search_nigerian_law(self, query: str) -> str:
        """Search the web for Nigerian legal information."""
        enhanced_query = f"{query} Nigerian law Nigeria"

        try:
            results = self.search.results(enhanced_query)
            formatted_results = ["WEB SEARCH RESULTS:\n"]

            if 'organic' in results:
                for i, result in enumerate(results['organic'][:3], 1):
                    formatted_results.append(
                        f"{i}. {result.get('title', 'No title')}\n"
                        f"   Source: {result.get('link', 'No link')}\n"
                        f"   Summary: {result.get('snippet', 'No description')}\n"
                    )

            return '\n'.join(formatted_results)

        except Exception as e:
            return f"Web search failed: {str(e)}"

    def assess_context_sufficiency(self, query: str, context: str) -> tuple[bool, str]:
        """Quick check if we have enough context."""
        if len(context) > 200:
            return True, "Context available from legal documents"
        return False, "Insufficient context, will search web"

    def generate_casual_response(self, query: str) -> str:
        """Generate a casual conversational response."""
        messages = [
            SystemMessage(content=self.casual_system_prompt),
            *self.chat_history.messages,
            HumanMessage(content=query)
        ]
        response = self.llm.invoke(messages)
        return response.content

    def generate_legal_response(
        self,
        query: str,
        context: str,
        sources: List[Dict],
        web_results: Optional[str] = None
    ) -> str:
        """Generate legal response using Gemini with context and sources."""
        if web_results:
            full_context = f"""LEGAL DOCUMENT CONTEXT:
{context}
{web_results}
NOTE: Prioritize information from legal documents. Use web results only as supplementary information."""
        else:
            full_context = f"""LEGAL DOCUMENT CONTEXT:
{context}"""

        messages = [
            SystemMessage(content=self.legal_system_prompt),
            *self.chat_history.messages,
            HumanMessage(content=f"""User Question: {query}
{full_context}
Based on the above legal context and sources, provide a comprehensive answer.""")
        ]

        response = self.llm.invoke(messages)
        answer = response.content

        if sources:
            answer += "\n\n📚 **LEGAL SOURCES CITED:**\n"
            for i, source in enumerate(sources, 1):
                answer += f"{i}. {source['document']} - {source['section']}"
                if source['part'] != 'N/A':
                    answer += f" ({source['part']})"
                answer += f" [Relevance: {(1 - source['distance']):.1%}]\n"

        return answer

    def chat(self, user_query: str) -> Dict[str, Any]:
        """Main chat function - orchestrates the entire pipeline."""
        is_legal, confidence, reasoning = self.is_legal_query(user_query)

        if not is_legal:
            answer = self.generate_casual_response(user_query)
            self.chat_history.add_user_message(user_query)
            self.chat_history.add_ai_message(answer)

            return {
                'answer': answer,
                'sources': [],
                'used_web_search': False,
                'context_sufficient': True,
                'is_legal_query': False,
                'query_classification': reasoning
            }

        retrieval_results = self.retrieve_legal_context(user_query, n_results=5)
        context = retrieval_results['context']
        sources = retrieval_results['sources']

        is_sufficient, reasoning = self.assess_context_sufficiency(user_query, context)

        web_results = None
        if not is_sufficient:
            web_results = self.web_search_nigerian_law(user_query)

        answer = self.generate_legal_response(user_query, context, sources, web_results)

        self.chat_history.add_user_message(user_query)
        self.chat_history.add_ai_message(answer)

        return {
            'answer': answer,
            'sources': sources,
            'used_web_search': web_results is not None,
            'context_sufficient': is_sufficient,
            'is_legal_query': True,
            'query_classification': reasoning
        }

    def clear_memory(self):
        """Clear conversation memory."""
        self.chat_history.clear()


def create_gradio_interface(chatbot: NigerianLegalAidChatbot):
    """Create a modern, beautiful Gradio interface for the chatbot."""

    def chat_interface(message, history):
        """Handle chat interactions"""
        if not message or not message.strip():
            return history

        try:
            history.append([message, None])
            result = chatbot.chat(message)
            response = result['answer']

            metadata = "\n\n---\n"
            if result['is_legal_query']:
                metadata += "⚖️ Legal Query | "
                metadata += f"{'🌐 Web Search Used' if result['used_web_search'] else '📚 Database Only'}"
            else:
                metadata += "💬 Casual Conversation"

            response += metadata
            history[-1][1] = response
            return history

        except Exception as e:
            error_msg = f"❌ Error: {str(e)}\n\nPlease try again or rephrase your question."
            history[-1][1] = error_msg
            return history

    def clear_chat():
        """Clear chat history."""
        chatbot.clear_memory()
        return None, []

    # Modern Custom CSS with gradient backgrounds and smooth animations
    custom_css = """
    @import url('https://fonts.googleapis.com/css2?family=Inter:wght@400;500;600;700&display=swap');
    
    * {
        font-family: 'Inter', -apple-system, BlinkMacSystemFont, sans-serif !important;
    }
    
    .gradio-container {
        max-width: 1400px !important;
        margin: 0 auto !important;
    }
    
    /* Modern Header with Nigerian Green */
    .header-container {
        background: linear-gradient(135deg, #008751 0%, #00b36b 50%, #008751 100%);
        padding: 40px 30px;
        border-radius: 20px;
        margin-bottom: 30px;
        box-shadow: 0 10px 40px rgba(0, 135, 81, 0.3);
        position: relative;
        overflow: hidden;
    }
    
    .header-container::before {
        content: '';
        position: absolute;
        top: -50%;
        right: -50%;
        width: 200%;
        height: 200%;
        background: radial-gradient(circle, rgba(255,255,255,0.1) 0%, transparent 70%);
        animation: rotate 20s linear infinite;
    }
    
    @keyframes rotate {
        from { transform: rotate(0deg); }
        to { transform: rotate(360deg); }
    }
    
    .header-content {
        position: relative;
        z-index: 1;
        text-align: center;
        color: white;
    }
    
    .header-title {
        font-size: 3em;
        font-weight: 700;
        margin: 0;
        text-shadow: 2px 2px 4px rgba(0,0,0,0.2);
        letter-spacing: -1px;
    }
    
    .header-subtitle {
        font-size: 1.2em;
        margin-top: 10px;
        opacity: 0.95;
        font-weight: 500;
    }
    
    .header-badge {
        display: inline-block;
        background: rgba(255,255,255,0.2);
        padding: 8px 20px;
        border-radius: 20px;
        margin-top: 15px;
        backdrop-filter: blur(10px);
        border: 1px solid rgba(255,255,255,0.3);
    }
    
    /* Modern Chat Container */
    .chat-container {
        background: white;
        border-radius: 20px;
        padding: 30px;
        box-shadow: 0 5px 30px rgba(0,0,0,0.08);
        border: 1px solid rgba(0,0,0,0.05);
    }
    
    /* Sidebar Styling */
    .sidebar-card {
        background: linear-gradient(135deg, #f8f9fa 0%, #ffffff 100%);
        border-radius: 15px;
        padding: 25px;
        box-shadow: 0 3px 15px rgba(0,0,0,0.06);
        border: 1px solid rgba(0,0,0,0.05);
        margin-bottom: 20px;
    }
    
    .sidebar-title {
        font-size: 1.3em;
        font-weight: 700;
        color: #008751;
        margin-bottom: 15px;
        display: flex;
        align-items: center;
        gap: 10px;
    }
    
    .topic-item {
        background: white;
        padding: 12px 15px;
        border-radius: 10px;
        margin-bottom: 10px;
        border-left: 4px solid #008751;
        box-shadow: 0 2px 8px rgba(0,0,0,0.04);
        transition: all 0.3s ease;
    }
    
    .topic-item:hover {
        transform: translateX(5px);
        box-shadow: 0 4px 12px rgba(0,135,81,0.15);
    }
    
    .topic-title {
        font-weight: 600;
        color: #008751;
        margin-bottom: 5px;
    }
    
    .topic-desc {
        font-size: 0.85em;
        color: #666;
        line-height: 1.4;
    }
    
    /* Modern Button Styling */
    .btn-primary {
        background: linear-gradient(135deg, #008751 0%, #00b36b 100%) !important;
        border: none !important;
        border-radius: 12px !important;
        padding: 12px 24px !important;
        font-weight: 600 !important;
        box-shadow: 0 4px 15px rgba(0,135,81,0.3) !important;
        transition: all 0.3s ease !important;
    }
    
    .btn-primary:hover {
        transform: translateY(-2px) !important;
        box-shadow: 0 6px 20px rgba(0,135,81,0.4) !important;
    }
    
    .btn-secondary {
        background: linear-gradient(135deg, #6c757d 0%, #495057 100%) !important;
        border: none !important;
        border-radius: 12px !important;
        padding: 12px 24px !important;
        font-weight: 600 !important;
        transition: all 0.3s ease !important;
    }
    
    /* Textbox Styling */
    textarea, input {
        border-radius: 12px !important;
        border: 2px solid #e0e0e0 !important;
        transition: all 0.3s ease !important;
    }
    
    textarea:focus, input:focus {
        border-color: #008751 !important;
        box-shadow: 0 0 0 3px rgba(0,135,81,0.1) !important;
    }
    
    /* Example Cards */
    .example-container {
        display: grid;
        grid-template-columns: repeat(auto-fit, minmax(280px, 1fr));
        gap: 15px;
        margin-top: 20px;
    }
    
    .example-card {
        background: linear-gradient(135deg, #f8f9fa 0%, #ffffff 100%);
        padding: 15px 20px;
        border-radius: 12px;
        border: 2px solid #e0e0e0;
        cursor: pointer;
        transition: all 0.3s ease;
        text-align: left;
    }
    
    .example-card:hover {
        border-color: #008751;
        transform: translateY(-3px);
        box-shadow: 0 5px 20px rgba(0,135,81,0.15);
    }
    
    /* Footer Styling */
    .footer {
        text-align: center;
        padding: 30px;
        margin-top: 40px;
        background: linear-gradient(135deg, #f8f9fa 0%, #ffffff 100%);
        border-radius: 15px;
        border: 1px solid rgba(0,0,0,0.05);
    }
    
    .footer-title {
        font-weight: 600;
        color: #333;
        margin-bottom: 10px;
    }
    
    .footer-text {
        color: #666;
        font-size: 0.9em;
    }
    
    /* Warning Box */
    .warning-box {
        background: linear-gradient(135deg, #fff3cd 0%, #fff8e1 100%);
        border-left: 4px solid #ffc107;
        padding: 15px 20px;
        border-radius: 10px;
        margin: 15px 0;
    }
    
    .warning-title {
        font-weight: 700;
        color: #856404;
        margin-bottom: 5px;
    }
    
    .warning-text {
        color: #856404;
        font-size: 0.9em;
    }
    
    /* Chatbot Messages */
    .message {
        border-radius: 15px !important;
        padding: 15px !important;
    }
    
    .user.message {
        background: linear-gradient(135deg, #008751 0%, #00b36b 100%) !important;
        color: white !important;
    }
    
    .bot.message {
        background: #f8f9fa !important;
        border: 1px solid #e0e0e0 !important;
    }
    """

    # Create the interface
    with gr.Blocks(css=custom_css, title="🇳🇬 Nigerian Legal Aid Chatbot", theme=gr.themes.Soft()) as interface:

        # Modern Header
        gr.HTML("""
            <div class="header-container">
                <div class="header-content">
                    <h1 class="header-title">🇳🇬 Nigerian Legal Aid Chatbot</h1>
                    <p class="header-subtitle">Your AI-Powered Legal Assistant for Nigerian Law</p>
                    <div class="header-badge">
                        ⚖️ FCCPC  •  🏠 Tenancy Law  •  💼 Labour Act
                    </div>
                </div>
            </div>
        """)

        with gr.Row():
            # Main Chat Area
            with gr.Column(scale=7):
                with gr.Group():
                    chatbot_ui = gr.Chatbot(
                        height=550,
                        label="💬 Chat with Your Legal Assistant",
                        bubble_full_width=False,
                        show_label=True,
                        avatar_images=(None, "https://huggingface.co/spaces/assets/huggingface-logo.svg")
                    )

                with gr.Row():
                    msg = gr.Textbox(
                        label="Your Message",
                        placeholder="💭 Ask about Nigerian law or start a conversation...",
                        lines=2,
                        scale=5,
                        show_label=False
                    )
                    with gr.Column(scale=1, min_width=100):
                        submit_btn = gr.Button("Send 📤", variant="primary", size="lg")
                        clear_btn = gr.Button("Clear 🗑️", variant="secondary", size="lg")

            # Modern Sidebar
            with gr.Column(scale=3):
                with gr.Group():
                    gr.HTML("""
                        <div class="sidebar-card">
                            <div class="sidebar-title">📚 Legal Topics Covered</div>
                            
                            <div class="topic-item">
                                <div class="topic-title">⚖️ FCCPC Act 2018</div>
                                <div class="topic-desc">Consumer rights, product safety, fair trading practices</div>
                            </div>
                            
                            <div class="topic-item">
                                <div class="topic-title">🏠 Tenancy Law 2011</div>
                                <div class="topic-desc">Rent regulations, eviction procedures, tenant rights</div>
                            </div>
                            
                            <div class="topic-item">
                                <div class="topic-title">💼 Labour Act</div>
                                <div class="topic-desc">Employment terms, worker rights, contracts</div>
                            </div>
                        </div>
                    """)

                with gr.Group():
                    gr.HTML("""
                        <div class="sidebar-card">
                            <div class="sidebar-title">💡 How to Get the Best Answers</div>
                            <ul style="color: #666; line-height: 1.8; padding-left: 20px;">
                                <li>Be specific with your questions</li>
                                <li>Include relevant details</li>
                                <li>Ask for section citations</li>
                                <li>Feel free to chat casually too!</li>
                            </ul>
                        </div>
                    """)

                with gr.Group():
                    gr.HTML("""
                        <div class="warning-box">
                            <div class="warning-title">⚠️ Important Disclaimer</div>
                            <div class="warning-text">
                                This chatbot provides legal information, not professional legal advice. 
                                Always consult a licensed Nigerian lawyer for official legal representation.
                            </div>
                        </div>
                    """)

        # Example Questions Section
        gr.HTML("""
            <div style="margin: 30px 0 20px 0;">
                <h3 style="color: #008751; font-weight: 700; font-size: 1.5em; text-align: center;">
                    📝 Try These Example Questions
                </h3>
            </div>
        """)
        
        with gr.Row():
            gr.Examples(
                examples=[
                    "What are my rights as a tenant in Lagos?",
                    "How much advance rent can a landlord demand?",
                    "What is the notice period for terminating employment?",
                    "Can a landlord increase rent without notice?",
                    "What are consumer protection rights in Nigeria?",
                    "Hello! What can you help me with?"
                ],
                inputs=msg,
                label=""
            )

        # Modern Footer
        gr.HTML("""
            <div class="footer">
                <div class="footer-title">🤖 Powered by Advanced AI Technology</div>
                <div class="footer-text">
                    Gemini AI • RAG with ChromaDB • Web Search Enhanced
                </div>
                <div class="footer-text" style="margin-top: 10px;">
                    Built with ❤️ for Nigerian Legal Aid • © 2024
                </div>
            </div>
        """)

        # Event handlers
        msg.submit(chat_interface, inputs=[msg, chatbot_ui], outputs=[chatbot_ui])
        submit_btn.click(chat_interface, inputs=[msg, chatbot_ui], outputs=[chatbot_ui])
        clear_btn.click(clear_chat, outputs=[chatbot_ui, chatbot_ui])

    return interface


# Main execution
if __name__ == "__main__":
    GOOGLE_API_KEY = os.getenv("GOOGLE_API_KEY")
    SERPER_API_KEY = os.getenv("SERPER_API_KEY")

    if not GOOGLE_API_KEY:
        GOOGLE_API_KEY = "your-google-api-key-here"
    if not SERPER_API_KEY:
        SERPER_API_KEY = "your-serper-api-key-here"

    if GOOGLE_API_KEY == "your-google-api-key-here" or SERPER_API_KEY == "your-serper-api-key-here":
        print("\n⚠️  WARNING: Using placeholder API keys!")
        print("For Hugging Face Spaces, set GOOGLE_API_KEY and SERPER_API_KEY in Space secrets.\n")

    print("\n" + "="*60)
    print("🚀 Initializing Nigerian Legal Aid Chatbot...")
    print("="*60 + "\n")

    try:
        chatbot = NigerianLegalAidChatbot(
            google_api_key=GOOGLE_API_KEY,
            serper_api_key=SERPER_API_KEY,
            chroma_persist_directory="./legal_chroma_db"
        )

        print("✓ Chatbot initialized successfully!")
        print("✓ Creating modern Gradio interface...\n")

        interface = create_gradio_interface(chatbot)

        print("="*60)
        print("🎉 Launching web interface...")
        print("="*60 + "\n")

        interface.launch(
            server_name="0.0.0.0",
            server_port=7860,
            show_error=True
        )

    except Exception as e:
        print(f"\n❌ Failed to initialize: {e}")
        print("\n📋 Checklist:")
        print("1. ✓ Valid API keys set?")
        print("2. ✓ ChromaDB folder exists?")
        print("3. ✓ Dependencies installed?")
        raise