#!/usr/bin/env python3
"""
Indian Legal AI Assistant using Hugging Face LLM
This app provides a chat interface for querying Indian laws using the
invincibleambuj/Ambuj-Tripathi-Indian-Legal-Llama-GGUF model.
"""

import os
import gradio as gr
from huggingface_hub import hf_hub_download
from llama_cpp import Llama

# Model configuration
MODEL_REPO = "invincibleambuj/Ambuj-Tripathi-Indian-Legal-Llama-GGUF"
MODEL_FILE = "ambuj-tripathi-indian-legal-llama.Q4_K_M.gguf"

# Global model instance
llm = None


def load_model():
    """Load the GGUF model from Hugging Face Hub"""
    global llm
    if llm is None:
        print(f"Downloading model from {MODEL_REPO}...")
        model_path = hf_hub_download(
            repo_id=MODEL_REPO,
            filename=MODEL_FILE,
            cache_dir="./models"
        )
        print(f"Model downloaded to: {model_path}")
        print("Loading model into memory...")
        llm = Llama(
            model_path=model_path,
            n_ctx=2048,  # Context window
            n_threads=4,  # CPU threads
            n_gpu_layers=0,  # Set to 0 for CPU-only, increase if GPU available
        )
        print("Model loaded successfully!")
    return llm


def chat_with_legal_ai(message, history):
    """
    Process user message and generate response using the Indian Legal LLM
    
    Args:
        message: User's input message
        history: Chat history (list of [user_msg, bot_msg] pairs)
    
    Returns:
        Response from the LLM
    """
    try:
        model = load_model()
        
        # Build conversation context
        conversation = ""
        if history:
            for user_msg, bot_msg in history:
                conversation += f"User: {user_msg}\nAssistant: {bot_msg}\n"
        
        # Add current message
        prompt = f"""You are an expert Indian legal AI assistant. You have deep knowledge of Indian laws, acts, and legal procedures. Provide accurate, helpful, and concise answers to legal queries.

{conversation}User: {message}
Assistant:"""
        
        # Generate response
        response = model(
            prompt,
            max_tokens=512,
            temperature=0.7,
            top_p=0.95,
            echo=False,
            stop=["User:", "\n\n\n"]
        )
        
        return response['choices'][0]['text'].strip()
    
    except Exception as e:
        return f"Error: {str(e)}\n\nPlease ensure the model is properly downloaded and configured."


def create_gradio_interface():
    """Create and return the Gradio chat interface"""
    
    # Custom CSS for better appearance
    custom_css = """
    .container {
        max-width: 900px;
        margin: auto;
        padding: 20px;
    }
    #title {
        text-align: center;
        color: #1e3a8a;
    }
    """
    
    with gr.Blocks(css=custom_css, title="Indian Legal AI Assistant") as demo:
        gr.Markdown(
            """
            # 🏛️ Indian Legal AI Assistant
            
            Ask questions about Indian laws, acts, and legal procedures. This AI assistant is powered by 
            the **Ambuj Tripathi Indian Legal Llama** model, trained on Indian legal documents.
            
            **Examples:**
            - What is the Indian Penal Code?
            - Explain Section 377 of IPC
            - What are the grounds for divorce under Hindu Marriage Act?
            - Explain the Right to Information Act
            """,
            elem_id="title"
        )
        
        chatbot = gr.Chatbot(
            height=500,
            label="Chat History",
            show_label=True,
            elem_id="chatbot"
        )
        
        with gr.Row():
            msg = gr.Textbox(
                label="Your Question",
                placeholder="Ask about Indian laws...",
                lines=2,
                scale=4
            )
            submit = gr.Button("Send", variant="primary", scale=1)
        
        with gr.Row():
            clear = gr.Button("Clear Chat")
        
        gr.Markdown(
            """
            ---
            **Disclaimer:** This AI assistant provides general information only and should not be considered 
            legal advice. For specific legal matters, please consult a qualified legal professional.
            
            **Model:** [invincibleambuj/Ambuj-Tripathi-Indian-Legal-Llama-GGUF](https://huggingface.co/invincibleambuj/Ambuj-Tripathi-Indian-Legal-Llama-GGUF)
            """
        )
        
        # Set up event handlers
        msg.submit(chat_with_legal_ai, [msg, chatbot], [chatbot])
        submit.click(chat_with_legal_ai, [msg, chatbot], [chatbot])
        msg.submit(lambda: "", None, [msg])
        submit.click(lambda: "", None, [msg])
        clear.click(lambda: None, None, [chatbot])
    
    return demo


if __name__ == "__main__":
    print("Starting Indian Legal AI Assistant...")
    print(f"Using model: {MODEL_REPO}/{MODEL_FILE}")
    
    # Pre-load the model to avoid delays on first query
    print("\nPre-loading model (this may take a few minutes on first run)...")
    load_model()
    
    # Launch the Gradio interface
    demo = create_gradio_interface()
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False,
        show_error=True
    )