Planning and Task Decomposition

Action: move(Block, From, To)
  Preconditions: on(Block, From), clear(Block), clear(To)
  Add effects:   on(Block, To), clear(From)
  Delete effects: on(Block, From), clear(To)

(define (domain logistics)
  (:requirements :strips)
  (:types city location package vehicle)

  (:action drive
    :parameters (?v - vehicle ?from - location ?to - location ?c - city)
    :precondition (and (in-city ?from ?c) (in-city ?to ?c) (at ?v ?from))
    :effect (and (at ?v ?to) (not (at ?v ?from)))
  )

  (:action load-package
    :parameters (?p - package ?v - vehicle ?l - location)
    :precondition (and (at ?v ?l) (at ?p ?l))
    :effect (and (in ?p ?v) (not (at ?p ?l)))
  )
)

PLANNING_PROMPT = """You are a planning agent. Given a goal, generate a step-by-step
plan to achieve it. Each step should be:
- Specific and actionable
- Ordered logically (dependencies respected)
- Achievable with the available tools

Goal: {goal}

Available tools:
{tools}

Generate a numbered plan:"""


def generate_plan(goal: str, tools: list[str], llm_call) -> list[str]:
    """Generate a natural language plan for a goal.

    Args:
        goal: The goal to achieve.
        tools: List of available tool descriptions.
        llm_call: Function to call the LLM.

    Returns:
        List of plan steps as strings.
    """
    tools_str = "\n".join(f"- {tool}" for tool in tools)
    prompt = PLANNING_PROMPT.format(goal=goal, tools=tools_str)
    response = llm_call(prompt=prompt)

    # Parse numbered steps
    steps = []
    for line in response.strip().split("\n"):
        line = line.strip()
        if line and line[0].isdigit():
            # Remove the number prefix
            step = line.lstrip("0123456789.)")
            steps.append(step.strip())

    return steps

class PlanAndExecuteAgent:
    """An agent that first plans, then executes step by step.

    This separates the planning phase (creative, high-level) from the
    execution phase (precise, tool-using), allowing each to be optimized
    independently.
    """

    def __init__(self, planner_llm, executor_llm, tools: dict[str, callable]):
        self.planner_llm = planner_llm
        self.executor_llm = executor_llm
        self.tools = tools

    def plan(self, goal: str) -> list[str]:
        """Generate a plan for the goal."""
        tools_desc = "\n".join(f"- {name}" for name in self.tools.keys())
        prompt = PLANNING_PROMPT.format(goal=goal, tools=tools_desc)
        response = self.planner_llm(prompt=prompt)
        return self._parse_steps(response)

    def execute_step(self, step: str, context: str) -> str:
        """Execute a single step of the plan.

        Args:
            step: The step to execute.
            context: Results from previous steps.

        Returns:
            The result of executing the step.
        """
        prompt = f"""Execute the following step of a plan.

Previous context:
{context}

Current step: {step}

Available tools: {list(self.tools.keys())}

Determine which tool to use (if any) and execute the step.
Respond with the result."""

        response = self.executor_llm(prompt=prompt)

        # Check if the executor wants to use a tool
        for tool_name, tool_fn in self.tools.items():
            if tool_name.lower() in response.lower():
                # Extract arguments and call the tool
                tool_result = tool_fn(step)
                return f"Tool '{tool_name}' result: {tool_result}"

        return response

    def run(self, goal: str) -> dict:
        """Plan and execute a goal.

        Returns:
            Dict with the plan, execution results, and final outcome.
        """
        # Phase 1: Plan
        plan = self.plan(goal)
        print(f"Plan generated with {len(plan)} steps:")
        for i, step in enumerate(plan, 1):
            print(f"  {i}. {step}")

        # Phase 2: Execute
        results = []
        context = ""
        for i, step in enumerate(plan):
            print(f"\nExecuting step {i + 1}: {step}")
            result = self.execute_step(step, context)
            results.append({"step": step, "result": result})
            context += f"\nStep {i + 1} ({step}): {result}"
            print(f"  Result: {result[:200]}")

        return {
            "goal": goal,
            "plan": plan,
            "results": results,
            "final_context": context,
        }

    def _parse_steps(self, response: str) -> list[str]:
        """Parse numbered steps from LLM response."""
        steps = []
        for line in response.strip().split("\n"):
            line = line.strip()
            if line and (line[0].isdigit() or line.startswith("-")):
                step = line.lstrip("0123456789.)- ")
                if step:
                    steps.append(step)
        return steps

