import json, os
from smolagents import Tool

PLAN_PATH = "./execution_plan.json"

def load_plan():
    if not os.path.exists(PLAN_PATH):
        return []
    with open(PLAN_PATH, "r") as f:
        return json.load(f)

def save_plan(plan):
    with open(PLAN_PATH, "w") as f:
        json.dump(plan, f, indent=2)

class UpdatePlanTool(Tool):
    """
    Tool to update the execution plan stored on disk.
    
    Maintains a persistent execution plan that can be updated during multi-step operations.
    The plan is stored in JSON format at './execution_plan.json'.
    
    Usage:
    - Set merge=False to overwrite the entire plan
    - Set merge=True to update specific items by ID while preserving others

    Example input:
    {
      "merge": true,
      "todos": [
        {"id": "1", "content": "Research market trends", "status": "completed"},
        {"id": "2", "content": "Analyze competitor data", "status": "in_progress"}
      ]
    }
    

    Returns:
    - {"result": "plan overwritten", "plan": updated_plan} when merge=False
    - {"result": "plan updated", "plan": updated_plan} when merge=True
    - {"error": "...", "plan": current_plan} when invariant violated

    output_type: object
    
    """
    name = "update_plan"
    description = (
        "Update the current execution plan on disk. "
        "Use for multi-step tasks only. Maintains 2–5 non-operational items. "
        "Requires exactly one item in_progress during execution."
    )
    inputs = {
        "merge": {"type": "boolean", "description": "Merge with existing plan"},
        "todos": {"type": "array", "description": "List of objects {id, content, status}"}
    }

    output_type = "string"

    def forward(self, merge: bool, todos: list) -> str:
        plan = load_plan()

        if not merge:
            save_plan(todos)
            return json.dumps({"result": "plan overwritten", "plan": todos})

        # merge: update items by id
        plan_by_id = {item["id"]: item for item in plan}
        for new_item in todos:
            plan_by_id[new_item["id"]] = new_item

        merged_plan = list(plan_by_id.values())

        # enforce invariant: at most one in_progress
        in_prog = [t for t in merged_plan if t["status"] == "in_progress"]
        if len(in_prog) > 1:
            return {
                "error": "Only one item may be in_progress at a time. Fix your update.",
                "plan": merged_plan
            }

        save_plan(merged_plan)
        return json.dumps({"result": "plan updated", "plan": merged_plan})