"""ProgramLoader — reads apishift_program.md and selects the section
relevant to the current step.

The Manager's observation includes a 'program_section' field. We don't
inject the entire manual every step (context bloat). Instead we pick the
section that most directly applies to the most recent action context.
"""

import os
import re
from typing import Dict, List


SECTION_RE = re.compile(r"^##\s+\d+\.\s+(.*?)\s*$", re.MULTILINE)


class ProgramLoader:

    def __init__(self, program_path: str = "apishift_program.md"):
        self.program_path = program_path
        self._sections: Dict[str, str] = {}
        self._load()

    def _load(self):
        if not os.path.exists(self.program_path):
            return
        with open(self.program_path, "r", encoding="utf-8") as f:
            text = f.read()
        # Split into sections at "## N. <title>" headers
        parts = re.split(r"(?m)^##\s+(\d+\.\s+.*?)$", text)
        # parts looks like: [intro, 'header_1', 'body_1', 'header_2', 'body_2', ...]
        for i in range(1, len(parts) - 1, 2):
            header = parts[i].strip()
            body = parts[i + 1].strip()
            self._sections[header] = body

    def section_for_step(self, last_command: str = "", steps_taken: int = 0) -> str:
        """Return the most relevant section text for the current step.
        Falls back to the setup ritual section at the start of an episode."""
        if not self._sections:
            return ""

        keywords = {
            "inspect":            "1. Setup ritual",
            "dispatch_diff":      "2. Action ordering rules",
            "classify_impact":    "2. Action ordering rules",
            "dispatch_patch":     "2. Action ordering rules",
            "dispatch_test":      "4. Failure handling",
            "dispatch_rollback":  "2. Action ordering rules",
            "read_memory":        "5. Memory usage rules",
            "submit":             "7. Step budget management",
        }

        if steps_taken == 0:
            target = "1. Setup ritual"
        else:
            target = keywords.get(last_command, "3. Simplicity criterion")

        for header, body in self._sections.items():
            if header.startswith(target.split(".")[0]):
                return f"## {header}\n{body}"
        # Fallback to first section
        first = next(iter(self._sections.items()))
        return f"## {first[0]}\n{first[1]}"

    def all_sections(self) -> List[str]:
        return [f"## {h}\n{b}" for h, b in self._sections.items()]