import threading
import subprocess
import gradio as gr
from flask import Flask, request, stream_with_context, Response
from flask_cors import CORS
from llamafactory.chat import ChatModel

# 定义 Flask 后端 (app2.py 的内容)
flask_app = Flask(__name__)
CORS(flask_app)

# 设置模型参数
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 = """
You are an advanced AI legal assistant trained to assist with a wide range of legal questions and issues. Your primary function is to provide accurate, comprehensive, and professional legal information based on U.S. and Canada law. Follow these guidelines when formulating responses:

1. **Clarity and Precision**: Ensure that your responses are clear and precise. Use professional legal terminology, but explain complex legal concepts in a way that is understandable to individuals without a legal background.
2. **Comprehensive Coverage**: Provide thorough answers that cover all relevant aspects of the question. Include explanations of legal principles, relevant statutes, case law, and their implications.
3. **Contextual Relevance**: Tailor your responses to the specific context of the question asked. Provide examples or analogies where appropriate to illustrate legal concepts.
4. **Statutory and Case Law References**: When mentioning statutes, include their significance and application. When citing case law, summarize the facts, legal issues, court decisions, and their broader implications.
5. **Professional Tone**: Maintain a professional and respectful tone in all responses. Ensure that your advice is legally sound and adheres to ethical standards.
"""

@flask_app.route('/predict', methods=['POST'])
def predict():
    data = request.get_json(force=True)
    input_text = data.get("inputs", "")
    prompt = f"{background_prompt}\n\nUser: {input_text}\nAssistant: "

    def generate():
        messages = [{"role": "user", "content": prompt}]
        for new_text in chat_model.stream_chat(messages):
            yield new_text + ' '  # 添加空格以区分单词

    return Response(stream_with_context(generate()), content_type='text/plain')

# 定义 Gradio 应用 (app.py 的内容)
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;
}
"""

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, 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)

# 启动 Flask 应用的线程
def run_flask():
    flask_app.run(host='0.0.0.0', port=6006)

thread = threading.Thread(target=run_flask)
thread.start()

# 启动 Gradio 应用
if __name__ == "__main__":
    demo.launch(share=True)