import gradio as gr
import os
from threading import Thread
from llamafactory.chat import ChatModel
from llamafactory.extras.misc import torch_gc

# Set an environment variable
HF_TOKEN = os.environ.get("HF_TOKEN", None)

DESCRIPTION = '''
<div>
<h1 style="text-align: center;">AI Lawyer</h1>
</div>
'''

LICENSE = """
<p/>
---
Built with model "StevenChen16/Llama3-8B-Lawyer", based on "meta-llama/Meta-Llama-3-8B"
"""

PLACEHOLDER = """
<div style="padding: 30px; text-align: center; display: flex; flex-direction: column; align-items: center;">
   <h1 style="font-size: 28px; margin-bottom: 2px; opacity: 0.55;">AI Lawyer</h1>
   <p style="font-size: 18px; margin-bottom: 2px; opacity: 0.65;">Ask me anything about US and Canada law...</p>
</div>
"""

css = """
h1 {
  text-align: center;
  display: block;
}
#duplicate-button {
  margin: auto;
  color: white;
  background: #1565c0;
  border-radius: 100vh;
}
"""

args = dict(
    model_name_or_path="StevenChen16/llama3-8b-Lawyer",
    template="llama3",
    finetuning_type="lora",
    quantization_bit=8,
    use_unsloth=True,
)
chat_model = ChatModel(args)

background_prompt = """
As an AI legal assistant, you are a highly trained expert in U.S. and Canadian law. Your purpose is to provide accurate, comprehensive, and professional legal information to assist users with a wide range of legal questions and issues. 

When responding to queries, adhere to the following guidelines:

1. Clarity and Precision: 
- Provide clear, concise answers using precise legal terminology.
- Explain complex legal concepts in a manner accessible to non-legal professionals.

2. Comprehensive Coverage:
- Offer thorough, well-rounded responses that address all relevant aspects of the question.
- Explain pertinent legal principles, statutes, case law, and their implications.

3. Contextual Relevance: 
- Tailor your advice to the specific context of each question.
- Utilize examples or analogies to illustrate legal concepts when appropriate.

4. Statutory and Case Law References:
- When citing statutes, explain their relevance and application to the matter at hand.
- When referencing case law, summarize the key facts, legal issues, court decisions, and the broader implications of the ruling.

5. Professional Tone:
- Maintain a professional, respectful demeanor in all interactions.
- Ensure your advice is legally sound and adheres to the highest ethical standards.

Remember, your role is to provide general legal information and analysis. 

This is a detailed description of the case or general questions, or detailed instructions for you:
"""

def query_model(user_input, history):
    combined_query = background_prompt + user_input
    messages = [{"role": "user", "content": combined_query}]
    
    response = ""
    for new_text in chat_model.stream_chat(messages, max_new_tokens=512, temperature=0.9):
        response += new_text
        yield response
    
# Gradio block
chatbot = gr.Chatbot(height=450, placeholder=PLACEHOLDER, label='Gradio ChatInterface')

with gr.Blocks(css=css) as demo:
    gr.Markdown(DESCRIPTION)
    gr.DuplicateButton(value="Duplicate Space for private use", elem_id="duplicate-button")
    gr.ChatInterface(
        fn=query_model,
        chatbot=chatbot,
        examples=[
            ['What are the key differences between a sole proprietorship and a partnership?'],
            ['What legal steps should I take if I want to start a business in the US?'],
            ['Can you explain the concept of "duty of care" in negligence law?'],
            ['What are the legal requirements for obtaining a patent in Canada?'],
            ['How can I protect my intellectual property when sharing my idea with potential investors?']
        ],
        cache_examples=False,
    )
    gr.Markdown(LICENSE)

if __name__ == "__main__":
    demo.launch()