Compound Task: prepare_dinner
  Method 1: home_cooking
    Subtasks: [choose_recipe, buy_ingredients, cook, serve]
  Method 2: order_delivery
    Subtasks: [choose_restaurant, place_order, wait, receive]

DECOMPOSITION_PROMPT = """Break down the following task into 3-7 subtasks.
Each subtask should be:
- More specific than the parent task
- Independent enough to be worked on separately
- Small enough to be completed in a focused work session

If a subtask is still complex, mark it with [COMPLEX] so it can be
further decomposed.

Task: {task}

Subtasks:"""


def hierarchical_decompose(
    task: str, llm_call, max_depth: int = 3, current_depth: int = 0
) -> dict:
    """Recursively decompose a task into a hierarchy of subtasks.

    Args:
        task: The task to decompose.
        llm_call: Function to call the LLM.
        max_depth: Maximum decomposition depth.
        current_depth: Current depth in the recursion.

    Returns:
        A tree structure representing the task hierarchy.
    """
    if current_depth >= max_depth:
        return {"task": task, "subtasks": [], "leaf": True}

    response = llm_call(prompt=DECOMPOSITION_PROMPT.format(task=task))
    subtasks = _parse_subtasks(response)

    tree = {"task": task, "subtasks": [], "leaf": False}

    for subtask_text in subtasks:
        is_complex = "[COMPLEX]" in subtask_text
        clean_text = subtask_text.replace("[COMPLEX]", "").strip()

        if is_complex and current_depth < max_depth - 1:
            # Recursively decompose complex subtasks
            subtree = hierarchical_decompose(
                clean_text, llm_call, max_depth, current_depth + 1
            )
            tree["subtasks"].append(subtree)
        else:
            tree["subtasks"].append({
                "task": clean_text,
                "subtasks": [],
                "leaf": True,
            })

    return tree


def _parse_subtasks(response: str) -> list[str]:
    """Parse subtasks from LLM response."""
    subtasks = []
    for line in response.strip().split("\n"):
        line = line.strip()
        if line and (line[0].isdigit() or line.startswith("-")):
            subtask = line.lstrip("0123456789.)- ")
            if subtask:
                subtasks.append(subtask)
    return subtasks


def print_task_tree(tree: dict, indent: int = 0) -> None:
    """Pretty-print a task tree."""
    prefix = "  " * indent
    marker = "[leaf]" if tree["leaf"] else "[node]"
    print(f"{prefix}{marker} {tree['task']}")
    for subtask in tree.get("subtasks", []):
        print_task_tree(subtask, indent + 1)

VALIDATION_PROMPT = """Review the following plan for achieving the stated goal.
Check for:
1. Missing steps: Are there any prerequisites or necessary actions not included?
2. Ordering errors: Are steps in the right order? Are dependencies respected?
3. Feasibility: Can each step be executed with the available tools?
4. Redundancy: Are there unnecessary or duplicate steps?
5. Specificity: Is each step specific enough to be actionable?

Goal: {goal}

Plan:
{plan}

Available tools: {tools}

Provide your assessment, listing any issues found and suggesting corrections."""


def validate_plan(
    goal: str, plan: list[str], tools: list[str], llm_call
) -> dict:
    """Validate a plan and identify issues.

    Returns:
        Dict with validation results and suggested fixes.
    """
    plan_str = "\n".join(f"{i+1}. {step}" for i, step in enumerate(plan))
    tools_str = ", ".join(tools)

    response = llm_call(
        prompt=VALIDATION_PROMPT.format(goal=goal, plan=plan_str, tools=tools_str)
    )

    return {
        "original_plan": plan,
        "validation_feedback": response,
        "needs_revision": any(
            keyword in response.lower()
            for keyword in ["missing", "error", "incorrect", "should be", "add"]
        ),
    }

CRITIC_PROMPT = """You are a plan reviewer. Your job is to find flaws in plans
before they are executed. Be thorough and skeptical.

Goal: {goal}
Plan:
{plan}

For each step, assess:
- Is this step necessary?
- Is it in the correct position?
- What could go wrong?
- Is it specific enough?

Rate the overall plan quality (1-10) and list all issues."""

