"""Gradio app showing the original rechtspraak texts and the GPT-NL versions side by side."""
import difflib
import html
import random
import re
from collections import defaultdict

import gradio as gr
from datasets import load_dataset

N = 1000
GPTNL_REPO = 'GPT-NL/GPT-NL_Public_Corpus'
DIRECT_REPO = 'sheggle/dutch-legal-documents'

# Stream instead of downloading the full datasets (~10 GB of parquet).
# Copies of an ECLI are adjacent (max 3), so read 2 extra rows to catch spill-over past row N
gptnl = list(load_dataset(GPTNL_REPO, 'rechtspraak', split='train', streaming=True).take(N + 2))
wanted = {doc['source'] for doc in gptnl[:N]}

# Only fetch the source rows we need, filtered on the parquet side
direct = load_dataset(DIRECT_REPO, split='train', streaming=True,
                      filters=[('identifier', 'in', sorted(wanted))])
direct_docs = {doc['identifier']: doc for doc in direct}

# One ECLI can occur in several (adjacent) GPT-NL rows: collect all of them
gptnl_rows = defaultdict(list)
for i, doc in enumerate(gptnl):
    if doc['source'] in wanted and doc['source'] in direct_docs:
        gptnl_rows[doc['source']].append((i, doc))
eclis = list(gptnl_rows)

STYLE = """
<style>
.diff-grid { display: grid; grid-template-columns: 1fr 1fr; gap: 16px; }
.diff-stack { max-height: 75vh; overflow-y: auto; }
.diff-stack .diff-col { max-height: none; margin-bottom: 12px; }
.diff-grid h3, .diff-grid h4 { margin: 8px 0 4px; white-space: nowrap; overflow: hidden; text-overflow: ellipsis; }
.diff-grid h3 small, .diff-grid h4 small { font-weight: normal; opacity: 0.75; }
.diff-col { white-space: pre-wrap; font-family: monospace; font-size: 13px; line-height: 1.5;
            border: 1px solid #8884; border-radius: 6px; padding: 10px; max-height: 75vh; overflow-y: auto; }
.diff-col del { background: #f8514940; text-decoration: none; }
.diff-col ins { background: #2ea04340; text-decoration: none; }
</style>
"""


def diff(left, right):
    """Word-level diff; returns (left html, right html, similarity)."""
    # Tokens are words plus their trailing whitespace, so newlines survive; compare on the words only
    lt, rt = re.findall(r'\S+\s*', left), re.findall(r'\S+\s*', right)
    sm = difflib.SequenceMatcher(None, [t.rstrip() for t in lt], [t.rstrip() for t in rt])
    l_html, r_html = [], []
    for tag, i1, i2, j1, j2 in sm.get_opcodes():
        l_chunk = html.escape(''.join(lt[i1:i2]))
        r_chunk = html.escape(''.join(rt[j1:j2]))
        if tag == 'equal':
            l_html.append(l_chunk)
            r_html.append(r_chunk)
        else:
            if l_chunk:
                l_html.append(f'<del>{l_chunk}</del>')
            if r_chunk:
                r_html.append(f'<ins>{r_chunk}</ins>')
    return ''.join(l_html), ''.join(r_html), sm.ratio()


def ecli_link(ecli):
    url = f'https://uitspraken.rechtspraak.nl/details?id={ecli}'
    return f'<a href="{url}" target="_blank" rel="noopener">{html.escape(ecli)}</a>'


def repo_link(repo):
    return f'<a href="https://huggingface.co/datasets/{repo}" target="_blank" rel="noopener">{repo}</a>'


def header(tag, title, *meta):
    """One-line header; meta items are (label, html) pairs."""
    meta = ' · '.join(f'{label} <code>{value}</code>' if label else value for label, value in meta)
    return f'<{tag}>{title} <small>{meta}</small></{tag}>'


def render(ecli):
    d = direct_docs[ecli]
    left = d['full_text']
    # Diff every GPT-NL row against the original; the left side shows the diff with the best row
    results = [(i, doc, *diff(left, doc['text'])) for i, doc in gptnl_rows[ecli]]
    best = max(results, key=lambda r: r[4])
    stats = f'<p><b>{ecli_link(ecli)}</b> — {len(results)} GPT-NL row(s)</p>'
    left_header = header('h3', repo_link(DIRECT_REPO),
                         ('id', html.escape(str(d['id']))),
                         ('', f'{len(left.split()):,} words'), ('', f'diff vs. row {best[0]}'))
    right_html = ''.join(
        header('h4', f'Row {i}',
               ('id', html.escape(doc['id'])),
               ('', f'{len(doc["text"].split()):,} words'), ('', f'similarity {sim:.3f}'))
        + f'<div class="diff-col">{r_html}</div>'
        # Shortest rows first, so the partial copies sit above the full document
        for i, doc, _, r_html, sim in sorted(results, key=lambda r: len(r[1]['text'])))
    body = (f'{STYLE}<div class="diff-grid">'
            f'<div>{left_header}<div class="diff-col">{best[2]}</div></div>'
            f'<div><h3>{repo_link(GPTNL_REPO)} <small>(rechtspraak)</small></h3><div class="diff-stack">{right_html}</div></div></div>')
    return stats, body


def random_doc():
    ecli = random.choice(eclis)
    return ecli, *render(ecli)


with gr.Blocks(title='Rechtspraak diff') as demo:
    with gr.Row():
        dropdown = gr.Dropdown(eclis, value=eclis[0], label='ECLI', scale=4)
        button = gr.Button('Random Document', scale=1)
    stats = gr.HTML()
    diff_view = gr.HTML()

    dropdown.input(render, dropdown, [stats, diff_view])
    button.click(random_doc, None, [dropdown, stats, diff_view])
    demo.load(render, dropdown, [stats, diff_view])

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