# app.py

import os
import gradio as gr
from langchain_community.vectorstores import FAISS
from langchain_community.embeddings import HuggingFaceEmbeddings
from huggingface_hub import InferenceClient

# Load FAISS index and embedding model
embedding_model = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
db = FAISS.load_local("faiss_index", embedding_model)

# Load Hugging Face Inference API client
client = InferenceClient(
    model="HuggingFaceH4/zephyr-7b-beta",
    token=os.getenv("HF_TOKEN")  # Make sure this is set in your environment
)

def ask_law_bot(query):
    try:
        results = db.similarity_search(query, k=5, filter={"section": "PPC"})
        if not results:
            return "❌ No relevant content found for this topic."

        context = "\n\n".join([doc.page_content for doc in results if len(doc.page_content.strip()) > 100])

        prompt = f"""You are a legal assistant helping users understand Pakistani law.
Respond to the question using the given legal context. Your answer must follow these rules:
- Use numbered bullet points (1. 2. 3.)
- Reference relevant law sections like (section 220(b))
- Be concise, clear, and avoid repetition
- Use "YES" or "NO" if the question requires binary response
Context:
{context}
Question: {query}
Answer:"""

        response = client.chat_completion(
            messages=[
                {"role": "system", "content": "You are a helpful and concise legal assistant for Pakistani law."},
                {"role": "user", "content": prompt}
            ],
            max_tokens=512
        )

        return response.choices[0].message["content"].strip()

    except Exception as e:
        return f"❌ Error: {e}"

# Gradio UI
gr.Interface(
    fn=ask_law_bot,
    inputs=gr.Textbox(lines=2, placeholder="e.g., What is the punishment for theft?"),
    outputs=gr.Textbox(label="📘 Legal Answer"),
    title="⚖️ Ask Pakistan Law — Powered by Zephyr 7B",
    description="Ask questions from Pakistan's law using FAISS retrieval + Zephyr-7B via Hugging Face API.",
    examples=[
        "What is the punishment for theft?",
        "What are the duties of the Commission?",
        "What is the process of appeal under this law?"
    ]
).launch(share=True, debug=True)