def simulate_plan_execution(
    plan: list[str], tools: dict[str, callable], dry_run: bool = True
) -> list[dict]:
    """Simulate plan execution to catch errors before real execution.

    Args:
        plan: List of plan steps.
        tools: Available tools.
        dry_run: If True, tools log what they would do without executing.

    Returns:
        List of simulation results for each step.
    """
    results = []
    state = {}

    for i, step in enumerate(plan):
        result = {
            "step": i + 1,
            "description": step,
            "status": "unknown",
            "issues": [],
        }

        # Check if the step references a known tool
        tool_found = False
        for tool_name in tools:
            if tool_name.lower() in step.lower():
                tool_found = True
                if dry_run:
                    result["status"] = "simulated"
                    result["output"] = f"Would execute: {tool_name}"
                else:
                    try:
                        output = tools[tool_name](step)
                        result["status"] = "success"
                        result["output"] = output
                    except Exception as e:
                        result["status"] = "error"
                        result["issues"].append(str(e))
                break

        if not tool_found:
            result["status"] = "no_tool"
            result["issues"].append("No matching tool found for this step.")

        results.append(result)

    return results

REFINEMENT_PROMPT = """Revise the following plan based on the feedback provided.
Make minimal changes to fix the identified issues while preserving the parts
that are correct.

Original plan:
{plan}

Feedback:
{feedback}

Revised plan:"""


def refine_plan(
    plan: list[str], feedback: str, llm_call
) -> list[str]:
    """Refine a plan based on validation feedback."""
    plan_str = "\n".join(f"{i+1}. {step}" for i, step in enumerate(plan))
    response = llm_call(
        prompt=REFINEMENT_PROMPT.format(plan=plan_str, feedback=feedback)
    )
    return _parse_subtasks(response)

def iterative_planning(
    goal: str, tools: list[str], llm_call, max_iterations: int = 3
) -> list[str]:
    """Generate and iteratively refine a plan.

    Args:
        goal: The goal to plan for.
        tools: Available tools.
        llm_call: Function to call the LLM.
        max_iterations: Maximum refinement iterations.

    Returns:
        The final refined plan.
    """
    # Generate initial plan
    plan = generate_plan(goal, tools, llm_call)
    print(f"Initial plan: {len(plan)} steps")

    for iteration in range(max_iterations):
        # Validate
        validation = validate_plan(goal, plan, tools, llm_call)

        if not validation["needs_revision"]:
            print(f"Plan validated after {iteration + 1} iteration(s).")
            return plan

        # Refine
        print(f"Iteration {iteration + 1}: Refining plan based on feedback...")
        plan = refine_plan(plan, validation["validation_feedback"], llm_call)

    print(f"Reached max iterations ({max_iterations}). Returning best plan.")
    return plan

class DynamicPlanExecutor:
    """Execute a plan with dynamic replanning on failure.

    When a step fails, the executor asks the LLM to generate an
    alternative approach or revise the remaining plan.
    """

    def __init__(self, llm_call, tools: dict[str, callable]):
        self.llm_call = llm_call
        self.tools = tools

    def execute_with_replanning(
        self, goal: str, initial_plan: list[str], max_replans: int = 3
    ) -> dict:
        """Execute a plan, replanning when steps fail.

        Args:
            goal: The original goal.
            initial_plan: The initial plan.
            max_replans: Maximum number of times to replan.

        Returns:
            Execution results including any replanning events.
        """
        plan = list(initial_plan)
        execution_log = []
        replans = 0
        step_idx = 0

        while step_idx < len(plan):
            step = plan[step_idx]
            print(f"Executing step {step_idx + 1}/{len(plan)}: {step}")

            try:
                result = self._execute_step(step)
                execution_log.append({
                    "step": step,
                    "status": "success",
                    "result": result,
                })
                print(f"  Success: {result[:100]}")
                step_idx += 1

            except Exception as e:
                error_msg = str(e)
                print(f"  Failed: {error_msg}")
                execution_log.append({
                    "step": step,
                    "status": "failed",
                    "error": error_msg,
                })

                if replans >= max_replans:
                    print("  Max replans reached. Aborting.")
                    break

                # Replan
                print("  Replanning...")
                remaining_steps = plan[step_idx:]
                completed_context = "\n".join(
                    f"- {log['step']}: {log.get('result', log.get('error', 'N/A'))[:100]}"
                    for log in execution_log
                )

                new_plan = self._replan(
                    goal, completed_context, remaining_steps, error_msg
                )

                if new_plan:
                    plan = plan[:step_idx] + new_plan
                    replans += 1
                    print(f"  New plan ({len(new_plan)} steps remaining)")
                else:
                    print("  Replanning failed. Skipping step.")
                    step_idx += 1

        return {
            "goal": goal,
            "execution_log": execution_log,
            "replans": replans,
            "completed": step_idx >= len(plan),
        }

    def _execute_step(self, step: str) -> str:
        """Execute a single step using available tools."""
        for tool_name, tool_fn in self.tools.items():
            if tool_name.lower() in step.lower():
                return tool_fn(step)
        # If no tool matches, use the LLM to reason about the step
        return self.llm_call(prompt=f"Execute this step and report the result: {step}")

    def _replan(
        self, goal: str, completed: str, failed_steps: list[str], error: str
    ) -> list[str] | None:
        """Generate a new plan given a failure.

        Args:
            goal: The original goal.
            completed: Summary of completed steps.
            failed_steps: The steps that failed or remain.
            error: The error that triggered replanning.

        Returns:
            A new list of steps, or None if replanning fails.
        """
        prompt = f"""A plan step failed during execution. Generate a revised plan
to achieve the goal, taking into account what has already been completed
and the error encountered.

Goal: {goal}

Completed steps:
{completed}

Failed step: {failed_steps[0] if failed_steps else 'N/A'}
Error: {error}

Remaining original steps:
{chr(10).join(f'- {s}' for s in failed_steps[1:])}

Generate a revised plan for the remaining work (avoid repeating completed steps):"""

        response = self.llm_call(prompt=prompt)
        steps = _parse_subtasks(response)
        return steps if steps else None

