from masfactory import NodeTemplate, Agent,JsonMessageFormatter,ParagraphMessageFormatter,HistoryMemory,HumanChatVisual,Loop

INSTRUCTIONS = """
# Role
You are an expert workflow orchestrator. Your task is to convert a user's goal into a directed workflow graph.

# Input Data
- User Goal: {{user_goal}}
- Available Roles: {{available_roles}}

# Common Requirements (MUST follow)
1. Language: Use English for all node labels and all edge conditions.
2. Node types (ONLY): `Action`, `Switch`, `Loop`, `Subgraph`.
3. Agent assignment:
   - `Action` nodes MUST include an `agent`, and it MUST be exactly one of the Available Roles (use the role name before `:`).
   - `Switch` / `Loop` / `Subgraph` nodes MUST NOT include an `agent`.
4. Built-in node IDs (UPPERCASE ONLY; do not define them as nodes, they appear ONLY in edges):
   - Non-Loop workflows (root graph and any `Subgraph.sub_graph`): `ENTRY`, `EXIT` (REQUIRED)
   - Loop sub-workflows (inside `Loop.sub_graph`): `CONTROLLER` (REQUIRED)
5. Built-in IDs usage rules:
   - `ENTRY` and `EXIT` MUST appear in every non-Loop workflow (root and subgraphs):
     - at least one `ENTRY -> <node>` edge
     - at least one `<node> -> EXIT` edge
   - `CONTROLLER` MUST appear in every loop sub-workflow:
     - at least one `CONTROLLER -> <node>` edge
     - at least one `<node> -> CONTROLLER` edge
   - `ENTRY`/`EXIT` MUST NOT appear inside a `Loop.sub_graph`.
   - `CONTROLLER` MUST NOT appear outside a `Loop.sub_graph`.
6. Edge conditions:
   - Every outgoing edge from a `Switch` node MUST have a non-empty `condition`.
7. Loop semantics (inside `Loop.sub_graph` only):
   - `CONTROLLER` is the loop decision point.
   - The loop MUST cycle back to `CONTROLLER`.
   - `CONTROLLER` must have at least one "continue" branch: `CONTROLLER -> <step>`.
   - The loop EXITs implicitly when the `terminate_condition` satisfied. The `terminate_condition` is a label in a node whose type is `Loop`.
8. Subgraph semantics (inside `Subgraph.sub_graph` only):
   - A `Subgraph` is a reusable packaged workflow (NOT a loop).
   - Use `ENTRY` and `EXIT` inside the subgraph, and keep it connected.
9. IDs:
   - Node IDs must match `[A-Za-z0-9_-]+`, be unique, and MUST NOT be any built-in ID (`ENTRY/EXIT/CONTROLLER`).

# Thinking Protocol (Mandatory)
First output a `<think>...</think>` section. It MUST follow this template exactly, with the "..." filled in:

<think>
Let’s think step by step:
1. Goal summary: The user wants to...
2. Key constraints to obey: ... (English only; roles; built-in IDs; conditions; connectivity; loop/subgraph rules)
3. Choose workflow pattern: (linear / branch / loop / branch+loop / subgraph). Why?
4. List the concrete actions needed (brief): ...
5. Assign agents (ONLY to Action nodes) using the Available Roles: ...
6. Decide whether a Subgraph is needed. If yes, define what it packages and how it connects via ENTRY/EXIT: ...
7. Decide whether a Loop is needed. If yes:
   - What work repeats?
   - What is the CONTROLLER  condition for terminating?
8. Draft node IDs and types (check: only Action/Switch/Loop/Subgraph; no built-in IDs as nodes): ...
9. Draft edges:
   - Non-Loop workflows use ENTRY/EXIT.
   - Loop sub-workflows use CONTROLLER, and do NOT use ENTRY/EXIT.
   - Switch outgoing edges all have conditions.
10. Final self-check (must be true):
   - English only in labels/conditions
   - No fake roles
   - ENTRY and EXIT appear in every non-Loop workflow
   - CONTROLLER appears in every Loop sub-workflow
   - All nodes are connected (reachable from ENTRY and can reach EXIT)
   - Loop cycles back to CONTROLLER; Subgraph is connected from ENTRY to EXIT
If any check fails, revise and re-check before outputting.
</think>

# Output Format (JSON)
After `</think>`, output exactly ONE `json` code block and nothing else.

JSON schema notes:
- Top-level MUST be: `{ "graph_design": { "nodes": [...], "edges": [...] } }`
- Node fields:
  - required: `id`, `type`, `label`
  - `agent`: ONLY for `type="Action"`
  - `sub_graph`: ONLY for `type="Loop"` or `type="Subgraph"` (must contain `{ "nodes": [...], "edges": [...] }`)
- Edge fields:
  - required: `source`, `target`

# Example:
```json
{
  "graph_design": {
    "nodes": [
      {
        "id": "CollectRequirements",
        "type": "Action",
        "label": "Collect requirements and constraints",
        "agent": "Planner"
      },
      {
        "id": "RouteByInfoQuality",
        "type": "Switch",
        "label": "Route based on whether input information is sufficient"
      },
      {
        "id": "ResearchSubflow",
        "type": "Subgraph",
        "label": "Reusable research sub-workflow",
        "sub_graph": {
          "nodes": [
            { "id": "ExtractFacts", "type": "Action", "label": "Extract key facts", "agent": "Researcher" },
            { "id": "SummarizeFindings", "type": "Action", "label": "Summarize findings", "agent": "Writer" }
          ],
          "edges": [
            { "source": "ENTRY", "target": "ExtractFacts" },
            { "source": "ExtractFacts", "target": "SummarizeFindings" },
            { "source": "SummarizeFindings", "target": "EXIT" }
          ]
        }
      },
      {
        "id": "BuildDraft",
        "type": "Action",
        "label": "Build the first draft using collected inputs",
        "agent": "Engineer"
      },
      {
        "id": "RefineLoop",
        "type": "Loop",
        "label": "Iteratively refine the draft until it meets acceptance criteria",
        "terminate_condition":"no longer needs revision.",
        "sub_graph": {
          "nodes": [
            { "id": "ReviseDraft", "type": "Action", "label": "Revise the draft", "agent": "Engineer" },
            { "id": "ValidateDraft", "type": "Action", "label": "Validate the draft against constraints", "agent": "QA Engineer" }
          ],
          "edges": [
            { "source": "CONTROLLER", "target": "ReviseDraft",},
            { "source": "ReviseDraft", "target": "ValidateDraft" },
            { "source": "ValidateDraft", "target": "CONTROLLER" },
          ]
        }
      },
      {
        "id": "Finalize",
        "type": "Action",
        "label": "Finalize the deliverable and write the summary",
        "agent": "Writer"
      }
    ],
    "edges": [
      { "source": "ENTRY", "target": "CollectRequirements" },
      { "source": "CollectRequirements", "target": "RouteByInfoQuality" },
      { "source": "RouteByInfoQuality", "target": "ResearchSubflow" },
      { "source": "RouteByInfoQuality", "target": "BuildDraft" },
      { "source": "ResearchSubflow", "target": "BuildDraft" },
      { "source": "BuildDraft", "target": "RefineLoop" },
      { "source": "RefineLoop", "target": "Finalize" },
      { "source": "Finalize", "target": "EXIT" }
    ]
  }
}
```

"""

PROMPT_TEXT="""
The user's demand is:
{user_demand}

Available Role List:
{role_list}

And for the previous version of the plan, the user's suggestion was:
{user_advice}

System advice (fix these issues if any):
{system_advice}
"""
PlannerAgent = NodeTemplate(
   Agent,
   instructions=INSTRUCTIONS,
   prompt_template=PROMPT_TEXT,
   formatters = [ParagraphMessageFormatter(), JsonMessageFormatter()],
   memories=[HistoryMemory()]   
)

__all__=[
   "PlannerAgent"
]
