"""AstrAI promo: Full architecture overview."""

from manim import *


class Architecture(Scene):
    """Animates the full AstrAI system stack layer by layer."""

    def construct(self):
        title = Text("AstrAI Architecture", font_size=48, color=BLUE)
        self.play(Write(title))
        self.wait(0.2)
        self.play(title.animate.to_edge(UP))

        layers_data = [
            (0.9, GREEN, "API Layer", ["FastAPI Server  •  OpenAI-Compatible API"]),
            (0.9, BLUE, "Inference Engine", ["Streaming  •  Async  •  Batch Modes"]),
            (1.6, YELLOW, "Continuous Batching Scheduler",
             ["Cleanup → Refill → Prefill → Decode",
              "Position-Grouped Decode",
              "Bitmask O(1) Slot Allocation"]),
            (1.2, ORANGE, "Prefix Cache + KV Cache",
             ["Radix Tree  •  Slot Versioning",
              "GPU copy_()  →  Zero-Copy Reuse"]),
            (1.2, PURPLE, "Transformer Model (1B params)",
             ["24-layer GQA  •  RoPE  •  SwiGLU",
              "bfloat16  •  100K vocab"]),
        ]

        layers = VGroup()
        for height, color, label, subs in layers_data:
            box = Rectangle(width=7.5, height=height, color=color, fill_opacity=0.1)
            lbl = Text(label, font_size=18, color=color)
            items = [lbl] + [Text(s, font_size=11, color=WHITE) for s in subs]
            content = VGroup(*items)
            content.arrange(DOWN, buff=0.22)
            content.move_to(box.get_center())
            layers.add(VGroup(box, content))

        layers.arrange(DOWN, buff=0.18)
        layers.next_to(title, DOWN, buff=0.3)

        for i in range(len(layers)):
            self.play(Create(layers[i]), run_time=0.35)
            if i > 0:
                # Use box-to-box for arrow endpoints (not content)
                prev_box = layers[i - 1][0]
                curr_box = layers[i][0]
                arrow = Arrow(
                    prev_box.get_bottom(),
                    curr_box.get_top(),
                    color=GRAY,
                    buff=0.1,
                    max_tip_length_to_length_ratio=0.15,
                )
                self.play(Create(arrow), run_time=0.15)

        self.wait(0.5)

        hl = SurroundingRectangle(layers[3], color=GREEN, buff=0.12)
        hl_note = Text("Zero-Copy Prefix Reuse", font_size=22, color=GREEN)
        hl_note.next_to(hl, RIGHT, buff=0.8)
        self.play(Create(hl), Write(hl_note))
        self.wait(1.5)
        self.play(FadeOut(hl), FadeOut(hl_note))

        self.play(FadeOut(layers))

        cta = VGroup(
            Text("AstrAI", font_size=52, color=BLUE),
            Text("Single GPU  •  Open Source  •  1B params", font_size=24, color=GRAY),
            Text("github.com/ViperEkura/AstrAI", font_size=20, color=YELLOW),
        ).arrange(DOWN, buff=0.35)
        self.play(Write(cta))
        self.wait(2)
        self.play(FadeOut(cta), FadeOut(title))