from dataclasses import dataclass, field
from enum import Enum


class TaskStatus(Enum):
    PENDING = "pending"
    READY = "ready"      # All dependencies met
    RUNNING = "running"
    COMPLETED = "completed"
    FAILED = "failed"
    SKIPPED = "skipped"


@dataclass
class TaskNode:
    """A node in the task graph."""

    id: str
    description: str
    dependencies: list[str] = field(default_factory=list)  # IDs of prerequisite tasks
    status: TaskStatus = TaskStatus.PENDING
    result: str | None = None
    error: str | None = None

    def is_ready(self, completed_tasks: set[str]) -> bool:
        """Check if all dependencies are satisfied."""
        return all(dep in completed_tasks for dep in self.dependencies)


class TaskGraph:
    """A directed acyclic graph (DAG) of tasks with dependency management.

    Supports parallel execution of independent tasks and respects
    ordering constraints.
    """

    def __init__(self):
        self.nodes: dict[str, TaskNode] = {}

    def add_task(
        self, task_id: str, description: str, dependencies: list[str] | None = None
    ) -> None:
        """Add a task to the graph."""
        deps = dependencies or []
        # Validate dependencies exist
        for dep in deps:
            if dep not in self.nodes:
                raise ValueError(f"Dependency '{dep}' not found. Add it first.")
        self.nodes[task_id] = TaskNode(
            id=task_id, description=description, dependencies=deps
        )

    def get_ready_tasks(self) -> list[TaskNode]:
        """Get all tasks that are ready to execute (dependencies satisfied)."""
        completed = {
            node_id for node_id, node in self.nodes.items()
            if node.status == TaskStatus.COMPLETED
        }
        ready = []
        for node in self.nodes.values():
            if node.status == TaskStatus.PENDING and node.is_ready(completed):
                node.status = TaskStatus.READY
                ready.append(node)
        return ready

    def mark_completed(self, task_id: str, result: str) -> None:
        """Mark a task as completed with its result."""
        self.nodes[task_id].status = TaskStatus.COMPLETED
        self.nodes[task_id].result = result

    def mark_failed(self, task_id: str, error: str) -> None:
        """Mark a task as failed."""
        self.nodes[task_id].status = TaskStatus.FAILED
        self.nodes[task_id].error = error
        # Also skip dependent tasks
        self._propagate_failure(task_id)

    def _propagate_failure(self, failed_id: str) -> None:
        """Skip all tasks that depend on a failed task."""
        for node in self.nodes.values():
            if failed_id in node.dependencies and node.status == TaskStatus.PENDING:
                node.status = TaskStatus.SKIPPED
                node.error = f"Skipped because dependency '{failed_id}' failed."
                self._propagate_failure(node.id)

    def is_complete(self) -> bool:
        """Check if all tasks are completed (or skipped/failed)."""
        return all(
            node.status in (TaskStatus.COMPLETED, TaskStatus.FAILED, TaskStatus.SKIPPED)
            for node in self.nodes.values()
        )

    def get_execution_order(self) -> list[list[str]]:
        """Get the topological execution order as batches of parallel tasks.

        Returns:
            List of batches, where each batch contains tasks that can run in parallel.
        """
        completed: set[str] = set()
        remaining = set(self.nodes.keys())
        batches = []

        while remaining:
            batch = []
            for task_id in remaining:
                node = self.nodes[task_id]
                if all(dep in completed for dep in node.dependencies):
                    batch.append(task_id)

            if not batch:
                raise ValueError("Circular dependency detected!")

            batches.append(batch)
            completed.update(batch)
            remaining -= set(batch)

        return batches

    def visualize(self) -> str:
        """Create a text visualization of the task graph."""
        lines = ["Task Graph:"]
        batches = self.get_execution_order()
        for i, batch in enumerate(batches):
            lines.append(f"\n  Batch {i + 1} (parallel):")
            for task_id in batch:
                node = self.nodes[task_id]
                deps = f" (depends on: {', '.join(node.dependencies)})" if node.dependencies else ""
                status = f" [{node.status.value}]"
                lines.append(f"    - {task_id}: {node.description}{deps}{status}")
        return "\n".join(lines)


