import os
import streamlit as st
from groq import Groq
from utils import load_law_data
from prompt_templates import legal_assistant_prompt

# Initialize Groq client
GROQ_API_KEY = "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"
client = Groq(api_key=GROQ_API_KEY)

# Load legal data
law_data = load_law_data("data/pakistani_law/")

# Define the chatbot response function
def get_response(user_input):
    chat_completion = client.chat.completions.create(
        messages=[
            {"role": "system", "content": legal_assistant_prompt},
            {"role": "user", "content": user_input},
        ],
        model="llama-3.3-70b-versatile",
    )
    return chat_completion.choices[0].message.content

# Streamlit app
def main():
    st.title("Pakistani Legal Chatbot")
    st.write("Ask any question related to Pakistani law, and I'll assist you!")

    # Input text box
    user_input = st.text_input("Your Question", placeholder="Type your legal query here...")

    if st.button("Get Response"):
        if user_input.strip():
            # Get chatbot response
            response = get_response(user_input)
            # Display chatbot response
            st.text_area("Chatbot's Response", value=response, height=200)
        else:
            st.warning("Please enter a valid question.")

if __name__ == "__main__":
    main()