graph TD
    A[System Layer] --> B[Identity & Constraints]
    A --> C[Output Schema Definition]
    A --> D[Tool Definitions]
    B --> E[Prompt Assembly]
    C --> E
    D --> E
    F[User Layer] --> G[Task Specification]
    F --> H[Dynamic Few-Shot Examples]
    G --> E
    H --> E
    I[Assistant Layer] --> J[Reasoning Primer]
    I --> K[Format Enforcement]
    J --> E
    K --> E
    E --> L[Validation Gate]
    L -->|Pass| M[Agent Execution]
    L -->|Fail| N[Prompt Revision]
    N --> E

interface PromptLayer {
  role: "system" | "user" | "assistant";
  sections: PromptSection[];
}

interface PromptSection {
  name: string;
  content: string;
  priority: number;
  tokenBudget: number;
}

function assemblePrompt(layers: PromptLayer[], maxTokens: number): Message[] {
  const messages: Message[] = [];
  for (const layer of layers) {
    const sections = layer.sections
      .sort((a, b) => b.priority - a.priority)
      .reduce((acc, section) => {
        const currentTokens = countTokens(acc.map(s => s.content).join("\n"));
        if (currentTokens + section.tokenBudget <= maxTokens * 0.4) {
          acc.push(section);
        }
        return acc;
      }, [] as PromptSection[]);

    messages.push({
      role: layer.role,
      content: sections.map(s => s.content).join("\n\n"),
    });
  }
  return messages;
}

SCHEMA_ENFORCEMENT_PROMPT = """
You MUST respond with valid JSON matching this exact schema:
{schema}

Rules:
- Every field is required unless marked optional
- String fields must not exceed {max_length} characters
- Numeric fields must be within specified ranges
- Do not include fields not in the schema
- Do not wrap the JSON in markdown code blocks

Begin your response with the opening brace {{.
"""

def build_constrained_prompt(schema: dict, task: str) -> list[dict]:
    return [
        {"role": "system", "content": SCHEMA_ENFORCEMENT_PROMPT.format(
            schema=json.dumps(schema, indent=2),
            max_length=500
        )},
        {"role": "user", "content": task},
        {"role": "assistant", "content": "{"}  # Prime the generation
    ]

class FewShotSelector:
    def __init__(self, examples: list[dict], embedder):
        self.examples = examples
        self.embedder = embedder
        self.embeddings = [embedder.encode(ex["input"]) for ex in examples]

    def select(self, query: str, k: int = 3) -> list[dict]:
        query_emb = self.embedder.encode(query)
        similarities = [
            cosine_similarity(query_emb, emb) for emb in self.embeddings
        ]
        top_indices = sorted(
            range(len(similarities)),
            key=lambda i: similarities[i],
            reverse=True
        )[:k]
        return [self.examples[i] for i in top_indices]

    def format_examples(self, examples: list[dict]) -> str:
        parts = []
        for ex in examples:
            parts.append(f"Input: {ex['input']}\nOutput: {ex['output']}")
        return "\n\n---\n\n".join(parts)