app.py

import os
import sys
import signal
import warnings
import gradio as gr
import config
import ingestion
import vector_store
import rag_engine
import logic

# Suppress warnings
warnings.filterwarnings("ignore", category=FutureWarning)

# Global variable to store the chain
rag_chain = None

def initialize_system_once():
    """
    Initialize the complete system only once.
    """
    global rag_chain
    if rag_chain is not None:
        return rag_chain

    print("Initializing LegalizeAI System...")
    
    # Initialize Embedding Model
    print("Loading embedding model...")
    embedding_model = vector_store.get_embedding_model()
    
    # Check if Vector Store exists
    if os.path.exists(config.CHROMA_DB_DIR) and os.listdir(config.CHROMA_DB_DIR):
        print(f"Loading existing vector store from {config.CHROMA_DB_DIR}...")
        v_store = vector_store.get_vector_store(embedding_model)
    else:
        print("No existing vector store found. Starting ingestion process...")
        docs = ingestion.load_documents()
        if not docs:
            # Create empty placeholder if no docs found to prevent crash, 
            # but warn hard.
            print("CRITICAL WARNING: No documents loaded. App will run but local search will fail.")
            # In a real app we might want to fail, but for UI it's better to stay up
            return None
            
        chunks = ingestion.split_documents(docs)
        print("Creating vector store...")
        v_store = vector_store.create_vector_store(chunks, embedding_model)
        
    # Setup Retriever
    retriever = vector_store.get_retriever(v_store)
    
    # Setup RAG Chain
    print("Initializing RAG chain...")
    rag_chain = rag_engine.create_rag_chain(retriever)
    
    print("System initialization complete!")
    return rag_chain

def chat_response(message, history):
    """
    Gradio chat function.
    """
    try:
        chain = initialize_system_once()
        if not chain:
            return "System Error: Failed to initialize AI chain. Please check server logs."
            
        response = logic.generate_hybrid_response(message, chain)
        return response
    except Exception as e:
        return f"An error occurred: {str(e)}"

# Custom CSS for a professional look
custom_css = """
body { background-color: #f0f2f5; }
footer { visibility: hidden !important; }

/* Custom Developer Footer */
.dev-footer {
    text-align: center;
    padding: 20px;
    margin-top: 30px;
    border-top: 1px solid #e5e7eb;
    color: #4b5563;
    background-color: transparent !important;
}

.dev-footer a {
    display: inline-flex;
    align-items: center;
    justify-content: center;
    margin: 0 10px;
    color: #4b5563;
    text-decoration: none;
    transition: color 0.2s;
}

.dev-footer a:hover {
    color: #1f2937;
}

.dev-footer svg {
    margin-right: 5px;
    width: 20px;
    height: 20px;
    fill: currentColor;
}

/* Chatbot styling */
#chatbot {
    min-height: 500px;
    max-height: 500px;
    overflow-y: auto;
}

/* Mobile responsiveness */
@media (max-width: 768px) {
    .dev-footer div {
        flex-direction: column;
        gap: 10px;
    }
    
    .dev-footer a {
        margin: 5px 0;
    }
    
    #chatbot {
        min-height: 400px;
        max-height: 400px;
    }
}
"""

