"""A quiet interface around the existing compression API."""
import html
from pathlib import Path

import gradio as gr
from highlight import render_highlight

CSS = Path(__file__).with_name('style.css').read_text()
EXAMPLE = ('Language models predict what comes next. When a model predicts a word well, '
           'we can describe that word using fewer bits. Arithmetic coding turns these '
           'predictions into a compact string, and the same model can recover the exact original text.')
IDLE = '<span class="quiet">Up to 8,192 tokens or 64 KiB</span>'
MODEL_URL = 'https://huggingface.co/Qwen/Qwen3.5-35B-A3B-Base'
GGUF_URL = 'https://huggingface.co/mradermacher/Qwen3.5-35B-A3B-Base-i1-GGUF'

HEADER = f'''<header class="page-header">
<h1>Text compressor</h1>
<a class="model-label" href="{MODEL_URL}" target="_blank" rel="noreferrer">Qwen3.5-35B-A3B Base · Q4_K_M</a>
<p class="intro">Lossless compression with a language model and arithmetic coding.</p>
</header>'''

FOOTER = f'''<footer class="page-footer"><details><summary>About</summary>
<p>Qwen3.5-35B-A3B Base, run from the published
<a href="{GGUF_URL}" target="_blank" rel="noreferrer">community i1-Q4_K_M GGUF</a>, predicts each next token.
Arithmetic coding turns those probabilities into bits, written out as basE91 ASCII.
Restoring replays the same predictions, so it needs this model and a compatible runtime.
Use ⇄ to switch to Restore, then paste a code to decode it and verify its checksum.</p>
<p>Both directions use ZeroGPU quota, and the first run may need to load the model.
This is an experiment, not an archival format; keep your original text.</p>
</details></footer>'''

TOKEN_JS = '''
const shown = s => s.replace(/ /g, '\\u2423').replace(/\\n/g, '\\u21b5').replace(/\\t/g, '\\u21e5');
let active = null;

function tip() {
    const section = element.querySelector('.analysis');
    if (!section) return null;
    let node = section.querySelector('.token-tip');
    if (!node) {
        node = document.createElement('div');
        node.className = 'token-tip';
        node.setAttribute('role', 'status');
        section.appendChild(node);
    }
    return node;
}

function show(token) {
    const node = tip();
    if (!token || !node) return;
    if (active && active !== token) active.classList.remove('selected');
    active = token;
    element.querySelectorAll('.cost-token').forEach(t => { t.tabIndex = t === token ? 0 : -1; });
    token.classList.add('selected');
    const text = document.createElement('code');
    text.textContent = shown(' '.repeat(Number(token.dataset.space || 0)) + token.textContent);
    const cost = document.createElement('span');
    cost.textContent = token.dataset.tip;
    node.replaceChildren(text, cost);
    node.classList.add('visible');
    // Center above the token's first line box; flip below when it would leave the viewport.
    const box = node.parentElement.getBoundingClientRect();
    const rect = token.getClientRects()[0] || token.getBoundingClientRect();
    const left = rect.left - box.left + rect.width / 2 - node.offsetWidth / 2;
    const above = rect.top - node.offsetHeight - 8;
    node.style.left = Math.max(0, Math.min(left, box.width - node.offsetWidth)) + 'px';
    node.style.top = (above < 0 ? rect.bottom + 8 : above) - box.top + 'px';
}

function hide() {
    if (active) active.classList.remove('selected');
    active = null;
    const node = element.querySelector('.token-tip');
    if (node) node.classList.remove('visible');
}

element.addEventListener('pointerover', e => {
    const token = e.target.closest('.cost-token');
    if (token) show(token);
});
element.addEventListener('pointerout', e => {
    const text = e.target.closest('.token-text');
    if (text && !text.contains(e.relatedTarget) && !text.contains(document.activeElement)) hide();
});
element.addEventListener('click', e => {
    const token = e.target.closest('.cost-token');
    token ? show(token) : hide();
});
element.addEventListener('focusin', e => show(e.target.closest('.cost-token')));
element.addEventListener('focusout', e => {
    if (!element.contains(e.relatedTarget)) hide();
});
element.addEventListener('scroll', hide, true);
element.addEventListener('keydown', e => {
    const token = e.target.closest('.cost-token');
    if (!token) return;
    if (e.key === 'Escape') { hide(); return; }
    if (!['ArrowLeft', 'ArrowRight'].includes(e.key)) return;
    const next = e.key === 'ArrowRight' ? token.nextElementSibling : token.previousElementSibling;
    if (next && next.classList.contains('cost-token')) {
        e.preventDefault(); next.focus();
    }
});
'''


