"""Check reference navigation and parse examples with the release compiler."""

from html.parser import HTMLParser
import subprocess
import unittest

from check_reference import REFERENCE, ROOT, ReferenceParser


OPT = ROOT / "build_release/Release/bin/onnx-mlir-opt"


class IntroductionParser(HTMLParser):
    def __init__(self, document):
        super().__init__()
        self.examples = {}
        self.current = None
        self.contents = []
        self.in_contents = False
        self.feed(document)

    def handle_starttag(self, tag, attrs):
        attrs = dict(attrs)
        if tag == "pre" and "data-example" in attrs:
            self.current = attrs["data-example"]
            self.examples[self.current] = ""
        if tag == "nav":
            self.in_contents = attrs.get("aria-label") == "Table of contents"
        if self.in_contents and tag == "a":
            self.contents.append(attrs["href"])

    def handle_data(self, data):
        if self.current:
            self.examples[self.current] += data

    def handle_endtag(self, tag):
        if tag == "pre":
            self.current = None
        if tag == "nav":
            self.in_contents = False


class IntroductionTest(unittest.TestCase):
    @classmethod
    def setUpClass(cls):
        cls.document = REFERENCE.read_text()
        cls.introduction = IntroductionParser(cls.document)

    def test_contents_covers_topics_and_indexes(self):
        self.assertLess(
            self.document.index('aria-label="Table of contents"'),
            self.document.index('<section id="reading">'),
        )
        parser = ReferenceParser()
        parser.feed(self.document)
        for target in self.introduction.contents:
            self.assertIn(target[1:], parser.ids)
        self.assertEqual(set(self.introduction.contents), {
            "#reading", "#compact", "#planning", "#compute-regions",
            "#graph-batches", "#input-representations", "#scheduling",
            "#destinations", "#memory-units", "#communication", "#contracts",
            "#spatial-index", "#pim-index",
        })

    @unittest.skipUnless(OPT.is_file(), "Release onnx-mlir-opt is unavailable")
    def test_examples_parse_and_print(self):
        arguments = {
            "addition": "%a: tensor<1x4xf32>, %b: tensor<1x4xf32>",
            "rows": "%matrix: tensor<2x4xf32>",
            "plan": "%a: tensor<1x2x3x4xf32>, %b: tensor<1x2x3x4xf32>",
            "weighted": "%x: tensor<1x4xf32>",
            "graph-batch": "%x: tensor<1x8xf32>",
            "reconstruction": "%fragments: tensor<2x1x4xf32>",
            "core-batch": None,
            "destination": "%x: tensor<1x4xf32>, %dst: tensor<1x4xf32>",
            "bufferized": "%x: memref<1x4xf32>, %dst: memref<1x4xf32>",
            "copy": "%dst: memref<1x4xf32>, %src: memref<1x4xf32>",
            "gather": "%dst: memref<1x4xf32>, %src: memref<1x8xf32>",
            "send": "%x: tensor<1x4xf32>, %target: index",
            "receive": "%dst: tensor<1x4xf32>, %source: index",
            "signal": "%target: index, %event: index",
            "wait": "%event: index, %count: index",
            "deferred-group": "%a: tensor<1x4xf32>, %b: tensor<1x4xf32>",
        }
        self.assertEqual(set(self.introduction.examples), set(arguments))
        for name, args in arguments.items():
            with self.subTest(example=name):
                snippet = self.introduction.examples[name]
                if name in {"gather", "receive"}:
                    snippet = "%zero = arith.constant 0 : index\n" + snippet
                if args is not None:
                    snippet = f"func.func @example({args}) {{\n{snippet}\nreturn\n}}"
                program = f"module {{\n{snippet}\n}}\n"
                command = [str(OPT), "--maccel=PIM"]
                if name == "deferred-group":
                    command.append("--mlir-print-op-generic")
                for _ in range(2):
                    result = subprocess.run(
                        command, input=program,
                        text=True, capture_output=True, timeout=30,
                    )
                    self.assertEqual(result.returncode, 0, result.stderr)
                    program = result.stdout


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