import unittest
from pathlib import Path

from lihont import avatar_patch

ROOT = Path(__file__).resolve().parents[1]
REAL = ROOT / "work" / "dl" / "avatar" / "X" / "1Bem.smali"

VIEW = """\
.class public LX/1Bem;
.super Lcom/bytedance/lighten/loader/SmartImageView;

.method public final onDetachedFromWindow()V
    .registers 1

    return-void
.end method

.method public onDraw(Landroid/graphics/Canvas;)V
    .registers 3

    :try_start_0
    invoke-virtual {p0, p1}, LX/1Bem;->LJIIIIZZ(Landroid/graphics/Canvas;)V
    goto :goto_9

    :catch_4
    move-exception v0
    throw v0

    :goto_9
    return-void
.end method
"""


class AvatarRewriteTest(unittest.TestCase):
    def test_the_hand_over_is_bracketed(self):
        patched = avatar_patch.rewrite_avatar(VIEW)
        begin = patched.index(avatar_patch.BEGIN)
        end = patched.index(avatar_patch.END)
        self.assertLess(begin, patched.index("invoke-virtual {p0, p1}, LX/1Bem;->LJIIIIZZ"))
        self.assertLess(patched.index("invoke-virtual {p0, p1}, LX/1Bem;->LJIIIIZZ"), end)
        self.assertLess(end, patched.index("    return-void\n.end method", end))
        self.assertIn("    .registers 3\n\n" + avatar_patch.BEGIN_CALL, patched)

    def test_only_the_draw_method_changes(self):
        patched = avatar_patch.rewrite_avatar(VIEW)
        self.assertEqual(1, patched.count(avatar_patch.BEGIN))
        self.assertEqual(1, patched.count(avatar_patch.END))
        self.assertEqual(VIEW.count(".method "), patched.count(".method "))
        self.assertTrue(patched.startswith(VIEW.split(".method public onDraw")[0]))
        # the calls use the parameter registers only, so no register is added to a method with three
        for register in ("v0", "v1", "v2"):
            self.assertNotIn(register, avatar_patch.BEGIN_CALL + avatar_patch.END_CALL)

    def test_second_passes_and_foreign_classes_are_refused(self):
        with self.assertRaises(ValueError):
            avatar_patch.rewrite_avatar(avatar_patch.rewrite_avatar(VIEW))
        with self.assertRaises(ValueError):
            avatar_patch.rewrite_avatar(".class public Lx;\n.super Ljava/lang/Object;\n")
        with self.assertRaises(ValueError):
            avatar_patch.rewrite_avatar(VIEW.replace("invoke-virtual", "invoke-direct"))
        with self.assertRaises(ValueError):
            avatar_patch.rewrite_avatar(VIEW.replace("    goto :goto_9", "    return-void"))

    def test_runtime_targets_are_what_the_rewrite_calls(self):
        patched = avatar_patch.rewrite_avatar(VIEW)
        for target in avatar_patch.runtime_targets():
            self.assertIn(target, patched)

    @unittest.skipUnless(REAL.exists(), "the disassembled X.1Bem is not in work/")
    def test_the_real_class_is_rewritten(self):
        text = REAL.read_text(encoding="utf-8").replace(chr(13) + chr(10), chr(10))
        patched = avatar_patch.rewrite_avatar(text)
        self.assertEqual(text.count(".method "), patched.count(".method "))
        self.assertEqual(1, patched.count(avatar_patch.BEGIN))


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