# Example usage
def build_example_graph() -> TaskGraph:
    """Build an example task graph for a data analysis project."""
    graph = TaskGraph()

    # Independent tasks (can run in parallel)
    graph.add_task("fetch_data", "Download dataset from API")
    graph.add_task("setup_env", "Set up Python virtual environment and install dependencies")

    # Depends on fetch_data
    graph.add_task("clean_data", "Clean and preprocess the raw data", ["fetch_data"])
    graph.add_task("explore_data", "Perform exploratory data analysis", ["fetch_data"])

    # Depends on clean_data and setup_env
    graph.add_task("train_model", "Train the ML model", ["clean_data", "setup_env"])

    # Depends on train_model and explore_data
    graph.add_task("evaluate", "Evaluate model on test set", ["train_model"])
    graph.add_task("generate_report", "Generate analysis report", ["evaluate", "explore_data"])

    return graph

TASK_GRAPH_PROMPT = """Analyze the following goal and generate a task graph.
For each task, identify:
1. A unique short ID
2. A description
3. Which other tasks it depends on (by ID)

Output format (one task per line):
ID | Description | Dependencies (comma-separated IDs, or "none")

Goal: {goal}

Tasks:"""


def generate_task_graph(goal: str, llm_call) -> TaskGraph:
    """Generate a task graph from a natural language goal."""
    response = llm_call(prompt=TASK_GRAPH_PROMPT.format(goal=goal))

    graph = TaskGraph()
    tasks_to_add = []

    # Parse the response
    for line in response.strip().split("\n"):
        parts = [p.strip() for p in line.split("|")]
        if len(parts) >= 3:
            task_id = parts[0].lower().replace(" ", "_")
            description = parts[1]
            deps_str = parts[2]
            deps = (
                [] if deps_str.lower() == "none"
                else [d.strip().lower().replace(" ", "_") for d in deps_str.split(",")]
            )
            tasks_to_add.append((task_id, description, deps))

    # Add tasks in dependency order
    added = set()
    max_passes = len(tasks_to_add) * 2
    passes = 0
    while tasks_to_add and passes < max_passes:
        remaining = []
        for task_id, description, deps in tasks_to_add:
            if all(d in added for d in deps):
                graph.add_task(task_id, description, deps)
                added.add(task_id)
            else:
                remaining.append((task_id, description, deps))
        tasks_to_add = remaining
        passes += 1

    return graph

def iterative_refinement(
    goal: str, llm_call, refinement_rounds: int = 3
) -> list[str]:
    """Iteratively refine a plan through multiple passes.

    Round 1: Generate a high-level outline (3-5 major phases)
    Round 2: Expand each phase into specific steps
    Round 3: Add details, error handling, and validation steps
    """
    # Round 1: High-level outline
    outline = llm_call(
        prompt=f"Create a high-level outline (3-5 phases) for: {goal}"
    )

    # Round 2: Expand each phase
    expanded = llm_call(
        prompt=(
            f"Expand each phase of this outline into 2-4 specific, actionable steps.\n\n"
            f"Outline:\n{outline}\n\nExpanded plan:"
        )
    )

    # Round 3: Add details
    detailed = llm_call(
        prompt=(
            f"Review this plan and add any missing steps, error handling, "
            f"and validation checkpoints.\n\n"
            f"Plan:\n{expanded}\n\nDetailed plan:"
        )
    )

    return _parse_subtasks(detailed)

Complex problem: "Calculate the total cost of a round trip from Madrid to Tokyo,
including flights, hotels for 5 nights, and daily meals."

Decomposition:
1. What is the average cost of a round-trip flight from Madrid to Tokyo?
2. What is the average cost of a hotel in Tokyo per night?
3. What is the average daily meal budget in Tokyo?
4. Given answers to 1-3, what is the total cost?

CONTINGENCY_PROMPT = """Generate a plan for the following goal, including
contingency branches for steps that might fail.

For each step, specify:
- The primary action
- What could go wrong
- The fallback action if it fails

Goal: {goal}

Plan with contingencies:"""


