"""AstrAI promo: Prefix Cache animation (Radix tree with branches)."""

from manim import *

Text.set_default(font="Times New Roman")


class PrefixCache(Scene):
    """Animates the radix-tree prefix cache with multiple distinct branches."""

    def _add_node(self, parent_pos, label, color, dx, dy):
        pos = parent_pos + np.array([dx, dy, 0])
        dot = Dot(point=pos, color=color, radius=0.1)
        txt = Text(label, font_size=13, color=color)
        txt.next_to(dot, UP, buff=0.1)
        grp = VGroup(dot, txt)
        edge = Line(parent_pos, pos, color=GRAY, stroke_width=1.5)
        return grp, edge, pos

    def _add_leaf(self, parent_pos, color, tag):
        leaf = Square(side_length=0.25, color=color, fill_opacity=0.4)
        leaf.move_to(parent_pos + DOWN * 0.7)
        edge = Line(parent_pos, leaf.get_top(), color=color, stroke_width=1.5)
        lbl = Text(tag, font_size=10, color=color).next_to(leaf, DOWN, buff=0.1)
        return VGroup(leaf, edge, lbl)

    def construct(self):
        title = Text("Prefix Cache", font_size=48, color=BLUE)
        self.play(Write(title))
        self.wait(0.2)
        self.play(title.animate.to_edge(UP).scale(0.6))

        # Root at top-left, tree stays visible throughout
        root_pos = np.array([-4.5, 2.0, 0])
        root = Circle(radius=0.25, color=BLUE, fill_opacity=0.2)
        root.move_to(root_pos)
        root_lbl = Text("root", font_size=10, color=GRAY).move_to(root)
        root_grp = VGroup(root, root_lbl)
        self.play(FadeIn(root_grp, scale=0.5), run_time=0.3)

        # Labels accumulate on the right side
        right_x = 3.5
        label_y = 2.5
        label_step = 0.5

        def show_label(text, color):
            nonlocal label_y
            lbl = Text(text, font_size=14, color=color)
            lbl.move_to([right_x, label_y, 0])
            label_y -= label_step
            self.play(Write(lbl))
            return lbl

        # ── R1: A → B → C ──
        r1_lbl = show_label('R1: "A B C"', GREEN)

        a_grp, a_edge, a_pos = self._add_node(root_pos, "A", GREEN, 0.6, -0.9)
        self.play(Create(a_edge), FadeIn(a_grp, scale=0.5), run_time=0.2)
        b_grp, b_edge, b_pos = self._add_node(a_pos, "B", GREEN, 0.6, -0.9)
        self.play(Create(b_edge), FadeIn(b_grp, scale=0.5), run_time=0.2)
        c_grp, c_edge, c_pos = self._add_node(b_pos, "C", GREEN, 0.6, -0.9)
        self.play(Create(c_edge), FadeIn(c_grp, scale=0.5), run_time=0.2)
        self.play(FadeIn(self._add_leaf(c_pos, GREEN, "slot 0"), scale=0.8), run_time=0.3)
        self.wait(0.3)

        # ── R2: shares A B, branches D E ──
        r2_lbl = show_label('R2: "A B D E"', ORANGE)

        for g in [a_grp, b_grp]:
            flash = SurroundingRectangle(g, color=YELLOW, buff=0.12)
            self.play(Create(flash), run_time=0.1)
            self.play(FadeOut(flash), run_time=0.08)

        d_grp, d_edge, d_pos = self._add_node(b_pos, "D", ORANGE, -0.6, -0.9)
        self.play(Create(d_edge), FadeIn(d_grp, scale=0.5), run_time=0.2)
        e_grp, e_edge, e_pos = self._add_node(d_pos, "E", ORANGE, -0.6, -0.9)
        self.play(Create(e_edge), FadeIn(e_grp, scale=0.5), run_time=0.2)
        self.play(FadeIn(self._add_leaf(e_pos, ORANGE, "slot 1"), scale=0.8), run_time=0.3)
        self.wait(0.3)

        # ── R3: shares A B, single F ──
        r3_lbl = show_label('R3: "A B F"', PINK)

        f_grp, f_edge, f_pos = self._add_node(b_pos, "F", PINK, 0.0, -1.2)
        self.play(Create(f_edge), FadeIn(f_grp, scale=0.5), run_time=0.2)
        self.play(FadeIn(self._add_leaf(f_pos, PINK, "slot 2"), scale=0.8), run_time=0.3)
        self.wait(0.3)

        # ── R4: new prefix from root ──
        r4_lbl = show_label('R4: "X Y"', TEAL)

        x_grp, x_edge, x_pos = self._add_node(root_pos, "X", TEAL, -1.0, -0.9)
        self.play(Create(x_edge), FadeIn(x_grp, scale=0.5), run_time=0.2)
        y_grp, y_edge, y_pos = self._add_node(x_pos, "Y", TEAL, -0.6, -0.9)
        self.play(Create(y_edge), FadeIn(y_grp, scale=0.5), run_time=0.2)
        self.play(FadeIn(self._add_leaf(y_pos, TEAL, "slot 3"), scale=0.8), run_time=0.3)
        self.wait(0.5)

        # ── highlight shared prefix (tree stays) ──
        reuse_box = SurroundingRectangle(VGroup(a_grp, b_grp), color=YELLOW, buff=0.15)
        reuse_note = Text(
            'Prefix "A B" shared\nby 3 requests — 0 copy',
            font_size=16,
            color=YELLOW,
        )
        reuse_note.next_to(reuse_box, LEFT, buff=1.0)
        self.play(Create(reuse_box), Write(reuse_note))
        self.wait(2)
        self.play(FadeOut(reuse_box), FadeOut(reuse_note))

        # ── summary below tree (tree stays visible) ──
        summary = VGroup(
            Text("KV cache reuse across requests", font_size=26, color=GREEN),
            Text("First-token latency: up to 50% reduction", font_size=18, color=GRAY),
        ).arrange(DOWN, buff=0.2)
        summary.to_edge(DOWN, buff=0.5)
        self.play(Write(summary))
        self.wait(2)
        self.play(FadeOut(summary), FadeOut(root_grp), FadeOut(title))