import os
os.system('pip install transformers')
os.system('pip install datasets')
os.system('pip install gradio')
os.system('pip install minijinja')
os.system('pip install PyMuPDF')
os.system('pip install PyPDF2')
os.system('pip install pdf2image')
os.system('pip install gradio_pdf')

import gradio as gr
from huggingface_hub import InferenceClient
from transformers import pipeline, AutoTokenizer, AutoModelForMaskedLM
from datasets import load_dataset
import fitz  # PyMuPDF
from pathlib import Path

dir_ = Path(__file__).parent

dataset = load_dataset("ibunescu/qa_legal_dataset_train")

# Load the BERT model and tokenizer
tokenizer = AutoTokenizer.from_pretrained("google-bert/bert-base-uncased")
model = AutoModelForMaskedLM.from_pretrained("google-bert/bert-base-uncased")

# Create the fill-mask pipeline
pipe = pipeline("fill-mask", model=model, tokenizer=tokenizer)

client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")

def respond(
    message,
    history: list[tuple[str, str]],
    system_message,
    max_tokens,
    temperature,
    top_p,
):
    messages = [{"role": "system", "content": system_message}]

    for val in history:
        if val[0]:
            messages.append({"role": "user", "content": val[0]})
        if val[1]:
            messages.append({"role": "assistant", "content": val[1]})

    messages.append({"role": "user", "content": message})

    try:
        response = ""
        for message in client.chat_completion(
            messages,
            max_tokens=max_tokens,
            stream=True,
            temperature=temperature,
            top_p=top_p,
        ):
            token = message.choices[0].delta.content
            if token is not None:
                response += token
            yield response, history + [(message, response)]
    except Exception as e:
        print(f"Error during chat completion: {e}")
        yield "An error occurred during the chat completion.", history

def generate_case_outcome(prosecutor_response, defense_response):
    prompt = f"Prosecutor's arguments: {prosecutor_response}\n\nDefense's arguments: {defense_response}\n\nProvide details on who won the case and why. Provide reasons for your decision and provide a link to the source of the case."
    evaluation = ""
    try:
        for message in client.chat_completion(
            [{"role": "system", "content": "You are a legal expert evaluating the details of the case presented by the prosecution and the defense."},
             {"role": "user", "content": prompt}],
            max_tokens=512,
            stream=True,
            temperature=0.6,
            top_p=0.95,
        ):
            token = message.choices[0].delta.content
            if token is not None:
                evaluation += token
    except Exception as e:
        print(f"Error during case outcome generation: {e}")
        return "An error occurred during the case outcome generation."
    return evaluation

def determine_outcome(outcome):
    prosecutor_count = outcome.split().count("Prosecutor")
    defense_count = outcome.split().count("Defense")
    if prosecutor_count > defense_count:
        return "Prosecutor Wins"
    elif defense_count > prosecutor_count:
        return "Defense Wins"
    else:
        return "No clear winner"

# Custom CSS for white background and black text for input and output boxes
custom_css = """
body {
    background-color: #ffffff;
    color: #000000;
    font-family: Arial, sans-serif;
}
.gradio-container {
    max-width: 1000px;
    margin: 0 auto;
    padding: 20px;
    background-color: #ffffff;
    border: 1px solid #e0e0e0;
    border-radius: 8px;
    box-shadow: 0 2px 5px rgba(0, 0, 0, 0.1);
}
.gr-button {
    background-color: #ffffff !important;
    border-color: #ffffff !important;
    color: #000000 !important;
    margin: 5px;
}
.gr-button:hover {
    background-color: #ffffff !important;
    border-color: #004085 !important;
}
.gr-input, .gr-textbox, .gr-slider, .gr-markdown, .gr-chatbox {
    border-radius: 4px;
    border: 1px solid #ced4da;
    background-color: #ffffff !important;
    color: #000000 !important;
}
.gr-input:focus, .gr-textbox:focus, .gr-slider:focus {
    border-color: #ffffff;
    outline: 0;
    box-shadow: 0 0 0 0.2rem rgba(255, 255, 255, 1.0);
}
#flagging-button {
    display: none;
}
footer {
    display: none;
}
.chatbox .chat-container .chat-message {
    background-color: #ffffff !important;
    color: #000000 !important;
}
.chatbox .chat-container .chat-message-input {
    background-color: #ffffff !important;
    color: #000000 !important;
}
.gr-markdown {
    background-color: #ffffff !important;
    color: #000000 !important;
}
.gr-markdown h1, .gr-markdown h2, .gr-markdown h3, .gr-markdown h4, .gr-markdown h5, .gr-markdown h6, .gr-markdown p, .gr-markdown ul, .gr-markdown ol, .gr-markdown li {
    color: #000000 !important;
}
.score-box {
    width: 60px;
    height: 60px;
    display: flex;
    align-items: center;
    justify-content: center;
    font-size: 12px;
    font-weight: bold;
    color: black;
    margin: 5px;
}
.scroll-box {
    max-height: 200px;
    overflow-y: scroll;
    border: 1px solid #ced4da;
    padding: 10px;
    border-radius: 4px;
}
"""

