Update backend/agents/puzzle_agent.py

backend/agents/puzzle_agent.py

@@ -19,7 +19,6 @@ THEME_PALETTES = {
     "default":    {"wall_color": "#8E24AA", "path_color": "#020817", "accent": "#FF477E"},
 }
 
-# Maps required_step keywords → tile IDs
 STEP_TILE_MAP = {
     "meet_manager":    8,
     "meet_scientist":  8,
@@ -33,81 +32,73 @@ STEP_TILE_MAP = {
     "verify_pin":      5,
 }
 
-# ─────────────────────────────────────────────────────────────
-# DAA LEARNING CONCEPTS — mapped to maze structure types
-# ─────────────────────────────────────────────────────────────
-DAA_CONCEPTS = [
-    {
-        "id": "collect_and_reach",
-        "title": "Conditional Logic: Collect Before Proceed",
-        "description": "You CANNOT reach the exit until you collect the required items. This teaches: if-else conditions and pre-conditions in algorithms.",
-        "maze_type": "imperfect",       # multiple paths
-        "min_dim": 7,                   # internal dim → 15x15 grid
-        "max_dim": 11,                  # → 23x23 grid
-        "force_npc": True,
-        "force_item": True,
-    },
-    {
-        "id": "shortest_path",
-        "title": "BFS: Find the Shortest Path",
-        "description": "Multiple routes exist. Explore carefully — some paths are dead ends. The shortest route wins. This teaches: BFS, graph traversal.",
-        "maze_type": "imperfect",       # extra walls removed = multiple routes
-        "min_dim": 8,
-        "max_dim": 12,
-        "force_npc": False,
-        "force_item": True,             # at least collect an item on the way
-    },
-    {
-        "id": "dfs_backtrack",
-        "title": "DFS: Explore and Backtrack",
-        "description": "The maze has many dead ends. Systematically explore each branch. When stuck, backtrack and try another. This teaches: DFS and recursion.",
-        "maze_type": "perfect",         # perfect = single path but MASSIVE dead ends
-        "min_dim": 10,
-        "max_dim": 15,
-        "force_npc": True,
-        "force_item": True,
-    },
-    {
-        "id": "sequence",
-        "title": "Sequencing: Visit in Correct Order",
-        "description": "Visit the Sentinel (🤵) FIRST, then collect the Key (🔑), THEN reach the Exit. Order matters! This teaches: sequential execution and ordered steps.",
-        "maze_type": "imperfect",
-        "min_dim": 7,
-        "max_dim": 11,
-        "force_npc": True,
-        "force_item": True,
-    },
-]
+MAZE_TOOLBOX_XML = (
+    '<category name="Movement" colour="230">'
+    '<block type="move_up"></block><block type="move_down"></block>'
+    '<block type="move_left"></block><block type="move_right"></block>'
+    '</category>'
+    '<category name="Loops" colour="120">'
+    '<block type="controls_repeat_ext">'
+    '<value name="TIMES"><block type="math_number"><field name="NUM">3</field></block></value>'
+    '</block></category>'
+)
+
+
 def _get_best_flash_model() -> str:
-    """Dynamically finds the best available Gemini Flash model."""
     try:
         available_models = []
         for m in genai.list_models():
-            if 'generateContent' in m.supported_generation_methods:
-                if 'flash' in m.name.lower():
-                    available_models.append(m.name)
-        
+            if 'generateContent' in m.supported_generation_methods and 'flash' in m.name.lower():
+                available_models.append(m.name)
         if available_models:
-            # Sort descending to prefer newer versions (e.g. gemini-2.0-flash over gemini-1.5-flash)
             available_models.sort(reverse=True)
             return available_models[0].replace('models/', '')
-            
-        # Fallback to any gemini model
         for m in genai.list_models():
             if 'generateContent' in m.supported_generation_methods and 'gemini' in m.name.lower():
                 return m.name.replace('models/', '')
-                
     except Exception as e:
         print(f"Warning: Could not list models ({e}). Using hardcoded default.")
-        
     return "gemini-2.0-flash"
 
 
+def _load_json_data(base_dir: str, filename: str) -> Any:
+    """Load JSON from backend/data with fallbacks for local dev and HF Docker."""
+    candidates = [
+        os.path.join(base_dir, "data", filename),
+        os.path.join(base_dir, "..", filename),
+        os.path.join(os.path.dirname(base_dir), filename),
+    ]
+    for path in candidates:
+        if os.path.isfile(path):
+            with open(path, "r", encoding="utf-8") as f:
+                return json.load(f)
+    raise FileNotFoundError(f"Could not find {filename} in {candidates}")
+
+
+def _is_prime(n: int) -> bool:
+    if n < 2:
+        return False
+    if n == 2:
+        return True
+    if n % 2 == 0:
+        return False
+    d = 3
+    while d * d <= n:
+        if n % d == 0:
+            return False
+        d += 2
+    return True
+
+
+def _is_palindrome(n: int) -> bool:
+    s = str(n)
+    return s == s[::-1]
+
+
 class PuzzleAgent:
     """
-    Logic Architect — DAA Learning Maze Engine.
-    Generates complex, educationally-rich mazes with mandatory checkpoints
-    and a specific CS/DAA learning concept per level.
+    Curriculum-driven maze puzzle generator for CodeCracker Puzzle mode.
+    Each 5-level arc teaches one concept per level with Cody & Zara continuity.
     """
 
     def __init__(self, api_key: Optional[str] = None):
@@ -121,36 +112,35 @@ class PuzzleAgent:
             print(f"[PuzzleAgent] Selected Model: {best_model}")
             self.model = genai.GenerativeModel(best_model)
 
-        # Load the 500-scenario dataset
-        # Look for the dataset inside the backend/data folder
-        base_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
-        path = os.path.join(base_dir, "data", "maze_scenarios_500.json")
+        self.base_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+
         self.scenarios: List[Dict] = []
         try:
-            with open(path, "r") as f:
-                self.scenarios = json.load(f).get("scenarios", [])
+            data = _load_json_data(self.base_dir, "maze_scenarios_500.json")
+            self.scenarios = data.get("scenarios", [])
             print(f"[PuzzleAgent] Loaded {len(self.scenarios)} scenarios.")
         except Exception as e:
-            print(f"[PuzzleAgent] Could not load dataset: {e}")
+            print(f"[PuzzleAgent] Could not load scenarios: {e}")
+
+        self.catalog: Dict = {"arcs": [], "concepts": {}}
+        self.characters: Dict = {}
+        try:
+            self.catalog = _load_json_data(self.base_dir, "maze_concept_catalog.json")
+            self.characters = _load_json_data(self.base_dir, "character_bible.json")
+            print(f"[PuzzleAgent] Loaded catalog with {len(self.catalog.get('arcs', []))} arcs.")
+        except Exception as e:
+            print(f"[PuzzleAgent] Could not load catalog/bible: {e}")
 
     # ─────────────────────────────────────────────────────────
-    # KAGGLE-IDENTICAL PERFECT MAZE (Recursive Backtracking)
+    # MAZE GENERATION
     # ─────────────────────────────────────────────────────────
     def _create_perfect_maze(self, dim: int) -> List[List[int]]:
-        """
-        Generates a perfect maze (one path between any two cells).
-        Identical algorithm to the Kaggle maze dataset.
-        dim = number of 'rooms' → grid size = 2*dim+1
-        """
         size = dim * 2 + 1
         maze = [[1] * size for _ in range(size)]
-
         x, y = 0, 0
         maze[2 * x + 1][2 * y + 1] = 0
-
         stack = [(x, y)]
         sys.setrecursionlimit(50000)
-
         while stack:
             x, y = stack[-1]
             directions = [(0, 1), (1, 0), (0, -1), (-1, 0)]
@@ -166,17 +156,9 @@ class PuzzleAgent:
                     break
             if not moved:
                 stack.pop()
-
         return maze
 
-    # ─────────────────────────────────────────────────────────
-    # IMPERFECT MAZE (perfect + extra wall removals = loops)
-    # ─────────────────────────────────────────────────────────
     def _create_imperfect_maze(self, dim: int, extra_wall_removals: float = 0.20) -> List[List[int]]:
-        """
-        Imperfect maze: removes ~20% of qualifying walls to create multiple routes.
-        Forces the user to evaluate multiple branching paths.
-        """
         maze = self._create_perfect_maze(dim)
         size = len(maze)
         wall_candidates = []
@@ -187,16 +169,11 @@ class PuzzleAgent:
                         wall_candidates.append((i, j))
                     elif maze[i][j - 1] == 0 and maze[i][j + 1] == 0:
                         wall_candidates.append((i, j))
-
         num_to_remove = int(len(wall_candidates) * extra_wall_removals)
         for i, j in random.sample(wall_candidates, min(num_to_remove, len(wall_candidates))):
             maze[i][j] = 0
-
         return maze
 
-    # ─────────────────────────────────────────────────────────
-    # BFS PATH FINDER
-    # ─────────────────────────────────────────────────────────
     def _bfs_path(self, grid: List[List[int]], start: Tuple, goal: Tuple) -> List[Tuple]:
         rows, cols = len(grid), len(grid[0])
         visited = {start}
@@ -205,152 +182,260 @@ class PuzzleAgent:
             (r, c), path = queue.popleft()
             if (r, c) == goal:
                 return path
-            for dr, dc in [(-1,0),(1,0),(0,-1),(0,1)]:
+            for dr, dc in [(-1, 0), (1, 0), (0, -1), (0, 1)]:
                 nr, nc = r + dr, c + dc
                 if 0 <= nr < rows and 0 <= nc < cols and (nr, nc) not in visited and grid[nr][nc] != 1:
                     visited.add((nr, nc))
                     queue.append(((nr, nc), path + [(nr, nc)]))
         return []
 
-    # ─────────────────────────────────────────────────────────
-    # CORE: Build one maze with GUARANTEED NPC + Item on path
-    # ─────────────────────────────────────────────────────────
-    def _build_level_maze(self, dim: int, maze_type: str, theme: str,
-                          force_npc: bool, force_item: bool,
-                          step_names: List[str]) -> Dict:
-        """
-        Builds a maze of internal size `dim` (grid = 2*dim+1).
-        ALWAYS places NPC (8) and Item (5) on the guaranteed BFS path.
-        """
-        # Generate maze
+    def _generate_cell_values(self, concept: Dict, path_middle: List[Tuple]) -> Tuple[Dict[str, int], Dict]:
+        """Place numbered floor tiles on the solution path for puzzle rules."""
+        rule = dict(concept.get("maze_rule") or {"type": "none"})
+        rule_type = rule.get("type", "none")
+        cell_values: Dict[str, int] = {}
+
+        if not path_middle or rule_type == "none" or rule_type == "ordered_checkpoints":
+            return cell_values, rule
+
+        slots = path_middle[::max(1, len(path_middle) // 4)][:5]
+        if not slots:
+            slots = path_middle[:3]
+
+        if rule_type == "parity":
+            pass_if = rule.get("pass_if", "even")
+            for idx, (r, c) in enumerate(slots):
+                val = random.choice([2, 4, 6, 8, 10, 12]) if pass_if == "even" else random.choice([3, 5, 7, 9, 11])
+                if idx % 2 == 1:
+                    val = random.choice([3, 5, 7, 9, 11]) if pass_if == "even" else random.choice([2, 4, 6, 8, 10])
+                cell_values[f"{r},{c}"] = val
+            return cell_values, rule
+
+        if rule_type == "prime_only":
+            pool = [4, 6, 8, 9, 10, 12, 14, 15]
+            primes = [3, 5, 7, 11, 13]
+            for i, (r, c) in enumerate(slots):
+                cell_values[f"{r},{c}"] = primes[i % len(primes)] if i % 2 == 0 else pool[i % len(pool)]
+            return cell_values, rule
+
+        if rule_type == "palindrome_only":
+            pals = [3, 5, 7, 11, 121]
+            non_pals = [12, 14, 18, 19]
+            for i, (r, c) in enumerate(slots):
+                cell_values[f"{r},{c}"] = pals[i % len(pals)] if i % 2 == 0 else non_pals[i % len(non_pals)]
+            return cell_values, rule
+
+        if rule_type == "max_label_only":
+            if len(slots) >= 2:
+                low, high = 8, 17
+                cell_values[f"{slots[0][0]},{slots[0][1]}"] = low
+                cell_values[f"{slots[1][0]},{slots[1][1]}"] = high
+                rule = {"type": "labeled_safe", "safe_values": [high]}
+            return cell_values, rule
+
+        return cell_values, rule
+
+    def _build_level_maze(
+        self,
+        dim: int,
+        concept: Dict,
+        step_names: List[str],
+    ) -> Dict:
+        maze_type = concept.get("maze_type", "imperfect")
+        force_npc = concept.get("force_npc", False)
+        force_item = concept.get("force_item", False)
+
         if maze_type == "imperfect":
             grid = self._create_imperfect_maze(dim)
         else:
             grid = self._create_perfect_maze(dim)
 
         size = len(grid)
-
-        # Start top-left corridor, Goal bottom-right corridor
         sr, sc = 1, 1
         gr, gc = size - 2, size - 2
-
         grid[sr][sc] = 0
         grid[gr][gc] = 0
 
-        # Get BFS path
         path = self._bfs_path(grid, (sr, sc), (gr, gc))
         if not path or len(path) < 5:
-            # Safety: rebuild with a simpler grid
             grid = self._create_perfect_maze(max(dim, 7))
             size = len(grid)
             sr, sc = 1, 1
             gr, gc = size - 2, size - 2
-            grid[sr][sc] = 0; grid[gr][gc] = 0
+            grid[sr][sc] = 0
+            grid[gr][gc] = 0
             path = self._bfs_path(grid, (sr, sc), (gr, gc))
 
-        path_middle = path[1:-1]  # exclude start and goal
-
-        # Build ordered tiles from step names + forced tiles
-        ordered = []  # [(step_label, tile_id), ...]
-
-        # Use scenario steps first
-        used_npc = False
-        used_item = False
+        path_middle = path[1:-1]
+        ordered = []
+        used_npc = used_item = False
         for step in step_names:
             tid = STEP_TILE_MAP.get(step)
             if tid == 8 and not used_npc:
-                ordered.append((step, 8)); used_npc = True
+                ordered.append((step, 8))
+                used_npc = True
             elif tid == 5 and not used_item:
-                ordered.append((step, 5)); used_item = True
+                ordered.append((step, 5))
+                used_item = True
 
-        # FORCE NPC and Item if not already present
         if force_npc and not used_npc:
             ordered.insert(0, ("meet_expert", 8))
         if force_item and not used_item:
             ordered.append(("collect_data", 5))
 
-        # Place ordered tiles evenly on the middle path
-        n = len(ordered)
         placed = []
+        n = len(ordered)
         for i, (step_label, tile_id) in enumerate(ordered):
             if not path_middle:
                 break
             idx = (i + 1) * len(path_middle) // (n + 1)
             idx = max(0, min(idx, len(path_middle) - 1))
             r, c = path_middle[idx]
-            # Avoid placing on same cell
             while (r, c) in [(pr, pc) for _, (pr, pc) in placed] and idx + 1 < len(path_middle):
                 idx += 1
                 r, c = path_middle[idx]
             grid[r][c] = tile_id
             placed.append((step_label, (r, c)))
 
-        # Place Start and Goal
         grid[sr][sc] = 2
         grid[gr][gc] = 3
 
-        ordered_step_names = [s for s, _ in ordered]
-        palette = THEME_PALETTES.get(theme, THEME_PALETTES["default"])
+        cell_values, maze_rule = self._generate_cell_values(concept, path_middle)
 
         return {
             "grid": grid,
             "rows": size,
             "cols": size,
-            "ordered_steps": ordered_step_names,
-            "palette": palette,
+            "ordered_steps": [s for s, _ in ordered],
+            "cell_values": cell_values,
+            "maze_rule": maze_rule,
         }
 
     # ─────────────────────────────────────────────────────────
-    # AI STORY ENRICHMENT
+    # ARC & SCENARIO SELECTION
+    # ─────────────────────────────────────────────────────────
+    def _pick_arc(self, diff_key: str, topic: Optional[str]) -> Dict:
+        arcs = self.catalog.get("arcs", [])
+        if topic:
+            t = topic.lower()
+            matched = [a for a in arcs if a.get("theme", "").lower() == t or t in a.get("title", "").lower()]
+            if matched:
+                arcs = matched
+        matched_diff = [a for a in arcs if a.get("difficulty", "easy") == diff_key]
+        if matched_diff:
+            return random.choice(matched_diff)
+        if arcs:
+            return random.choice(arcs)
+        return {
+            "arc_id": "fallback_arc",
+            "title": "Training Protocol",
+            "theme": "space",
+            "difficulty": diff_key,
+            "concept_ids": ["sequence", "even_odd", "collect_and_reach", "shortest_path", "loops_path"],
+        }
+
+    def _get_concept(self, concept_id: str) -> Dict:
+        concepts = self.catalog.get("concepts", {})
+        if concept_id in concepts:
+            return concepts[concept_id]
+        return {
+            "id": concept_id,
+            "title": "Maze Navigation",
+            "learning_outcome": "Navigate checkpoints in order.",
+            "description": "Reach the exit by following the mission steps.",
+            "maze_type": "imperfect",
+            "min_dim": 7,
+            "max_dim": 10,
+            "force_npc": True,
+            "force_item": True,
+            "maze_rule": {"type": "ordered_checkpoints"},
+            "hint_template": "Follow the checklist step by step.",
+        }
+
+    def _pick_scenarios_for_arc(self, arc: Dict, diff_key: str) -> List[Dict]:
+        theme = arc.get("theme", "space").lower()
+        filtered = [
+            s for s in self.scenarios
+            if s.get("difficulty", "").lower() == diff_key and s.get("theme", "").lower() == theme
+        ]
+        if len(filtered) < 5:
+            filtered = [s for s in self.scenarios if s.get("difficulty", "").lower() == diff_key]
+        if len(filtered) < 5:
+            filtered = self.scenarios
+        if len(filtered) >= 5:
+            return random.sample(filtered, 5)
+        pool = filtered if filtered else [{}]
+        return [random.choice(pool) for _ in range(5)]
+
     # ─────────────────────────────────────────────────────────
-    def _enrich_with_ai(self, scenarios: List[Dict], concept: Dict,
-                        compiled_grids: List[Dict]) -> Dict:
+    # AI STORY (Cody & Zara, one arc)
+    # ─────────────────────────────────────────────────────────
+    def _enrich_with_ai(
+        self,
+        arc: Dict,
+        level_concepts: List[Dict],
+        scenarios: List[Dict],
+        compiled_grids: List[Dict],
+    ) -> Dict:
         if not self.model:
             raise Exception("No AI model")
 
-        # Build progression details
+        cast = self.characters.get("cast", [])
+        cody = next((c for c in cast if c.get("id") == "cody"), {"name": "Captain Cody", "emoji": "🧑‍🚀"})
+        zara = next((c for c in cast if c.get("id") == "zara"), {"name": "Engineer Zara", "emoji": "👩‍🔬"})
+        setting = self.characters.get("setting", "Planet Antaria")
+
         progression_str = ""
-        for i, (scen, cg) in enumerate(zip(scenarios, compiled_grids)):
-            t = scen.get("theme", "space").capitalize()
-            title = scen.get("title", f"Level {i+1}")
-            base_story = scen.get("story", "Solve the maze.")
-            s = " → ".join(cg["ordered_steps"]) or "reach the exit"
-            progression_str += f"Level {i+1}:\n- Theme: {t}\n- Basic Concept: {title} - {base_story}\n- Steps: {s}\n\n"
+        for i, (concept, scen, cg) in enumerate(zip(level_concepts, scenarios, compiled_grids)):
+            steps = " → ".join(cg["ordered_steps"]) or "reach the exit"
+            nums = ", ".join(str(v) for v in cg.get("cell_values", {}).values())
+            progression_str += (
+                f"Level {i + 1} — Teach: {concept['title']}\n"
+                f"  Learning: {concept['learning_outcome']}\n"
+                f"  Puzzle rule: {concept['description']}\n"
+                f"  Numbered tiles on path: {nums or 'none'}\n"
+                f"  Mission steps: {steps}\n"
+                f"  Setting hint: {scen.get('title', 'Mission')}\n\n"
+            )
 
         prompt = f"""
-You are the Mission Architect for CodeCracker, a STEM education app for kids.
-Create a 5-level story arc. Each level uses a completely different theme and mission steps.
+You are the Mission Architect for CodeCracker Puzzle mode (maze-only teaching, ages 10-16).
 
-MISSION PROGRESSION:
-{progression_str}
+CHARACTERS (use in EVERY level):
+- {cody['name']} {cody['emoji']} — navigator
+- {zara['name']} {zara['emoji']} — explains logic rules
+
+SETTING: {setting}
+STORY ARC: "{arc.get('title', 'Mission')}" (arc_id: {arc.get('arc_id')})
+CRITICAL: All 5 levels are ONE continuous story on the SAME planet/ship crisis. Do NOT switch to unrelated themes like random bank/hospital.
 
-DAA CONCEPT: {concept['title']}
-LEARNING GOAL: {concept['description']}
+LEVEL PROGRESSION (each level teaches ONE new programming idea via the maze):
+{progression_str}
 
-STRICT RULES:
-1. The story MUST explain the DAA concept clearly to a student (age 10-16).
-2. Reference the exact mission steps for that level in its goal_description.
-3. CRITICAL: Read the "Basic Concept" provided for each level in the progression, and expand it into a UNIQUE, fully-fleshed out programming/logic problem in the `story` field for EACH level. Do NOT just repeat the basic concept.
-4. The `title` of each level should match its assigned Theme (e.g., "Bank Heist", "Space Station").
-5. CRITICAL: Generate a COMPLETELY UNIQUE `hint` for EVERY level. The hint MUST be different for Level 1, Level 2, etc., and must directly reference the specific theme/logic of that exact level! Do NOT repeat hints!
+RULES:
+1. `story` must star Cody and Zara and explain THAT level's concept in plain language (even/odd, primes, sequence, BFS, loops, etc.).
+2. `goal_description` must list the exact mission steps for that level.
+3. `concept_tutorial` = one short line shown in the app banner (use the learning outcome).
+4. `hint` must be unique per level; use the puzzle logic, not generic advice.
+5. `npc_emoji` should be {zara['emoji']} when meeting the expert; `item_emoji` 🔑; `goal_emoji` 🏁.
+6. Exactly 5 levels in order.
 
 OUTPUT only valid JSON (no markdown):
 {{
-  "story_arc_title": "Grand Mission Array",
-  "daa_concept_title": "{concept['title']}",
-  "daa_description": "{concept['description']}",
+  "story_arc_title": "{arc.get('title', 'Antaria Mission')}",
+  "arc_id": "{arc.get('arc_id', 'arc')}",
   "levels": [
     {{
       "type": "maze",
-      "level_id": "unique_id",
-      "title": "Level Title Matching Theme",
-      "story": "Present the unique logical problem or scenario for this level here...",
-      "goal_description": "Step 1: [action] → Step 2: [action] → Reach Exit",
-      "hint": "Specific, highly useful hint for this specific puzzle logic",
-      "npc_emoji": "🏦",
-      "item_emoji": "💵",
-      "goal_emoji": "🏁",
-      "allowed_blocks": ["move_forward", "turn_left", "turn_right"]
+      "level_id": "arc_lvl_1",
+      "title": "Episode title",
+      "story": "Cody and Zara ...",
+      "goal_description": "Step 1 → Step 2 → Exit",
+      "concept_tutorial": "Short learning line",
+      "hint": "Level-specific hint",
+      "npc_emoji": "{zara['emoji']}",
+      "item_emoji": "🔑",
+      "goal_emoji": "🏁"
     }}
   ]
 }}
@@ -363,140 +448,112 @@ OUTPUT only valid JSON (no markdown):
             text = text.split("```")[1].split("```")[0]
         return json.loads(text.strip())
 
+    def _build_fallback_arc(
+        self,
+        arc: Dict,
+        level_concepts: List[Dict],
+        scenarios: List[Dict],
+        compiled_grids: List[Dict],
+        error_msg: str = "Unknown",
+    ) -> Dict:
+        cast = self.characters.get("cast", [])
+        cody = next((c for c in cast if c.get("id") == "cody"), {"name": "Captain Cody", "emoji": "🧑‍🚀"})
+        zara = next((c for c in cast if c.get("id") == "zara"), {"name": "Engineer Zara", "emoji": "👩‍🔬"})
+
+        levels = []
+        for i, (concept, cg) in enumerate(zip(level_concepts, compiled_grids)):
+            steps = cg["ordered_steps"]
+            step_str = " → ".join(steps) + " → Exit" if steps else "Reach the Exit"
+            levels.append({
+                "type": "maze",
+                "level_id": f"{arc.get('arc_id', 'arc')}_{i + 1}",
+                "title": f"{concept['title']} — Part {i + 1}",
+                "story": (
+                    f"{cody['emoji']} {cody['name']} and {zara['emoji']} {zara['name']} continue "
+                    f'"{arc.get("title", "the mission")}".\n\n'
+                    f"Level {i + 1}: {concept['description']}\n\n"
+                    f"(Story AI offline: {error_msg})"
+                ),
+                "goal_description": step_str,
+                "concept_tutorial": concept.get("learning_outcome", concept["title"]),
+                "hint": concept.get("hint_template", "Plan your route before moving."),
+                "npc_emoji": zara.get("emoji", "👩‍🔬"),
+                "item_emoji": "🔑",
+                "goal_emoji": "🏁",
+            })
+        return {
+            "story_arc_title": arc.get("title", "Training Arc"),
+            "arc_id": arc.get("arc_id"),
+            "levels": levels,
+        }
+
     # ─────────────────────────────────────────────────────────
     # PUBLIC ENTRY POINT
     # ─────────────────────────────────────────────────────────
     def generate_level(self, difficulty: str, topic: Optional[str] = None) -> Dict[str, Any]:
-        """
-        Main pipeline:
-        1. Pick a scenario from the dataset
-        2. Pick a DAA concept matching difficulty
-        3. Build 5 progressively larger mazes with FORCED NPC+Item checkpoints
-        4. Enrich with AI story
-        5. Return complete arc
-        """
-        # Normalise: app sends "Beginner"/"Intermediate"/"Expert"
-        #            dataset uses  "easy"/"medium"/"hard"
         _diff_norm = {
-            "beginner":     "easy",
-            "easy":         "easy",
-            "intermediate": "medium",
-            "medium":       "medium",
-            "expert":       "hard",
-            "hard":         "hard",
+            "beginner": "easy", "easy": "easy",
+            "intermediate": "medium", "medium": "medium",
+            "expert": "hard", "hard": "hard",
         }
         diff_key = _diff_norm.get(difficulty.lower(), "easy")
 
-        # Step 1: Pick 5 completely DISTINCT scenarios matching normalised difficulty
-        filtered = [s for s in self.scenarios if s.get("difficulty", "").lower() == diff_key]
-        if len(filtered) < 5:
-            # If not enough filtered, use the entire dataset to guarantee 5 unique scenarios
-            filtered = self.scenarios
-            
-        arc_scenarios = random.sample(filtered, min(5, len(filtered)))
-        # Fallback if dataset has less than 5 elements
-        while len(arc_scenarios) < 5:
-            arc_scenarios.append(random.choice(filtered) if filtered else {})
-
-        # Ensure each level gets a distinct visual theme color palette
-        palette_keys = list(THEME_PALETTES.keys())
-        random.shuffle(palette_keys)
-
-        # Step 2: Pick DAA concept — harder = more complex maze type + larger dim
-        concept_map = {
-            "easy":   [DAA_CONCEPTS[0], DAA_CONCEPTS[3]],   # collect+reach, sequence
-            "medium": [DAA_CONCEPTS[1], DAA_CONCEPTS[0]],   # shortest_path, collect+reach
-            "hard":   [DAA_CONCEPTS[2], DAA_CONCEPTS[1]],   # dfs_backtrack, shortest_path
-        }
-        concept_pool = concept_map.get(diff_key, DAA_CONCEPTS)
-        concept = random.choice(concept_pool)
+        arc = self._pick_arc(diff_key, topic)
+        concept_ids = arc.get("concept_ids", [])[:5]
+        while len(concept_ids) < 5:
+            concept_ids.append("sequence")
+
+        level_concepts = [self._get_concept(cid) for cid in concept_ids]
+        arc_scenarios = self._pick_scenarios_for_arc(arc, diff_key)
+        theme = arc.get("theme", "space").lower()
+        palette = THEME_PALETTES.get(theme, THEME_PALETTES["default"])
 
-        # Step 3: Build 5 mazes, growing in complexity, each with its own scenario
         compiled_grids = []
-        for i in range(5):
-            dim = random.randint(concept["min_dim"] + i, concept["max_dim"] + i)
-            lvl_scenario = arc_scenarios[i]
-            lvl_theme = lvl_scenario.get("theme", "space").lower()
-            lvl_steps = lvl_scenario.get("required_steps", [])
-            
-            grid_data = self._build_level_maze(
-                dim=dim,
-                maze_type=concept["maze_type"],
-                theme=lvl_theme,
-                force_npc=concept["force_npc"],
-                force_item=concept["force_item"],
-                step_names=lvl_steps,
+        for i, concept in enumerate(level_concepts):
+            dim = random.randint(
+                concept.get("min_dim", 7) + i // 2,
+                concept.get("max_dim", 10) + i // 2,
             )
-            compiled_grids.append(grid_data)
+            steps = arc_scenarios[i].get("required_steps", []) if i < len(arc_scenarios) else []
+            compiled_grids.append(self._build_level_maze(dim, concept, steps))
 
-        # Step 4: AI enrichment
         try:
-            arc = self._enrich_with_ai(arc_scenarios, concept, compiled_grids)
+            arc_json = self._enrich_with_ai(arc, level_concepts, arc_scenarios, compiled_grids)
         except Exception as e:
-            error_msg = str(e)
-            print(f"[PuzzleAgent] AI enrichment failed: {error_msg}. Using fallback.")
-            arc = self._build_fallback_arc(arc_scenarios, concept, compiled_grids, error_msg)
+            print(f"[PuzzleAgent] AI enrichment failed: {e}. Using fallback.")
+            arc_json = self._build_fallback_arc(arc, level_concepts, arc_scenarios, compiled_grids, str(e))
 
-        # Step 5: Inject compiled grids and dynamic palettes into each level
-        levels = arc.get("levels", [])
+        levels = arc_json.get("levels", [])
         for i, level in enumerate(levels):
-            cg = compiled_grids[i] if i < len(compiled_grids) else compiled_grids[-1]
-            level["grid_layout"]    = cg["grid"]
-            level["maze_rows"]      = cg["rows"]
-            level["maze_cols"]      = cg["cols"]
-            level["ordered_steps"]  = cg["ordered_steps"]
-            
-            # Assign a distinct color palette
-            palette_key = palette_keys[i % len(palette_keys)]
-            base_palette = THEME_PALETTES.get(palette_key, THEME_PALETTES["default"]).copy()
-            
-            # Inject dynamic AI emojis into the palette so the Flutter UI can read them
-            base_palette["npc_emoji"] = level.get("npc_emoji", "🤵")
+            if i >= len(compiled_grids):
+                break
+            cg = compiled_grids[i]
+            concept = level_concepts[i]
+            level["grid_layout"] = cg["grid"]
+            level["maze_rows"] = cg["rows"]
+            level["maze_cols"] = cg["cols"]
+            level["ordered_steps"] = cg["ordered_steps"]
+            level["cell_values"] = cg["cell_values"]
+            level["maze_rule"] = cg["maze_rule"]
+            level["concept_id"] = concept["id"]
+            level["concept_title"] = concept["title"]
+            level["track"] = concept.get("track", "PF")
+            level["learning_outcome"] = concept.get("learning_outcome", concept["title"])
+            level["concept_tutorial"] = level.get("concept_tutorial") or concept.get("learning_outcome", "")
+            level["type"] = "maze"
+            level["toolbox_xml"] = MAZE_TOOLBOX_XML
+            level["allowed_blocks"] = ["move_up", "move_down", "move_left", "move_right", "controls_repeat_ext"]
+
+            base_palette = palette.copy()
+            base_palette["npc_emoji"] = level.get("npc_emoji", "👩‍🔬")
             base_palette["item_emoji"] = level.get("item_emoji", "🔑")
-            base_palette["goal_emoji"] = level.get("goal_emoji", "🚀")
-            
-            level["theme_palette"]  = base_palette
-            level["daa_concept"]    = concept["title"]
-            level["daa_description"]= concept["description"]
-            level["type"]           = "maze"
-
-        arc["levels"] = levels
-        # Remove top-level theme_palette so we don't override the distinct ones later
-        return arc
-
-    # ─────────────────────────────────────────────────────────
-    # FALLBACK (no AI or API failure)
-    # ─────────────────────────────────────────────────────────
-    def _build_fallback_arc(self, scenarios: List[Dict], concept: Dict,
-                            compiled_grids: List[Dict], error_msg: str = "Unknown") -> Dict:
-        levels = []
-        for i, cg in enumerate(compiled_grids):
-            theme = scenarios[i].get("theme", "mission").capitalize()
-            steps = cg["ordered_steps"]
-            step_str = " → ".join(steps) or "Reach the Exit"
+            base_palette["goal_emoji"] = level.get("goal_emoji", "🏁")
+            level["theme_palette"] = base_palette
 
-            levels.append({
-                "type": "maze",
-                "level_id": f"fallback_{i+1:03d}",
-                "title": f"{theme} Protocol — Level {i+1}",
-                "story": (
-                    f"SERVER ERROR: {error_msg}\n\n"
-                    f"Mission {i+1}: Apply {concept['title']} to navigate the {theme.lower()} labyrinth. "
-                    f"{concept['description']}"
-                ),
-                "goal_description": step_str + " → Exit",
-                "hint": "Try planning your path by starting from the end and tracing backwards.",
-                "npc_emoji": "🤵",
-                "item_emoji": "🔑",
-                "goal_emoji": "🚀",
-                "daa_concept": concept["title"],
-                "daa_description": concept["description"],
-                "allowed_blocks": ["move_forward", "turn_left", "turn_right"],
-            })
-        return {
-            "story_arc_title": f"Grand {concept['title']}",
-            "levels": levels,
-        }
+        arc_json["levels"] = levels[:5]
+        arc_json["arc_id"] = arc.get("arc_id")
+        arc_json["theme"] = theme
+        return arc_json
 
     def learn_from_feedback(self, level_data: Dict[str, Any], rating: int,
                             developer_feedback: str = None):