import math

LEGEND_HTML = """
<div style='margin:10px 0; font-family:monospace;'>
  <b>Confidence:</b>&nbsp;
  <span style='background:#ff0000; color:white; padding:4px 8px; border-radius:3px;'>Low (0%)</span>&nbsp;
  <span style='background:#7f7f00; color:white; padding:4px 8px; border-radius:3px;'>Moderate (50%)</span>&nbsp;
  <span style='background:#00ff00; color:black; padding:4px 8px; border-radius:3px;'>High (100%)</span>
  &nbsp;&nbsp;<em style='font-size:0.85em;'>Hover over tokens/phrases to see values.</em>
</div>
"""

TOOLTIP_CSS = """
span[data-tooltip] { position: relative; }
span[data-tooltip]:hover::after {
    content: attr(data-tooltip);
    position: absolute;
    bottom: 130%;
    left: 50%;
    transform: translateX(-50%);
    background: #333;
    color: white;
    padding: 4px 8px;
    border-radius: 4px;
    font-size: 0.75em;
    white-space: nowrap;
    z-index: 9999;
    pointer-events: none;
}
"""


def hex_to_rgb(hex_code):
    r = int(hex_code[0:2], 16)
    g = int(hex_code[2:4], 16)
    b = int(hex_code[4:6], 16)
    return r, g, b


def prob_to_color(avg_prob):
    p = max(0.0, min(1.0, float(avg_prob)))
    r = int((1 - p) * 255)
    g = int(p * 255)
    return f"{r:02x}{g:02x}00"


def _span(text, color_hex, tooltip):
    r, g, b = hex_to_rgb(color_hex)
    brightness = (r * 299 + g * 587 + b * 114) / 1000
    text_color = "black" if brightness > 128 else "white"
    safe_text = text.replace("&", "&amp;").replace("<", "&lt;").replace(">", "&gt;").replace(" ", "&nbsp;")
    return (
        f'<span data-tooltip="{tooltip}" '
        f'style="background-color:#{color_hex}; color:{text_color}; '
        f'padding:2px 4px; border-radius:3px; margin:1px; cursor:pointer;">'
        f'{safe_text}</span>'
    )


def _wrap(inner_html):
    return f"<p style='font-family: monospace; font-size: 1.1em; line-height: 2.2em;'>{inner_html}</p>"


def _token_span(text, color_hex, tooltip):
    """Inline highlight (no pill padding/margin) so tokens flow as continuous text."""
    r, g, b = hex_to_rgb(color_hex)
    brightness = (r * 299 + g * 587 + b * 114) / 1000
    text_color = "black" if brightness > 128 else "white"
    safe_text = text.replace("&", "&amp;").replace("<", "&lt;").replace(">", "&gt;")
    return (
        f'<span data-tooltip="{tooltip}" '
        f'style="background-color:#{color_hex}; color:{text_color};">'
        f'{safe_text}</span>'
    )


def _wrap_flow(inner_html):
    """Box matching the semantic-entropy view: words flow and wrap inside it."""
    return (
        "<div style='font-family: ui-sans-serif, system-ui, -apple-system, sans-serif; "
        "font-size: 1.05em; line-height: 2em; white-space: pre-wrap; "
        "overflow-wrap: anywhere; word-break: normal;'>"
        f"{inner_html}</div>"
    )


def render_token_probs_html(token_logprobs):
    """token_logprobs: list of {"token": str, "logprob": float}"""
    if not token_logprobs:
        return "<p>No token probability data available.</p>"
    spans = ""
    for item in token_logprobs:
        prob = math.exp(item["logprob"])
        color = prob_to_color(prob)
        tooltip = f"p={prob:.3f} logprob={item['logprob']:.3f}"
        spans += _token_span(item["token"], color, tooltip)
    return _wrap_flow(spans)


def color_code_answer_html(answer_text, sp_to_entropy, sentence_parts):
    """Render answer with sentence-part spans colored by semantic-entropy confidence."""
    if not sentence_parts or not sp_to_entropy or not answer_text:
        safe = answer_text.replace("&", "&amp;").replace("<", "&lt;").replace(">", "&gt;")
        return f"<p>{safe}</p>"

    # Find each sentence_part as a substring and record its position
    colored_parts = []
    for sp in sentence_parts:
        if not sp:
            continue
        idx = answer_text.find(sp)
        if idx == -1:
            continue
        entropy = sp_to_entropy.get(sp, 0.0)
        confidence = math.exp(-entropy)
        color = prob_to_color(confidence)
        tooltip = f"entropy={entropy:.3f}  confidence={confidence:.3f}"
        colored_parts.append((idx, idx + len(sp), sp, color, tooltip))

    colored_parts.sort(key=lambda x: x[0])

    html = ""
    pos = 0
    for start, end, text, color, tooltip in colored_parts:
        if start > pos:
            gap = answer_text[pos:start]
            html += gap.replace("&", "&amp;").replace("<", "&lt;").replace(">", "&gt;")
        html += _span(text, color, tooltip)
        pos = end
    if pos < len(answer_text):
        tail = answer_text[pos:]
        html += tail.replace("&", "&amp;").replace("<", "&lt;").replace(">", "&gt;")

    return _wrap(html)