@dataclass
class ContingencyStep:
    """A plan step with contingency handling."""

    primary_action: str
    potential_failure: str
    fallback_action: str

    def to_text(self) -> str:
        return (
            f"DO: {self.primary_action}\n"
            f"  IF FAILS ({self.potential_failure}): {self.fallback_action}"
        )

Belief: "The file is probably in /data/ (80%) or /backup/ (20%)"

Plan:
1. Check /data/ for the file.
2. IF found: proceed with processing.
3. IF not found: check /backup/.
4. IF still not found: ask the user for the file location.

def tree_of_thoughts(
    problem: str,
    llm_call,
    num_thoughts: int = 3,
    max_depth: int = 4,
    beam_width: int = 2,
) -> str:
    """Simplified Tree of Thoughts implementation.

    Uses beam search over multiple reasoning paths.

    Args:
        problem: The problem to solve.
        llm_call: Function to call the LLM.
        num_thoughts: Number of candidate thoughts to generate at each step.
        max_depth: Maximum depth of the thought tree.
        beam_width: Number of paths to keep at each step.

    Returns:
        The best solution found.
    """
    # Initialize with the problem
    beams = [{"path": [], "text": problem, "score": 0.0}]

    for depth in range(max_depth):
        all_candidates = []

        for beam in beams:
            # Generate candidate next thoughts
            thoughts = llm_call(
                prompt=(
                    f"Problem: {problem}\n\n"
                    f"Progress so far:\n{beam['text']}\n\n"
                    f"Generate {num_thoughts} different possible next steps. "
                    f"Number each option."
                )
            )

            # Parse and evaluate each thought
            for thought in _parse_thoughts(thoughts):
                # Evaluate this partial solution
                evaluation = llm_call(
                    prompt=(
                        f"Rate how promising this partial solution is (1-10):\n\n"
                        f"Problem: {problem}\n"
                        f"Progress: {beam['text']}\n"
                        f"Next step: {thought}\n\n"
                        f"Score (just the number):"
                    )
                )

                try:
                    score = float(evaluation.strip().split()[0])
                except (ValueError, IndexError):
                    score = 5.0

                all_candidates.append({
                    "path": beam["path"] + [thought],
                    "text": f"{beam['text']}\n{thought}",
                    "score": score,
                })

        # Keep top beam_width candidates
        all_candidates.sort(key=lambda x: x["score"], reverse=True)
        beams = all_candidates[:beam_width]

    # Return the best path
    return beams[0]["text"] if beams else "No solution found."


def _parse_thoughts(response: str) -> list[str]:
    """Parse numbered thoughts from LLM response."""
    thoughts = []
    for line in response.strip().split("\n"):
        line = line.strip()
        if line and line[0].isdigit():
            thought = line.lstrip("0123456789.)- ").strip()
            if thought:
                thoughts.append(thought)
    return thoughts

"""
A task decomposition agent that:
1. Takes a complex goal
2. Decomposes it into subtasks with dependencies
3. Builds a task graph
4. Executes tasks in dependency order
5. Replans on failure

Requirements:
    pip install openai  (or any LLM client)
"""

from dataclasses import dataclass, field
from datetime import datetime


@dataclass
class SubTask:
    """A subtask produced by decomposition."""

    id: str
    description: str
    dependencies: list[str] = field(default_factory=list)
    estimated_complexity: str = "medium"  # low, medium, high
    status: str = "pending"
    result: str | None = None


DECOMPOSE_SYSTEM = """You are a task decomposition expert. Given a complex goal,
break it down into concrete, actionable subtasks.

Rules:
1. Each subtask should be completable in a single focused step.
2. Identify dependencies between subtasks explicitly.
3. Rate each subtask's complexity (low/medium/high).
4. Aim for 4-8 subtasks (not too few, not too many).
5. Make subtasks specific enough to execute without further clarification.

Output format (one task per line):
TASK_ID | Description | Dependencies (comma-separated IDs or "none") | Complexity
"""

EXECUTE_SYSTEM = """You are a task execution agent. Given a subtask and context
from previously completed subtasks, execute the task and report the result.

Be specific and thorough in your execution. If you produce any output
(text, code, analysis), include it in your response.

If you cannot complete the task, explain why specifically."""