def theme():
    base = gr.themes.Base(
        primary_hue='neutral', secondary_hue='neutral', neutral_hue='neutral',
        font=[gr.themes.GoogleFont('Inter', weights=(400, 500, 600)), 'ui-sans-serif', 'system-ui', 'sans-serif'],
        font_mono=[gr.themes.GoogleFont('JetBrains Mono', weights=(400, 500)), 'ui-monospace', 'monospace'],
        radius_size='sm',
    )
    values = dict(
        body_background_fill='#fafaf8', body_text_color='#1d201c',
        body_text_color_subdued='#5d645a', background_fill_primary='#ffffff',
        background_fill_secondary='#f3f4f0', border_color_primary='#dcdfd7',
        block_background_fill='transparent', block_border_width='0px',
        block_shadow='none', block_label_background_fill='transparent',
        block_label_text_color='#5d645a', block_label_border_width='0px',
        block_label_text_size='13px', block_title_text_weight='500',
        input_background_fill='#ffffff', input_border_color='#dcdfd7',
        input_border_width='1px', input_shadow='none', input_shadow_focus='none',
        input_border_color_focus='#2f5d45', input_text_size='15px',
        button_primary_background_fill='#1d201c', button_primary_background_fill_hover='#343a32',
        button_primary_border_color='#1d201c', button_primary_text_color='#ffffff',
        button_secondary_background_fill='transparent', button_secondary_text_color='#1d201c',
        button_secondary_border_color='#dcdfd7',
        button_primary_shadow='none', button_primary_shadow_hover='none', button_primary_shadow_active='none',
        button_secondary_shadow='none', button_secondary_shadow_hover='none', button_secondary_shadow_active='none',
        button_large_text_size='14px', button_large_text_weight='500',
        loader_color='#2f5d45', color_accent='#2f5d45', color_accent_soft='#e6efe8',
        layout_gap='16px',
    )
    return base.set(**values, **{key+'_dark': value for key, value in values.items()
                               if hasattr(base, key+'_dark')})


def message(exc):
    # ZeroGPU's HTML errors render their plain-text variant through str().
    return f'<span class="error">{html.escape(str(getattr(exc, "message", exc)))}</span>'


def figures(lead, *rest):
    # Items never break internally; lines wrap after a separator.
    items = [f'<span class="figure">{lead}</span>', *(f'<span class="stat">{item}</span>' for item in rest)]
    return '<span class="sep">&nbsp;· </span>'.join(items)


def compressed_status(stats):
    original, code = stats['original_utf8_bytes'], stats['ascii_characters_and_bytes']
    share = f'{100*code/original:.1f}% of original' if original else 'empty input'
    return figures(f'{original:,} → {code:,} bytes', share, f'{stats["tokens"]:,} tokens',
                   f'{stats["encode_seconds"]:.2f} s')


def restored_status(stats):
    return figures('Checksum verified',
                   f'{stats["original_utf8_bytes"]:,} bytes restored in {stats["decode_seconds"]:.2f} s')


def swap_mode(mode, source, result, drafts):
    drafts = {**drafts, mode: source}
    mode = 'restore' if mode == 'compress' else 'compress'
    # Carry the result across, or restore this mode's unsubmitted draft.
    next_input = result if result else drafts.get(mode, '')
    restoring = mode == 'restore'
    return (
        mode, drafts,
        gr.update(value=next_input, label='Compressed' if restoring else 'Original text',
                  placeholder='Paste a code from this compressor' if restoring else 'Paste text to compress',
                  elem_classes=['ascii-panel'] if restoring else []),
        gr.update(value='', label='Restored text' if restoring else 'Compressed',
                  placeholder='The original text appears here' if restoring else 'The compressed code appears here',
                  elem_classes=['result-box'] if restoring else ['result-box', 'ascii-panel']),
        gr.update(visible=not restoring), gr.update(visible=restoring),
        IDLE, '', {},
    )


