import os
import requests
from bs4 import BeautifulSoup
import gradio as gr
from huggingface_hub import InferenceClient
from transformers import pipeline
from datasets import load_dataset
import fitz  # PyMuPDF

# Load dataset
dataset = load_dataset("ibunescu/qa_legal_dataset_train")

# Different pipelines for different tasks
qa_pipeline = pipeline("question-answering", model="deepset/roberta-base-squad2")
summarization_pipeline = pipeline("summarization", model="facebook/bart-large-cnn")
mask_filling_pipeline = pipeline("fill-mask", model="nlpaueb/legal-bert-base-uncased")

# Inference client for chat completion
client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")

def respond(message, history, 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})

    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
    return response, history + [(message, response)]

def generate_case_outcome(prosecutor_response, defense_response):
    prompt = f"Prosecutor's argument: {prosecutor_response}\nDefense Attorney's argument: {defense_response}\nBased on verified sources, provide the case details and give the outcome along with reasons."
    evaluation = ""
    for message in client.chat_completion(
        [{"role": "system", "content": "Analyze the case and provide the outcome based on verified sources."},
         {"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
    return evaluation

def determine_winner(outcome):
    if "Prosecutor" in outcome and "Defense" in outcome:
        if outcome.count("Prosecutor") > outcome.count("Defense"):
            return "Prosecutor Wins"
        else:
            return "Defense Wins"
    elif "Prosecutor" in outcome:
        return "Prosecutor Wins"
    elif "Defense" in outcome:
        return "Defense Wins"
    else:
        return "No clear winner"

def chat_between_bots(system_message1, system_message2, max_tokens, temperature, top_p, history1, history2, shared_history, message):
    prosecutor_response, history1 = respond(message, history1, system_message1, max_tokens, temperature, top_p)
    defense_response, history2 = respond(message, history2, system_message2, max_tokens, temperature, top_p)
    shared_history.append(f"Prosecutor: {prosecutor_response}")
    shared_history.append(f"Defense Attorney: {defense_response}")
    
    outcome = generate_case_outcome(prosecutor_response, defense_response)
    winner = determine_winner(outcome)
    
    return prosecutor_response, defense_response, history1, history2, shared_history, outcome, winner

def extract_text_from_pdf(pdf_file):
    text = ""
    doc = fitz.open(pdf_file)
    for page in doc:
        text += page.get_text()
    return text

def ask_about_pdf(pdf_text, question):
    result = qa_pipeline(question=question, context=pdf_text)
    return result['answer']

def update_pdf_gallery_and_extract_text(pdf_files):
    if len(pdf_files) > 0:
        pdf_text = extract_text_from_pdf(pdf_files[0].name)
    else:
        pdf_text = ""
    return pdf_files, pdf_text

def get_top_10_cases():
    url = "https://www.courtlistener.com/?order_by=dateFiled+desc"
    response = requests.get(url)
    soup = BeautifulSoup(response.content, 'html.parser')
    
    cases = []
    for item in soup.select('.search-result', limit=10):
        case_name = item.select_one('.search-result-title a').text.strip()
        case_number = item.select_one('.search-result-meta').text.strip().split()[-1]
        cases.append(f"{case_name} - Case Number: {case_number}")
    
    return "\n".join(cases)

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=False)

def bot(history):
    system_message = "You are a helpful assistant."
    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]})
    response = ""
    for message in client.chat_completion(
        messages,
        max_tokens=150,
        stream=True,
        temperature=0.6,
        top_p=0.95,
    ):
        token = message.choices[0].delta.content
        if token is not None:
            response += token
        history[-1][1] = response
    return history

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):
    result = qa_pipeline(question=question, context=shared_history)
    return result['answer']

# Custom CSS for a clean layout
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;
}
"""

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

    with gr.Tab("Argument Evaluation"):
        gr.Markdown("# Argument Evaluation", elem_classes=["gr-title"])
        gr.Markdown("## Prosecutor vs. Defense Attorney", elem_classes=["gr-subtitle"])
        with gr.Row():
            with gr.Column(scale=1):
                top_10_btn = gr.Button("Give me the top 10 cases")
                top_10_output = gr.Markdown(elem_classes=["scroll-box"])
                top_10_btn.click(get_top_10_cases, outputs=top_10_output)
            with gr.Column(scale=2):
                message = gr.Textbox(label="Enter Case Details to Argue", placeholder="Enter case details here...")
                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.6)
                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()
                
                shared_argument = gr.Textbox(label="Case Outcome", interactive=True, elem_classes=["scroll-box"])
                winner = gr.Textbox(label="Winner", interactive=False, 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_argument, winner])
                clear_btn.click(reset_conversation, outputs=[history1, history2, shared_history, prosecutor_response, defense_response, shared_argument, winner])
                save_btn.click(save_conversation, inputs=[history1, history2, shared_history], outputs=[history1, history2, shared_history])
                
                # Inner HTML for asking about the case outcome
                with gr.Row():
                    case_question = gr.Textbox(label="Ask a Question about the Case Outcome", placeholder="Enter your question here...")
                    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)
    
    with gr.Tab("PDF Management"):
        gr.Markdown("# PDF Management", elem_classes=["gr-title"])
        pdf_upload = gr.File(label="Upload Case Files (PDF)", file_types=[".pdf"])
        pdf_gallery = gr.Gallery(label="PDF Gallery")
        pdf_view = gr.Textbox(label="PDF Content", interactive=False, elem_classes=["scroll-box"])
        pdf_question = gr.Textbox(label="Ask a Question about the PDF", placeholder="Enter your question here...")
        pdf_answer = gr.Textbox(label="Answer", interactive=False, elem_classes=["scroll-box"])
        pdf_upload_btn = gr.Button("Update PDF Gallery")
        pdf_ask_btn = gr.Button("Ask")

        pdf_upload_btn.click(update_pdf_gallery_and_extract_text, inputs=[pdf_upload], outputs=[pdf_gallery, pdf_text])
        pdf_text.change(fn=lambda x: x, inputs=pdf_text, outputs=pdf_view)
        pdf_ask_btn.click(ask_about_pdf, inputs=[pdf_text, pdf_question], outputs=pdf_answer)
    
    with gr.Tab("Chatbot"):
        gr.Markdown("# Chatbot", elem_classes=["gr-title"])
        chatbot = gr.Chatbot(
            [],
            elem_id="chatbot",
            bubble_full_width=False
        )

        chat_input = gr.MultimodalTextbox(interactive=True, file_types=["image"], placeholder="Enter message or upload file...", show_label=False)

        chat_msg = chat_input.submit(add_message, [chatbot, chat_input], [chatbot, chat_input])
        bot_msg = chat_msg.then(bot, chatbot, chatbot, api_name="bot_response")
        bot_msg.then(lambda: gr.MultimodalTextbox(interactive=True), None, [chat_input])

        chatbot.like(print_like_dislike, None, None)

demo.queue()
demo.launch()