def chat_between_bots(system_message1, system_message2, max_tokens, temperature, top_p, history1, history2, shared_history, message):
    response1, history1 = list(respond(message, history1, system_message1, max_tokens, temperature, top_p))[-1]
    response2, history2 = list(respond(message, history2, system_message2, max_tokens, temperature, top_p))[-1]
    shared_history.append(f"Prosecutor: {response1}")
    shared_history.append(f"Defense Attorney: {response2}")
    
    max_length = max(len(response1), len(response2))
    response1 = response1[:max_length]
    response2 = response2[:max_length]
    
    outcome = generate_case_outcome(response1, response2)
    winner = determine_outcome(outcome)
    
    return response1, response2, history1, history2, shared_history, outcome

def get_top_5_cases():
    prompt = "List 5 random high-profile legal cases that have received significant media attention and are currently ongoing. Just a list of case names and numbers."
    response = ""
    for message in client.chat_completion(
        [{"role": "system", "content": "You are a legal research expert, able to provide information about high-profile legal cases."},
         {"role": "user", "content": prompt}],
        max_tokens=512,
        stream=True,
        temperature=0.6,
        top_p=0.95,
    ):
        token = message.choices[0].delta.content
        if token is not None:
            response += token
    return response

def add_message(history, message):
    for x in message["files"]:
        history.append(((x,), None))
    if message["text"] is not None:
        history.append((message["text"], None))
    return history, gr.MultimodalTextbox(value=None, interactive=True)

def print_like_dislike(x: gr.LikeData):
    print(x.index, x.value, x.liked)

def reset_conversation():
    return [], [], "", "", ""

def save_conversation(history1, history2, shared_history):
    return history1, history2, shared_history

def ask_about_case_outcome(shared_history, question):
    prompt = f"Case Outcome: {shared_history}\n\nQuestion: {question}\n\nAnswer:"
    response = ""
    for message in client.chat_completion(
        [{"role": "system", "content": "Give the real case details of the case being argued, from a verifed source."},
         {"role": "user", "content": prompt}],
        max_tokens=512,
        stream=True,
        temperature=0.7,
        top_p=0.95,
    ):
        token = message.choices[0].delta.content
        if token is not None:
            response += token
    return response

with gr.Blocks(css=custom_css) as demo:
    history1 = gr.State([])
    history2 = gr.State([])
    shared_history = gr.State([])
    top_10_cases = gr.State("")

    with gr.Tab("Argument Evaluation"):
        with gr.Row():
            with gr.Column(scale=1):
                top_5_btn = gr.Button("Give me the top 5 cases")
                top_5_output = gr.Textbox(label="Top 5 Cases", interactive=False, elem_classes=["scroll-box"])
                top_5_btn.click(get_top_5_cases, outputs=top_5_output)
            with gr.Column(scale=2):
                message = gr.Textbox(label="Case to Argue")
                system_message1 = gr.State("You are an expert Prosecutor. Give your best arguments for the case on behalf of the prosecution.")
                system_message2 = gr.State("You are an expert Defense Attorney. Give your best arguments for the case on behalf of the Defense.")
                max_tokens = gr.State(512)
                temperature = gr.State(0.5)
                top_p = gr.State(0.95)
                
                with gr.Row():
                    with gr.Column(scale=4):
                        prosecutor_response = gr.Textbox(label="Prosecutor's Response", interactive=True, elem_classes=["scroll-box"])
                    with gr.Column(scale=1):
                        prosecutor_score_color = gr.HTML()
                    
                    with gr.Column(scale=4):
                        defense_response = gr.Textbox(label="Defense Attorney's Response", interactive=True, elem_classes=["scroll-box"])
                    with gr.Column(scale=1):
                        defense_score_color = gr.HTML()
                
                outcome = gr.Textbox(label="Outcome", interactive=True, elem_classes=["scroll-box"])
                
                with gr.Row():
                    submit_btn = gr.Button("Argue")
                    clear_btn = gr.Button("Clear and Reset")
                    save_btn = gr.Button("Save Conversation")
                
                submit_btn.click(chat_between_bots, inputs=[system_message1, system_message2, max_tokens, temperature, top_p, history1, history2, shared_history, message], outputs=[prosecutor_response, defense_response, history1, history2, shared_history, outcome])
                clear_btn.click(reset_conversation, outputs=[history1, history2, shared_history, prosecutor_response, defense_response, outcome])
                save_btn.click(save_conversation, inputs=[history1, history2, shared_history], outputs=[history1, history2, shared_history])
    
    with gr.Tab("Case Outcome Chat"):
        case_question = gr.Textbox(label="Ask a Question about the Case Outcome")
        case_answer = gr.Textbox(label="Answer", interactive=False, elem_classes=["scroll-box"])
        ask_case_btn = gr.Button("Ask")

        ask_case_btn.click(ask_about_case_outcome, inputs=[shared_history, case_question], outputs=case_answer)

demo.queue()
demo.launch()