# HTML for the footer
footer_html = """
<div class="dev-footer">
    <p>Developed by <strong>Muhammad Hashir Lodhi</strong></p>
    <div>
        <a href="https://github.com/HashirLodhi" target="_blank">
            <svg viewBox="0 0 24 24" xmlns="http://www.w3.org/2000/svg"><path d="M12 0C5.37 0 0 5.37 0 12c0 5.31 3.435 9.795 8.205 11.385.6.105.825-.255.825-.57 0-.285-.015-1.05-.015-2.055-3.33.72-4.035-1.605-4.035-1.605-.54-1.38-1.32-1.74-1.32-1.74-1.095-.75.09-.735.09-.735 1.2.09 1.845 1.245 1.845 1.245 1.065 1.83 2.805 1.305 3.495.99.105-.78.42-1.305.765-1.605-2.67-.3-5.46-1.335-5.46-5.925 0-1.305.465-2.385 1.23-3.225-.12-.3-.54-1.53.12-3.18 0 0 1.005-.315 3.3 1.23.96-.27 1.98-.405 3-.405 1.02 0 2.04.135 3 .405 2.28-1.545 3.285-1.23 3.285-1.23.66 1.65.24 2.88.12 3.18.765.84 1.23 1.905 1.23 3.225 0 4.605-2.805 5.625-5.475 5.925.435.375.81 1.095.81 2.22 0 1.605-.015 2.895-.015 3.285 0 .315.225.69.825.57A12.02 12.02 0 0024 12c0-6.63-5.37-12-12-12z"/></svg>
            GitHub
        </a>
        <a href="https://medium.com/@hashirlodhi145" target="_blank">
            <svg viewBox="0 0 24 24" xmlns="http://www.w3.org/2000/svg"><path d="M13.54 12a6.8 6.8 0 01-6.77 6.82A6.8 6.8 0 010 12a6.8 6.8 0 016.77-6.82A6.8 6.8 0 0113.54 12zM20.96 12c0 3.54-1.51 6.42-3.38 6.42-1.87 0-3.39-2.88-3.39-6.42s1.52-6.42 3.39-6.42 3.38 2.88 3.38 6.42M24 12c0 3.17-.53 5.75-1.19 5.75-.66 0-1.19-2.58-1.19-5.75s.53-5.75 1.19-5.75C23.47 6.25 24 8.83 24 12z"/></svg>
            Medium
        </a>
        <a href="https://www.linkedin.com/in/hashir-lodhi/" target="_blank">
            <svg viewBox="0 0 24 24" xmlns="http://www.w3.org/2000/svg"><path d="M20.447 20.452h-3.554v-5.569c0-1.328-.027-3.037-1.852-3.037-1.853 0-2.136 1.445-2.136 2.939v5.667H9.351V9h3.414v1.561h.046c.477-.9 1.637-1.85 3.37-1.85 3.601 0 4.267 2.37 4.267 5.455v6.286zM5.337 7.433c-1.144 0-2.063-.926-2.063-2.065 0-1.138.92-2.063 2.063-2.063 1.14 0 2.064.925 2.064 2.063 0 1.139-.925 2.065-2.064 2.065zm1.782 13.019H3.555V9h3.564v11.452zM22.225 0H1.771C.792 0 0 .774 0 1.729v20.542C0 23.227.792 24 1.771 24h20.451C23.2 24 24 23.227 24 22.271V1.729C24 .774 23.2 0 22.222 0h.003z"/></svg>
            LinkedIn
        </a>
    </div>
</div>
"""

