"""Display coding costs without changing the arithmetic stream."""
import html
import math


def token_cost(cdf, token):
    probability = int(cdf[token + 1] - cdf[token]) / int(cdf[-1])
    return {'probability': probability, 'bits': -math.log2(probability)}


def text_spans(text, ids, costs, pieces):
    """Join byte-split Unicode tokens until they form valid display text."""
    if not len(ids) == len(costs) == len(pieces):
        raise ValueError('Token highlight lengths differ.')
    if b''.join(pieces) != text.encode('utf-8'):
        raise ValueError('Token highlights do not match the original text.')
    spans, pending, tokens = [], b'', []
    for token, cost, piece in zip(ids, costs, pieces):
        pending += piece
        tokens.append({'id': int(token), **cost})
        try:
            display = pending.decode('utf-8')
        except UnicodeDecodeError:
            continue
        if not display:
            continue
        spans.append({'text': display, 'tokens': tokens})
        pending, tokens = b'', []
    if pending or tokens:
        raise ValueError('Incomplete token highlight text.')
    return spans


# Fixed 2-bit bins make the same cost have the same color across inputs.
COLORS = ['#eef3ec', '#dcebdf', '#bcdcc6', '#8fc3a5', '#388060', '#266a4e', '#13432f']


def percent(probability):
    value = 100 * probability
    if value >= 99.95:
        return '>99.9%'
    if value >= 10:
        return f'{value:.0f}%'
    if value >= 0.01:
        return f'{value:.2g}%'
    return '<0.01%'


def describe(tokens):
    bits = sum(token['bits'] for token in tokens)
    if len(tokens) == 1:
        return f'{bits:.2f} bits · p {percent(tokens[0]["probability"])}'
    return f'{bits:.2f} bits · {len(tokens)} byte tokens'


def render_highlight(stats):
    if not stats or 'token_spans' not in stats:
        return ''
    parts = []
    for index, span in enumerate(stats['token_spans']):
        tokens = span['tokens']
        level = min(6, int(sum(token['bits'] for token in tokens) / len(tokens) / 2))
        foreground = '#ffffff' if level >= 4 else '#1d201c'
        # Leading spaces sit outside the shading so words read naturally; the tooltip restores them.
        text = span['text']
        body = text.lstrip(' ')
        if not body.strip():
            body = text
        lead = text[:len(text) - len(body)]
        # Line breaks have no width, so a leading one gets a visible mark.
        newline = ' data-newline' if body.startswith('\n') else ''
        space = f' data-space="{len(lead)}"' if lead else ''
        parts.append(
            f'{lead}<span class="cost-token" tabindex="{0 if index == 0 else -1}"{newline}{space} '
            f'data-tip="{html.escape(describe(tokens), quote=True)}" '
            f'style="background:{COLORS[level]};color:{foreground}">'
            f'{html.escape(body)}</span>'
        )
    swatches = ''.join(f'<i style="background:{color}"></i>' for color in COLORS)
    body = ''.join(parts) or '<span class="quiet">Empty input has no tokens.</span>'
    return (
        '<section class="analysis"><div class="analysis-heading"><h2>Token costs</h2>'
        f'<div class="legend" aria-hidden="true"><span>0</span>{swatches}<span>12+ bits</span></div></div>'
        '<p class="analysis-hint">Darker tokens took more bits to encode. Hover or tap one to inspect it.</p>'
        f'<div class="token-text" aria-label="Original text with token costs">{body}</div>'
        '<details class="fine-print"><summary>How to read this</summary>'
        '<p>Cost is −log₂ p, using the exact integer probabilities sent to the arithmetic coder. '
        'These are ideal token costs; the copied code also includes framing and ASCII overhead. '
        'A character split across several byte tokens is shaded by their average cost.</p></details></section>'
    )