from datetime import datetime

import gradio as gr
import pandas as pd
from datasets import load_dataset

# Load the dataset
ds = load_dataset("harshildarji/openlegaldata", "cases", split="main")
df = pd.DataFrame(ds)

# Precompute additional columns for filtering
df["state"] = df["court"].apply(
    lambda x: x.get("state", "Unknown") if isinstance(x, dict) else "Unknown"
)
df["court_name"] = df["court"].apply(
    lambda x: x.get("name", "Unknown") if isinstance(x, dict) else "Unknown"
)

# Build unique lists for the filters
state_list = sorted(df["state"].dropna().unique().tolist())
state_list.insert(0, "All")
court_list = sorted(df["court_name"].dropna().unique().tolist())
court_list.insert(0, "All")


def filter_cases(state, court, date_from, date_to):
    """
    Filter the cases based on state, court name, and date range.
    Returns a DataFrame with selected columns.
    """
    filtered = df.copy()
    if state and state != "All":
        filtered = filtered[filtered["state"] == state]
    if court and court != "All":
        filtered = filtered[filtered["court_name"] == court]
    if date_from:
        try:
            date_from_dt = datetime.strptime(date_from, "%Y-%m-%d")
            filtered = filtered[pd.to_datetime(filtered["date"]) >= date_from_dt]
        except Exception:
            pass
    if date_to:
        try:
            date_to_dt = datetime.strptime(date_to, "%Y-%m-%d")
            filtered = filtered[pd.to_datetime(filtered["date"]) <= date_to_dt]
        except Exception:
            pass

    return filtered[["id", "file_number", "date", "court_name", "type"]]


def get_case_details(case_id):
    """Return an HTML formatted string with details for a given case id."""
    try:
        case_id = int(case_id)
    except:
        return "<p style='color:red;'>Invalid case ID</p>"
    case_row = df[df["id"] == case_id]
    if case_row.empty:
        return "<p style='color:red;'>Case not found</p>"
    case_data = case_row.iloc[0].to_dict()

    html = "<h2>Case Details</h2>"
    html += f"<p><strong>ID:</strong> {case_data.get('id')}</p>"
    html += f"<p><strong>File Number:</strong> {case_data.get('file_number')}</p>"
    html += f"<p><strong>Date:</strong> {case_data.get('date')}</p>"

    # Display court information
    court_info = case_data.get("court", {})
    html += f"<p><strong>Court:</strong> {court_info.get('name', 'N/A')} ({court_info.get('state', 'N/A')})</p>"
    html += f"<p><strong>Type:</strong> {case_data.get('type', 'N/A')}</p>"
    html += '<hr style="margin: 15px 0;">'

    # Display tenor if available
    tenors = case_data.get("tenor", [])
    if tenors:
        html += "<h3>Tenor</h3>"
        for t in tenors:
            html += f"<p>{t}</p>"
        html += '<hr style="margin: 15px 0;">'

    # Display tatbestand if available
    tatbestand = case_data.get("tatbestand", [])
    if tatbestand:
        html += "<h3>Tatbestand</h3>"
        for t in tatbestand:
            html += f"<p>{t}</p>"
        html += '<hr style="margin: 15px 0;">'

    # Display Gründe if available
    gründe = case_data.get("gründe", [])
    if gründe:
        html += "<h3>Gründe</h3>"
        for g in gründe:
            html += f"<p>{g}</p>"
        html += '<hr style="margin: 15px 0;">'

    # Display Entscheidungsgründe if available
    entscheidungsgründe = case_data.get("entscheidungsgründe", [])
    if entscheidungsgründe:
        html += "<h3>Entscheidungsgründe</h3>"
        for e in entscheidungsgründe:
            html += f"<p>{e}</p>"
        html += '<hr style="margin: 15px 0;">'

    return html


# Build multi-tabs interface
with gr.Blocks(title="German Legal Case Viewer", fill_width=True) as demo:
    gr.Markdown("# German Legal Case Viewer")
    gr.Markdown(
        "Explore case information from the processed [Open Legal Data dataset](https://huggingface.co/datasets/harshildarji/openlegaldata)."
    )

    with gr.Tabs():
        # First Tab - Browse and filter cases
        with gr.Tab("Browse Cases"):
            gr.Markdown("## Filter Cases")
            with gr.Row():
                state_input = gr.Dropdown(
                    choices=state_list, label="State", value="All"
                )
                court_input = gr.Dropdown(
                    choices=court_list, label="Court", value="All"
                )
                date_from_input = gr.Textbox(
                    label="From Date (YYYY-MM-DD)", placeholder="2022-01-01"
                )
                date_to_input = gr.Textbox(
                    label="To Date (YYYY-MM-DD)", placeholder="2022-12-31"
                )
            filter_button = gr.Button("Apply Filters")
            output_table = gr.DataFrame(label="Cases Overview")
            filter_button.click(
                fn=filter_cases,
                inputs=[state_input, court_input, date_from_input, date_to_input],
                outputs=output_table,
            )

        # Second Tab - Show details of a specific case
        with gr.Tab("Case Details"):
            with gr.Row():
                case_id_input = gr.Textbox(
                    label="Enter Case ID", placeholder="e.g., 346915"
                )
                details_button = gr.Button("Get Details")
            case_details_output = gr.HTML(label="Case Details")
            details_button.click(
                fn=get_case_details,
                inputs=case_id_input,
                outputs=case_details_output,
            )

demo.launch()