# Create the Gradio Interface
def create_ui():
    initialize_system_once()

    # Define interaction functions
    def interact(message, history):
        if not message:
            return "", history
        
        # Initialize history if strictly None (though usually empty list)
        if history is None: 
            history = []

        # Add user message
        history.append({"role": "user", "content": message})
        
        # Get response
        try:
            response = chat_response(message, history)
        except Exception as e:
            response = f"Error: {str(e)}"
            
        # Add assistant response
        history.append({"role": "assistant", "content": response})
        return "", history

    def retry_last(history):
        if not history:
            return history, ""
        
        # Pop last message if it's assistant
        if history and history[-1]["role"] == "assistant":
            history.pop()
        
        # Pop user message to edit
        if history and history[-1]["role"] == "user":
            last_msg = history.pop()
            return history, last_msg["content"]
            
        return history, ""

    def regenerate_response(history):
        """Regenerate the last assistant response."""
        if not history or len(history) < 2:
            return history
        
        # Get the last user message
        last_user_msg = None
        for msg in reversed(history):
            if msg["role"] == "user":
                last_user_msg = msg["content"]
                break
        
        if last_user_msg:
            # Remove the last assistant response
            if history[-1]["role"] == "assistant":
                history.pop()
            
            # Generate new response
            try:
                response = chat_response(last_user_msg, history)
                history.append({"role": "assistant", "content": response})
            except Exception as e:
                history.append({"role": "assistant", "content": f"Error: {str(e)}"})
        
        return history

    # Create the Gradio Blocks
    with gr.Blocks(title="⚖️ LegalizeAI") as demo:
        gr.Markdown(
            """
            # ⚖️ LegalizeAI
            **Professional Assistant for Pakistani Law**
            
            Consulting Constitution of Pakistan, Pakistan Penal Code, and Real-time Web Sources.
            
            *Disclaimer: This AI provides legal information but is not a substitute for professional legal advice.*
            """
        )
        
        chatbot = gr.Chatbot(
            height=500,
            elem_id="chatbot",
            avatar_images=(None, "⚖️")
            # Removed show_copy_button for compatibility with older Gradio versions
        )
        
        with gr.Row():
            txt = gr.Textbox(
                scale=4,
                show_label=False,
                placeholder="Ask a legal question...",
                container=False,
                autofocus=True,
                max_lines=3
            )
            submit_btn = gr.Button("Send 🚀", scale=1, variant="primary")
        
        with gr.Row():
            retry_btn = gr.Button("Retry 🔄", size="sm")
            clear_btn = gr.Button("Clear 🗑️", size="sm")
            regenerate_btn = gr.Button("Regenerate 🔁", size="sm")

        # Example buttons logic
        examples = [
            "What is the punishment for theft in Pakistan?",
            "Explain Article 62 of the Constitution.",
            "Who is the current Prime Minister?",
            "What are my fundamental rights?",
            "Explain the process of filing a criminal case.",
            "What is defamation under Pakistani law?"
        ]
        
        gr.Examples(
            examples=examples,
            inputs=txt,
            examples_per_page=3
        )

        # Footer
        gr.HTML(footer_html)

        # Event Wiring
        submit_btn.click(interact, [txt, chatbot], [txt, chatbot])
        txt.submit(interact, [txt, chatbot], [txt, chatbot])
        
        retry_btn.click(retry_last, [chatbot], [chatbot, txt])
        clear_btn.click(lambda: None, None, chatbot, queue=False)
        regenerate_btn.click(regenerate_response, [chatbot], [chatbot])

    return demo

def cleanup(signum, frame):
    """Cleanup function for graceful shutdown."""
    print("\n🛑 Shutting down LegalizeAI gracefully...")
    sys.exit(0)

def main():
    """Main entry point."""
    try:
        # Register cleanup handlers
        signal.signal(signal.SIGINT, cleanup)
        signal.signal(signal.SIGTERM, cleanup)
        
        # Initialize system once before UI creation
        print("🚀 Starting LegalizeAI...")
        initialize_system_once()
        
        # Determine if running on Hugging Face Spaces
        is_huggingface = os.getenv("SPACE_ID") is not None
        server_name = "0.0.0.0" if is_huggingface else "127.0.0.1"
        
        if is_huggingface:
            print("🌐 Running on Hugging Face Spaces")
        else:
            print("💻 Running locally")
        
        # Using a professional theme
        theme = gr.themes.Soft(
            primary_hue="slate",
            secondary_hue="stone",
            neutral_hue="zinc",
            font=[gr.themes.GoogleFont("Inter"), "ui-sans-serif", "system-ui", "sans-serif"]
        )
        
        # Create UI
        demo = create_ui()
        
        # Launch with appropriate settings
        print(f"🌍 Server starting on {server_name}:7860")
        demo.launch(
            server_name=server_name,
            server_port=7860,
            share=False,  # Don't use share=True on Spaces
            favicon_path=None,
            theme=theme,
            show_error=True,
            quiet=False,
            show_api=False,
            css=custom_css  # CSS moved here for Gradio 6.0+ compatibility
        )
        
    except KeyboardInterrupt:
        print("\n👋 Shutdown requested by user")
        sys.exit(0)
    except Exception as e:
        print(f"❌ Fatal Error: {e}")
        import traceback
        traceback.print_exc()
        sys.exit(1)

if __name__ == "__main__":
    # Set environment variables for better performance
    os.environ["TOKENIZERS_PARALLELISM"] = "false"
    os.environ["PYTHONWARNINGS"] = "ignore"
    
    # Run the application
    main()