class TaskDecompositionAgent:
    """An agent that decomposes and executes complex tasks."""

    def __init__(self, llm_call):
        self.llm_call = llm_call
        self.subtasks: dict[str, SubTask] = {}
        self.execution_log: list[dict] = []

    def decompose(self, goal: str) -> list[SubTask]:
        """Decompose a goal into subtasks.

        Args:
            goal: The complex goal to decompose.

        Returns:
            List of SubTask objects.
        """
        response = self.llm_call(
            system=DECOMPOSE_SYSTEM,
            prompt=f"Decompose this goal into subtasks:\n\n{goal}",
        )

        subtasks = self._parse_decomposition(response)
        for task in subtasks:
            self.subtasks[task.id] = task

        return subtasks

    def execute(self, goal: str, max_replans: int = 2) -> dict:
        """Decompose and execute a goal.

        Args:
            goal: The goal to achieve.
            max_replans: Maximum replanning attempts on failure.

        Returns:
            Execution report.
        """
        # Step 1: Decompose
        print(f"Goal: {goal}")
        print(f"{'='*60}")
        subtasks = self.decompose(goal)
        print(f"\nDecomposed into {len(subtasks)} subtasks:")
        for task in subtasks:
            deps = f" (after: {', '.join(task.dependencies)})" if task.dependencies else ""
            print(f"  [{task.id}] {task.description}{deps} [{task.estimated_complexity}]")

        # Step 2: Execute in dependency order
        print(f"\n{'='*60}")
        print("Executing...\n")
        completed = set()
        replans = 0

        while True:
            # Find ready tasks
            ready = [
                task for task in self.subtasks.values()
                if task.status == "pending"
                and all(dep in completed for dep in task.dependencies)
            ]

            if not ready:
                if all(t.status in ("completed", "failed", "skipped") for t in self.subtasks.values()):
                    break
                else:
                    print("Deadlock: no tasks are ready but some are still pending.")
                    break

            for task in ready:
                print(f"Executing [{task.id}]: {task.description}")

                # Build context from completed tasks
                context = self._build_context(completed)

                try:
                    result = self._execute_subtask(task, context)
                    task.status = "completed"
                    task.result = result
                    completed.add(task.id)
                    self.execution_log.append({
                        "task_id": task.id,
                        "status": "completed",
                        "result": result[:200],
                        "timestamp": datetime.now().isoformat(),
                    })
                    print(f"  Completed: {result[:150]}...")

                except Exception as e:
                    error = str(e)
                    print(f"  Failed: {error}")
                    task.status = "failed"
                    self.execution_log.append({
                        "task_id": task.id,
                        "status": "failed",
                        "error": error,
                        "timestamp": datetime.now().isoformat(),
                    })

                    # Replan if possible
                    if replans < max_replans:
                        print("  Attempting to replan...")
                        new_task = self._replan_task(task, error, context)
                        if new_task:
                            self.subtasks[new_task.id] = new_task
                            replans += 1
                            print(f"  Created replacement task [{new_task.id}]: {new_task.description}")

        # Step 3: Generate report
        report = self._generate_report(goal)
        return report

    def _execute_subtask(self, task: SubTask, context: str) -> str:
        """Execute a single subtask."""
        prompt = (
            f"Execute the following subtask:\n\n"
            f"Task: {task.description}\n\n"
            f"Context from previous steps:\n{context}\n\n"
            f"Execute the task and provide the result:"
        )
        return self.llm_call(system=EXECUTE_SYSTEM, prompt=prompt)

    def _build_context(self, completed_ids: set[str]) -> str:
        """Build context string from completed tasks."""
        if not completed_ids:
            return "(No previous steps completed yet.)"

        lines = []
        for task_id in completed_ids:
            task = self.subtasks[task_id]
            result_preview = task.result[:300] if task.result else "N/A"
            lines.append(f"[{task_id}] {task.description}\n  Result: {result_preview}")
        return "\n".join(lines)

    def _replan_task(self, failed_task: SubTask, error: str, context: str) -> SubTask | None:
        """Generate a replacement task for a failed one."""
        prompt = (
            f"The following task failed:\n"
            f"Task: {failed_task.description}\n"
            f"Error: {error}\n\n"
            f"Context:\n{context}\n\n"
            f"Generate an alternative approach to accomplish the same objective. "
            f"Respond with a single task description."
        )
        response = self.llm_call(prompt=prompt)
        if response.strip():
            new_id = f"{failed_task.id}_retry"
            return SubTask(
                id=new_id,
                description=response.strip(),
                dependencies=failed_task.dependencies,
                estimated_complexity=failed_task.estimated_complexity,
            )
        return None

    def _parse_decomposition(self, response: str) -> list[SubTask]:
        """Parse the LLM's decomposition into SubTask objects."""
        subtasks = []
        for line in response.strip().split("\n"):
            parts = [p.strip() for p in line.split("|")]
            if len(parts) >= 3:
                task_id = parts[0].lower().replace(" ", "_")
                description = parts[1]
                deps_str = parts[2]
                deps = (
                    [] if deps_str.lower() == "none"
                    else [d.strip().lower().replace(" ", "_") for d in deps_str.split(",")]
                )
                complexity = parts[3].strip().lower() if len(parts) > 3 else "medium"
                subtasks.append(SubTask(
                    id=task_id,
                    description=description,
                    dependencies=deps,
                    estimated_complexity=complexity,
                ))
        return subtasks

    def _generate_report(self, goal: str) -> dict:
        """Generate a final execution report."""
        completed = [t for t in self.subtasks.values() if t.status == "completed"]
        failed = [t for t in self.subtasks.values() if t.status == "failed"]
        skipped = [t for t in self.subtasks.values() if t.status == "skipped"]

        return {
            "goal": goal,
            "total_subtasks": len(self.subtasks),
            "completed": len(completed),
            "failed": len(failed),
            "skipped": len(skipped),
            "success_rate": len(completed) / max(len(self.subtasks), 1),
            "execution_log": self.execution_log,
            "subtask_results": {
                task_id: {
                    "description": task.description,
                    "status": task.status,
                    "result_preview": (task.result[:200] if task.result else None),
                }
                for task_id, task in self.subtasks.items()
            },
        }