def build_demo(compress, decompress):
    def run_compress(text, progress=gr.Progress()):
        try:
            code, _, stats = compress(text, progress)
        except (gr.Error, ValueError, RuntimeError) as exc:
            return '', message(exc), {}
        return code, compressed_status(stats), stats

    def run_decompress(code, progress=gr.Progress()):
        try:
            text, _, stats = decompress(code, progress)
        except (gr.Error, ValueError, RuntimeError) as exc:
            return '', message(exc), {}
        return text, restored_status(stats), stats

    with gr.Blocks(title='Text compressor', analytics_enabled=False) as demo:
        mode = gr.State('compress')
        drafts = gr.State({'compress': EXAMPLE, 'restore': ''})
        gr.HTML(HEADER, apply_default_css=False)
        with gr.Row(elem_classes='converter-row'):
            with gr.Column(min_width=0, elem_classes='editor-panel'):
                source = gr.Textbox(label='Original text', value=EXAMPLE, lines=10, max_lines=16,
                                    placeholder='Paste text to compress', elem_id='input-panel',
                                    html_attributes={'spellcheck': 'false'})
            with gr.Column(scale=0, min_width=0, elem_classes='swap-column'):
                swap_button = gr.Button('Swap direction', variant='secondary', min_width=0, scale=0,
                                        elem_id='swap-mode')
            with gr.Column(min_width=0, elem_classes='editor-panel'):
                result = gr.Textbox(label='Compressed', lines=10, max_lines=16,
                                    placeholder='The compressed code appears here', interactive=False,
                                    buttons=['copy'], elem_id='output-panel',
                                    elem_classes=['result-box', 'ascii-panel'])
        with gr.Row(elem_classes='action-row'):
            encode_button = gr.Button('Compress', variant='primary', min_width=0, scale=0, elem_classes='run')
            decode_button = gr.Button('Restore', variant='primary', min_width=0, scale=0, elem_classes='run',
                                      visible=False)
            status = gr.HTML(IDLE, elem_id='result-status', apply_default_css=False)
        stats = gr.JSON(visible=False)
        heatmap = gr.HTML('', elem_id='heatmap', js_on_load=TOKEN_JS, apply_default_css=False)
        gr.HTML(FOOTER, apply_default_css=False)

        swap_button.click(
            swap_mode, [mode, source, result, drafts],
            [mode, drafts, source, result, encode_button, decode_button, status, heatmap, stats],
            queue=False, api_visibility='private', show_progress='hidden')

        controls = [source, swap_button, encode_button, decode_button]

        def lock():
            return ('', '', '', gr.update(interactive=False), gr.update(interactive=False),
                    gr.update(interactive=False, value='Compressing…'),
                    gr.update(interactive=False, value='Restoring…'))

        def unlock():
            return (gr.update(interactive=True), gr.update(interactive=True),
                    gr.update(interactive=True, value='Compress'),
                    gr.update(interactive=True, value='Restore'))

        for button, fn in [(encode_button, run_compress), (decode_button, run_decompress)]:
            run = button.click(
                lock, outputs=[status, heatmap, result, *controls],
                queue=False, api_visibility='private', show_progress='hidden').then(
                    fn, [source], [result, status, stats], api_visibility='private',
                    concurrency_id='gpu', concurrency_limit=1, show_progress_on=[result])
            if fn is run_compress:
                run.success(render_highlight, [stats], [heatmap], queue=False, api_visibility='private')
            run.then(unlock, outputs=controls, queue=False, api_visibility='private', show_progress='hidden')
            # Gradio skips .then() after an exception, so failure has to unlock the controls itself.
            run.failure(lambda: (message('Something went wrong. Please try again.'), *unlock()),
                        outputs=[status, *controls], queue=False, api_visibility='private', show_progress='hidden')

        # Editing invalidates the result and analysis for the previous input.
        source.input(lambda: ('', IDLE, '', {}), outputs=[result, status, heatmap, stats],
                     queue=False, api_visibility='private', show_progress='hidden')

        # The public API keeps the original endpoints, which raise errors instead of printing them.
        report = gr.Markdown(visible=False)
        for name, fn in [('compress', compress), ('decompress', decompress)]:
            gr.Button(visible=False).click(fn, [source], [result, report, stats], api_name=name,
                                           concurrency_id='gpu', concurrency_limit=1)
    return demo