import math
import unittest
from html.parser import HTMLParser

from highlight import describe, percent, token_cost, text_spans, render_highlight


class SpanText(HTMLParser):
    """Collect the rendered token text: token spans plus the spaces kept between them."""
    def __init__(self):
        super().__init__()
        self.depth = 0
        self.text = ''

    def handle_starttag(self, tag, attrs):
        if self.depth or dict(attrs).get('class') == 'token-text':
            self.depth += 1

    def handle_endtag(self, tag):
        if self.depth:
            self.depth -= 1

    def handle_data(self, data):
        if self.depth:
            self.text += data


class HighlightTests(unittest.TestCase):
    def test_exact_coding_probability(self):
        self.assertEqual(token_cost([0, 1, 4], 0), {'bits': 2.0, 'probability': .25})
        self.assertAlmostEqual(token_cost([0, 1, 4], 1)['bits'], -math.log2(.75))

    def test_byte_split_unicode_whitespace_and_html(self):
        text = '  Hi 🌍\n你好\t<img src=x onerror="bad()"> & café  '
        pieces = [bytes([b]) for b in text.encode()]
        costs = [{'bits': 2.0, 'probability': .25}] * len(pieces)
        spans = text_spans(text, list(range(len(pieces))), costs, pieces)
        self.assertEqual(''.join(s['text'] for s in spans), text)
        globe = next(s for s in spans if s['text'] == '🌍')
        self.assertEqual(len(globe['tokens']), 4)
        rendered = render_highlight({'token_spans': spans})
        self.assertNotIn('<img', rendered)
        parser = SpanText()
        parser.feed(rendered)
        self.assertEqual(parser.text, text)

    def test_token_details(self):
        self.assertEqual(describe([{'bits': 1.0, 'probability': .5}]), '1.00 bits · p 50%')
        self.assertEqual(describe([{'bits': 2.0, 'probability': .25}] * 3), '6.00 bits · 3 byte tokens')
        self.assertEqual([percent(p) for p in [.9999, .046, .0001234, 1e-7]], ['>99.9%', '4.6%', '0.012%', '<0.01%'])
        spans = [{'text': 'a', 'tokens': [{'bits': 1.0, 'probability': .5}]},
                 {'text': '\n\n', 'tokens': [{'bits': 3.0, 'probability': .125}]},
                 {'text': '  word', 'tokens': [{'bits': 3.0, 'probability': .125}]},
                 {'text': '  ', 'tokens': [{'bits': 3.0, 'probability': .125}]}]
        rendered = render_highlight({'token_spans': spans})
        self.assertEqual(rendered.count('data-newline'), 1)
        self.assertIn('\n\n</span>  <span', rendered)
        self.assertIn('data-space="2"', rendered)
        self.assertIn('>word</span><span', rendered)
        self.assertIn('>  </span>', rendered)

    def test_empty_and_mismatched_text(self):
        self.assertEqual(text_spans('', [], [], []), [])
        self.assertIn('no tokens', render_highlight({'token_spans': []}))
        with self.assertRaises(ValueError):
            text_spans('b', [1], [{'bits': 1}], [b'a'])


if __name__ == '__main__':
    unittest.main()