# ── Usage Example ─────────────────────────────────────────────────

def main():
    """Demonstrate the task decomposition agent."""

    def mock_llm_call(system: str = "", prompt: str = "") -> str:
        """Simulate LLM calls for demonstration."""
        if "Decompose" in prompt:
            return """t1 | Research existing sentiment analysis approaches | none | low
t2 | Select and download a suitable dataset | none | low
t3 | Set up Python environment with required libraries | none | low
t4 | Implement data preprocessing pipeline | t2, t3 | medium
t5 | Train a baseline sentiment classifier | t4 | medium
t6 | Evaluate model and analyze errors | t5 | medium
t7 | Write up results and conclusions | t6, t1 | low"""
        elif "Execute" in prompt or "Execute" in system:
            if "research" in prompt.lower():
                return (
                    "Surveyed recent sentiment analysis approaches. Key findings: "
                    "BERT-based models achieve ~95% accuracy on SST-2. DistilBERT "
                    "offers a good speed/accuracy trade-off. Recent work explores "
                    "aspect-based sentiment analysis for finer-grained results."
                )
            elif "dataset" in prompt.lower():
                return (
                    "Selected the Stanford Sentiment Treebank (SST-2) dataset. "
                    "Downloaded 67,349 training samples and 872 test samples. "
                    "Binary sentiment labels (positive/negative)."
                )
            elif "environment" in prompt.lower():
                return (
                    "Created virtual environment. Installed: torch 2.1, "
                    "transformers 4.36, datasets 2.16, scikit-learn 1.3."
                )
            elif "preprocessing" in prompt.lower():
                return (
                    "Implemented preprocessing: tokenization with BertTokenizer, "
                    "max_length=128, padding and truncation. Created DataLoader "
                    "with batch_size=32."
                )
            elif "train" in prompt.lower():
                return (
                    "Trained DistilBERT for 3 epochs. Training loss decreased from "
                    "0.42 to 0.12. Validation accuracy: 91.2%."
                )
            elif "evaluate" in prompt.lower():
                return (
                    "Test accuracy: 90.8%. Precision: 0.91, Recall: 0.90, F1: 0.905. "
                    "Error analysis: model struggles with sarcasm and negation."
                )
            elif "write" in prompt.lower():
                return (
                    "Report completed. Key findings: DistilBERT achieves 90.8% accuracy "
                    "on SST-2 with 3 epochs of fine-tuning. Main failure modes are "
                    "sarcasm detection and complex negation patterns."
                )
        return "Task completed successfully."

    # Create and run the agent
    agent = TaskDecompositionAgent(mock_llm_call)
    report = agent.execute(
        "Build a sentiment analysis model: research approaches, train a model, "
        "evaluate it, and write up the results."
    )

    print(f"\n{'='*60}")
    print("EXECUTION REPORT")
    print(f"{'='*60}")
    print(f"Goal: {report['goal']}")
    print(f"Subtasks: {report['total_subtasks']}")
    print(f"Completed: {report['completed']}")
    print(f"Failed: {report['failed']}")
    print(f"Success rate: {report['success_rate']:.0%}")


if __name__ == "__main__":
    main()