ptc_runner

defmodule PtcRunner.Plan do
  @moduledoc """
  Parsed execution plan for multi-agent workflows.

  Plans are generated by `PtcRunner.MetaPlanner` as JSON and parsed into
  a normalized structure. This module handles parsing, validation, and
  sanitization of LLM-generated plans.

  ## LLM-Generated JSON Structure

  The MetaPlanner produces JSON with this structure:

      %{
        "tasks" => [
          %{
            "id" => "fetch_data",
            "agent" => "researcher",
            "input" => "Find quarterly revenue",
            "depends_on" => ["prior_task_id"],
            "output" => "ptc_lisp",
            "signature" => "{revenue :float}",
            "verification" => "(number? (get data/result \"revenue\"))",
            "on_verification_failure" => "retry",
            "on_failure" => "stop",
            "max_retries" => 1,
            "critical" => true,
            "type" => "task"
          }
        ],
        "agents" => %{
          "researcher" => %{
            "prompt" => "You are a researcher...",
            "tools" => ["search"]
          }
        }
      }

  ## Task Fields

  | Field | Type | Default | Description |
  |-------|------|---------|-------------|
  | `id` | string | `"task_N"` | Unique task identifier |
  | `agent` | string | `"default"` | Agent to execute the task |
  | `input` | string | `""` | Task description or PTC-Lisp code (for `"direct"` agent) |
  | `depends_on` | list | `[]` | IDs of upstream tasks (also: `requires`, `after`) |
  | `output` | string | `nil` | `"ptc_lisp"`, `"json"`, or nil (auto-detect) |
  | `signature` | string | `nil` | Output signature (e.g., `"{name :string}"`) |
  | `verification` | string | `nil` | PTC-Lisp predicate to validate output |
  | `on_verification_failure` | string | `"replan"` | `"stop"`, `"skip"`, `"retry"`, or `"replan"` |
  | `on_failure` | string | `"stop"` | `"stop"`, `"skip"`, `"retry"`, or `"replan"` |
  | `max_retries` | integer | `1` | Max retry attempts |
  | `critical` | boolean | `true` | Whether failure stops the plan |
  | `type` | string | `"task"` | `"task"`, `"synthesis_gate"`, or `"human_review"` |

  ## Parsing

  `parse/1` handles common LLM variations:
  - Task keys: `tasks`, `steps`, `workflow`, `plan.steps`
  - Agent keys: `agents`, `workers`
  - Dependency keys: `depends_on`, `requires`, `after`

  ## Sanitization

  `sanitize/1` validates verification predicates using `PtcRunner.Lisp.validate/1`.
  Invalid predicates (undefined variables, parse errors, arity mistakes) are
  stripped with a warning rather than failing the plan.
  """

  @type agent_spec :: %{
          prompt: String.t(),
          tools: [String.t()]
        }

  @type on_failure :: :stop | :skip | :retry | :replan
  @type on_verification_failure :: :stop | :skip | :retry | :replan

  alias PtcRunner.Lisp
  alias PtcRunner.SubAgent.Signature

  @type task_type :: :task | :synthesis_gate | :human_review

  @type output_mode :: :ptc_lisp | :text | nil

  @type task :: %{
          id: String.t(),
          agent: String.t(),
          input: term(),
          depends_on: [String.t()],
          on_failure: on_failure(),
          max_retries: non_neg_integer(),
          critical: boolean(),
          type: task_type(),
          # Output mode: :ptc_lisp, :text, or nil (auto-detect based on tools)
          output: output_mode(),
          # Output signature for text mode (e.g., "{stocks [{symbol :string, price :float}]}")
          signature: String.t() | nil,
          # Verification (Phase 1)
          verification: String.t() | nil,
          on_verification_failure: on_verification_failure(),
          # Per-task quality gate override (nil = use global setting)
          quality_gate: boolean() | nil
        }

  @type t :: %__MODULE__{
          agents: %{String.t() => agent_spec()},
          tasks: [task()]
        }

  defstruct agents: %{}, tasks: []

  @doc """
  Parse an LLM-generated plan into a normalized structure.

  Handles common variations in how LLMs structure plans:
  - Tasks can be under `tasks`, `steps`, `workflow`, or `plan.steps`
  - Agents can be under `agents` or `workers`
  - Dependencies can be `depends_on`, `requires`, or `after`

  ## Examples

      iex> raw = %{"tasks" => [%{"id" => "t1", "agent" => "researcher", "input" => "test"}]}
      iex> {:ok, plan} = PtcRunner.Plan.parse(raw)
      iex> length(plan.tasks)
      1

      iex> raw = %{"steps" => [%{"id" => "s1", "action" => "search"}]}
      iex> {:ok, plan} = PtcRunner.Plan.parse(raw)
      iex> hd(plan.tasks).id
      "s1"

  """
  @spec parse(map()) :: {:ok, t()} | {:error, term()}
  def parse(raw_plan) when is_map(raw_plan) do
    # Unwrap "object" key if present (text mode artifact)
    plan = unwrap_object(raw_plan)

    agents = extract_agents(plan)
    tasks = extract_tasks(plan)

    {:ok,
     %__MODULE__{
       agents: agents,
       tasks: tasks
     }}
  end

  def parse(_), do: {:error, :invalid_plan_format}

  defp unwrap_object(plan), do: plan

  # Extract agents from various possible keys
  defp extract_agents(plan) do
    nested_plan = plan["plan"] || %{}

    agents_raw =
      plan["agents"] ||
        plan["workers"] ||
        nested_plan["agents"] ||
        %{}

    # Normalize agent specs
    agents_raw
    |> Enum.map(fn
      {id, spec} when is_map(spec) ->
        {to_string(id), normalize_agent_spec(spec)}

      {id, prompt} when is_binary(prompt) ->
        # Simple string prompt
        {to_string(id), %{prompt: prompt, tools: []}}
    end)
    |> Map.new()
  end

  defp normalize_agent_spec(spec) when is_map(spec) do
    prompt = spec["prompt"] || spec["description"] || spec["role"] || ""

    tools =
      (spec["tools"] || spec["capabilities"] || [])
      |> Enum.map(&to_string/1)

    %{prompt: prompt, tools: tools}
  end

  # Extract tasks from various possible keys
  defp extract_tasks(plan) do
    nested_plan = plan["plan"] || %{}

    tasks_raw =
      plan["tasks"] ||
        plan["steps"] ||
        plan["workflow"] ||
        plan["phases"] ||
        nested_plan["steps"] ||
        nested_plan["tasks"] ||
        nested_plan["workflow"] ||
        nested_plan["phases"] ||
        []

    # Normalize various LLM output formats to a flat list of task maps
    tasks_raw
    |> normalize_tasks_structure()
    |> Enum.with_index()
    |> Enum.map(fn {task, idx} -> normalize_task(task, idx) end)
  end

  # Handle various LLM output structures for tasks
  defp normalize_tasks_structure(tasks) when is_list(tasks) do
    # Could be a list of tasks or a list of phases with nested tasks
    Enum.flat_map(tasks, fn
      # Phase object with nested tasks array
      %{"tasks" => nested_tasks} = phase when is_list(nested_tasks) ->
        phase_deps = phase["depends_on"] || []
        # Add phase-level dependencies to each task
        Enum.map(nested_tasks, fn task ->
          task_deps = task["depends_on"] || []
          Map.put(task, "depends_on", task_deps ++ phase_deps)
        end)

      # Regular task map
      task when is_map(task) ->
        [task]

      # Simple string task
      task when is_binary(task) ->
        [task]

      # Skip other items (tuples, etc.)
      _ ->
        []
    end)
  end

  defp normalize_tasks_structure(workflow) when is_map(workflow) do
    # Workflow is a map of named phases like %{"parallel_phase_1" => %{...}, ...}
    # Extract tasks from each phase
    workflow
    |> Enum.flat_map(fn
      {_phase_name, %{"tasks" => tasks}} when is_list(tasks) ->
        tasks

      {_phase_name, phase} when is_map(phase) ->
        # Single task disguised as a phase
        [phase]

      _ ->
        []
    end)
  end

  defp normalize_tasks_structure(_), do: []

  defp normalize_task(task, idx) when is_map(task) do
    %{
      id: to_string(extract_task_id(task, idx)),
      agent: to_string(extract_task_agent(task)),
      input: extract_task_input(task),
      depends_on: extract_depends_on(task),
      on_failure: normalize_on_failure(task["on_failure"]),
      max_retries: task["max_retries"] || task["retries"] || 1,
      critical: normalize_critical(task["critical"]),
      type: normalize_task_type(task["type"]),
      output: normalize_output_mode(task["output"]),
      signature: task["signature"] || task["output_signature"],
      verification: task["verification"],
      on_verification_failure: normalize_on_verification_failure(task["on_verification_failure"]),
      quality_gate: normalize_quality_gate(task["quality_gate"])
    }
  end

  defp normalize_task(task, idx) when is_binary(task) do
    # Simple string task
    %{
      id: "task_#{idx + 1}",
      agent: "default",
      input: task,
      depends_on: [],
      on_failure: :stop,
      max_retries: 1,
      critical: true,
      type: :task,
      output: nil,
      signature: nil,
      verification: nil,
      on_verification_failure: :replan,
      quality_gate: nil
    }
  end

  defp extract_task_id(task, idx) do
    task["id"] || task["step_id"] || task["name"] || "task_#{idx + 1}"
  end

  defp extract_task_agent(task) do
    task["agent"] || task["worker"] || task["assigned_to"] || "default"
  end

  defp extract_task_input(task) do
    task["input"] || task["query"] || task["action"] || task["description"] || ""
  end

  defp extract_depends_on(task) do
    (task["depends_on"] || task["requires"] || task["after"] || task["dependencies"] || [])
    |> List.wrap()
    |> Enum.map(&to_string/1)
  end

  defp normalize_critical(false), do: false
  defp normalize_critical("false"), do: false
  defp normalize_critical(_), do: true

  defp normalize_task_type(nil), do: :task
  defp normalize_task_type("synthesis_gate"), do: :synthesis_gate
  defp normalize_task_type("gate"), do: :synthesis_gate
  defp normalize_task_type("synthesis"), do: :synthesis_gate
  defp normalize_task_type("compress"), do: :synthesis_gate
  defp normalize_task_type("summarize"), do: :synthesis_gate
  defp normalize_task_type("human_review"), do: :human_review
  defp normalize_task_type("review"), do: :human_review
  defp normalize_task_type("approval"), do: :human_review
  defp normalize_task_type(_), do: :task

  # Normalize output mode - nil means auto-detect based on tools
  defp normalize_output_mode(nil), do: nil
  defp normalize_output_mode("ptc_lisp"), do: :ptc_lisp
  defp normalize_output_mode("text"), do: :text
  defp normalize_output_mode(:ptc_lisp), do: :ptc_lisp
  defp normalize_output_mode(:text), do: :text
  defp normalize_output_mode(_), do: nil

  defp normalize_on_failure(nil), do: :stop
  defp normalize_on_failure("skip"), do: :skip
  defp normalize_on_failure("retry"), do: :retry
  defp normalize_on_failure("stop"), do: :stop
  defp normalize_on_failure("replan"), do: :replan
  defp normalize_on_failure(_), do: :stop

  defp normalize_on_verification_failure(nil), do: :replan
  defp normalize_on_verification_failure("stop"), do: :stop
  defp normalize_on_verification_failure("skip"), do: :skip
  defp normalize_on_verification_failure("retry"), do: :retry
  defp normalize_on_verification_failure("replan"), do: :replan
  defp normalize_on_verification_failure(_), do: :stop

  defp normalize_quality_gate(true), do: true
  defp normalize_quality_gate(false), do: false
  defp normalize_quality_gate(_), do: nil

  @type validation_issue :: %{
          severity: :error | :warning,
          category: atom(),
          message: String.t(),
          task_id: String.t() | nil
        }

  @doc """
  Validate a parsed plan for structural integrity.

  Checks for:
  - Circular dependencies (cycles in task graph)
  - Missing dependency references (depends_on non-existent task)
  - Missing agent references (task references non-existent agent)
  - Duplicate task IDs

  ## Examples

      iex> plan = %PtcRunner.Plan{tasks: [%{id: "a", depends_on: ["b"], agent: "x", input: ""}]}
      iex> {:error, issues} = PtcRunner.Plan.validate(plan)
      iex> hd(issues).category
      :missing_dependency

      iex> plan = %PtcRunner.Plan{tasks: [
      ...>   %{id: "a", depends_on: ["b"], agent: "x", input: ""},
      ...>   %{id: "b", depends_on: ["a"], agent: "x", input: ""}
      ...> ]}
      iex> {:error, issues} = PtcRunner.Plan.validate(plan)
      iex> hd(issues).category
      :cycle_detected

  """
  @spec validate(t()) :: :ok | {:error, [validation_issue()]}
  def validate(%__MODULE__{} = plan) do
    issues =
      []
      |> check_empty_tasks(plan)
      |> check_duplicate_ids(plan)
      |> check_missing_dependencies(plan)
      |> check_missing_agents(plan)
      |> check_cycles(plan)
      |> check_lisp_syntax(plan)

    case Enum.filter(issues, &(&1.severity == :error)) do
      [] -> :ok
      errors -> {:error, errors}
    end
  end

  @doc """
  Sanitize a plan by removing invalid verification predicates.

  Checks each task's verification predicate and removes any that:
  - Fail to parse
  - Use undefined variables (e.g., `result` instead of `data/result`)

  This is a pragmatic approach: rather than failing on LLM-generated invalid
  predicates, we strip them so the plan can execute without verification.

  Returns `{plan, warnings}` where warnings list any predicates that were removed.
  """
  @spec sanitize(t()) :: {t(), [validation_issue()]}
  def sanitize(%__MODULE__{} = plan) do
    {sanitized_tasks, warnings} =
      Enum.map_reduce(plan.tasks, [], fn task, acc ->
        # Validate verification predicate
        {task, acc} =
          case validate_predicate(task[:verification]) do
            :ok ->
              {task, acc}

            {:invalid, reason} ->
              warning = %{
                severity: :warning,
                category: :invalid_verification,
                message: "Task '#{task.id}' verification removed: #{reason}",
                task_id: task.id
              }

              {Map.put(task, :verification, nil), [warning | acc]}
          end

        # Validate signature
        case validate_signature(task[:signature]) do
          :ok ->
            {task, acc}

          {:invalid, reason} ->
            warning = %{
              severity: :warning,
              category: :invalid_signature,
              message: "Task '#{task.id}' signature removed: #{reason}",
              task_id: task.id
            }

            {Map.put(task, :signature, nil), [warning | acc]}
        end
      end)

    {%{plan | tasks: sanitized_tasks}, Enum.reverse(warnings)}
  end

  defp validate_signature(nil), do: :ok
  defp validate_signature(""), do: :ok

  defp validate_signature(signature) when is_binary(signature) do
    case Signature.parse(signature) do
      {:ok, _} -> :ok
      {:error, reason} -> {:invalid, reason}
    end
  end

  defp validate_signature(_), do: :ok

  defp validate_predicate(nil), do: :ok
  defp validate_predicate(""), do: :ok

  defp validate_predicate(predicate) when is_binary(predicate) do
    case Lisp.validate(predicate) do
      :ok -> :ok
      {:error, messages} -> {:invalid, Enum.join(messages, "; ")}
    end
  end

  defp validate_predicate(_), do: :ok

  # Check that the plan has at least one task
  defp check_empty_tasks(issues, plan) do
    if plan.tasks == [] do
      [
        %{
          severity: :error,
          category: :empty_plan,
          message: "Plan has no tasks. The LLM may have failed to generate tasks.",
          task_id: nil
        }
        | issues
      ]
    else
      issues
    end
  end

  # Check for duplicate task IDs
  defp check_duplicate_ids(issues, plan) do
    ids = Enum.map(plan.tasks, & &1.id)
    duplicates = ids -- Enum.uniq(ids)

    Enum.reduce(Enum.uniq(duplicates), issues, fn id, acc ->
      [
        %{
          severity: :error,
          category: :duplicate_id,
          message: "Duplicate task ID: '#{id}'",
          task_id: id
        }
        | acc
      ]
    end)
  end

  # Check that all depends_on references exist
  defp check_missing_dependencies(issues, plan) do
    task_ids = MapSet.new(plan.tasks, & &1.id)

    Enum.reduce(plan.tasks, issues, fn task, acc ->
      missing =
        task.depends_on
        |> Enum.reject(&MapSet.member?(task_ids, &1))

      Enum.reduce(missing, acc, fn dep_id, inner_acc ->
        [
          %{
            severity: :error,
            category: :missing_dependency,
            message: "Task '#{task.id}' depends on non-existent task '#{dep_id}'",
            task_id: task.id
          }
          | inner_acc
        ]
      end)
    end)
  end

  # Check that all agent references exist (or are "default")
  defp check_missing_agents(issues, plan) do
    agent_ids = MapSet.new(Map.keys(plan.agents))

    Enum.reduce(plan.tasks, issues, fn task, acc ->
      if task.agent in ["default", "direct"] or MapSet.member?(agent_ids, task.agent) do
        acc
      else
        [
          %{
            severity: :error,
            category: :missing_agent,
            message: "Task '#{task.id}' references non-existent agent '#{task.agent}'",
            task_id: task.id
          }
          | acc
        ]
      end
    end)
  end

  # Check for cycles using DFS with "in progress" tracking
  defp check_cycles(issues, plan) do
    task_map = Map.new(plan.tasks, fn t -> {t.id, t} end)
    task_ids = Map.keys(task_map)

    # DFS to find cycles
    {_visited, _in_progress, cycle_issues} =
      Enum.reduce(task_ids, {MapSet.new(), MapSet.new(), []}, fn id,
                                                                 {visited, in_progress, acc} ->
        if MapSet.member?(visited, id) do
          {visited, in_progress, acc}
        else
          detect_cycle(id, task_map, visited, in_progress, acc, [])
        end
      end)

    issues ++ cycle_issues
  end

  # DFS cycle detection with path tracking
  defp detect_cycle(id, task_map, visited, in_progress, issues, path) do
    cond do
      MapSet.member?(in_progress, id) ->
        # Found a cycle! Build the cycle path for the error message
        cycle_start_idx = Enum.find_index(path, &(&1 == id))
        cycle_path = Enum.slice(path, cycle_start_idx..-1//1) ++ [id]
        cycle_str = Enum.join(cycle_path, " -> ")

        issue = %{
          severity: :error,
          category: :cycle_detected,
          message: "Circular dependency detected: #{cycle_str}",
          task_id: id
        }

        {visited, in_progress, [issue | issues]}

      MapSet.member?(visited, id) ->
        # Already fully processed
        {visited, in_progress, issues}

      true ->
        # Mark as in progress
        in_progress = MapSet.put(in_progress, id)
        task = Map.get(task_map, id)
        new_path = path ++ [id]

        # Visit all dependencies
        {visited, in_progress, issues} =
          if task do
            Enum.reduce(task.depends_on, {visited, in_progress, issues}, fn dep_id,
                                                                            {v, ip, iss} ->
              detect_cycle(dep_id, task_map, v, ip, iss, new_path)
            end)
          else
            {visited, in_progress, issues}
          end

        # Mark as fully visited (remove from in_progress, add to visited)
        in_progress = MapSet.delete(in_progress, id)
        visited = MapSet.put(visited, id)

        {visited, in_progress, issues}
    end
  end

  # Check PTC-Lisp syntax in direct task inputs and verification predicates
  defp check_lisp_syntax(issues, plan) do
    Enum.reduce(plan.tasks, issues, fn task, acc ->
      acc
      |> check_direct_input(task)
      |> check_verification_predicate(task)
    end)
  end

  defp check_direct_input(issues, %{agent: "direct", input: input, id: id})
       when is_binary(input) and input != "" do
    case Lisp.validate(input) do
      :ok -> issues
      {:error, messages} -> [lisp_issue(id, :invalid_direct_input, messages) | issues]
    end
  end

  defp check_direct_input(issues, _task), do: issues

  defp check_verification_predicate(issues, %{verification: pred, id: id})
       when is_binary(pred) and pred != "" do
    case Lisp.validate(pred) do
      :ok -> issues
      {:error, messages} -> [lisp_issue(id, :invalid_verification, messages) | issues]
    end
  end

  defp check_verification_predicate(issues, _task), do: issues

  defp lisp_issue(task_id, category, messages) do
    label =
      case category do
        :invalid_direct_input -> "direct task input"
        :invalid_verification -> "verification predicate"
      end

    %{
      severity: :error,
      category: category,
      message: "Task '#{task_id}' has invalid #{label}: #{Enum.join(messages, "; ")}",
      task_id: task_id
    }
  end

  @doc """
  Sort tasks by dependencies (topological sort).

  Returns tasks in an order where all dependencies come before dependents.
  Raises if there's a cycle.

  ## Examples

      iex> tasks = [
      ...>   %{id: "t2", depends_on: ["t1"], agent: "a", input: ""},
      ...>   %{id: "t1", depends_on: [], agent: "a", input: ""}
      ...> ]
      iex> sorted = PtcRunner.Plan.topological_sort(tasks)
      iex> Enum.map(sorted, & &1.id)
      ["t1", "t2"]

  """
  @spec topological_sort([task()]) :: [task()]
  def topological_sort(tasks) do
    task_map = Map.new(tasks, fn t -> {t.id, t} end)
    visited = MapSet.new()
    result = []

    {sorted, _visited} =
      Enum.reduce(tasks, {result, visited}, fn task, {acc, vis} ->
        visit(task.id, task_map, vis, acc)
      end)

    Enum.reverse(sorted)
  end

  defp visit(id, task_map, visited, result) do
    if MapSet.member?(visited, id) do
      {result, visited}
    else
      visited = MapSet.put(visited, id)
      task = Map.get(task_map, id)

      {result, visited} =
        if task do
          Enum.reduce(task.depends_on, {result, visited}, fn dep_id, {acc, vis} ->
            visit(dep_id, task_map, vis, acc)
          end)
        else
          {result, visited}
        end

      if task do
        {[task | result], visited}
      else
        {result, visited}
      end
    end
  end

  @doc """
  Group tasks into parallel execution phases by dependency level.

  Tasks with no dependencies are level 0 (can run first in parallel).
  Tasks whose dependencies are all level 0 are level 1, etc.

  Returns a list of phases, where each phase is a list of tasks
  that can execute in parallel.

  ## Examples

      iex> tasks = [
      ...>   %{id: "a", depends_on: [], agent: "x", input: "", on_failure: :stop, max_retries: 1, critical: true},
      ...>   %{id: "b", depends_on: [], agent: "x", input: "", on_failure: :stop, max_retries: 1, critical: true},
      ...>   %{id: "c", depends_on: ["a", "b"], agent: "x", input: "", on_failure: :stop, max_retries: 1, critical: true}
      ...> ]
      iex> phases = PtcRunner.Plan.group_by_level(tasks)
      iex> length(phases)
      2
      iex> phases |> hd() |> Enum.map(& &1.id) |> Enum.sort()
      ["a", "b"]
      iex> phases |> Enum.at(1) |> Enum.map(& &1.id)
      ["c"]

  """
  @spec group_by_level([task()]) :: [[task()]]
  def group_by_level([]), do: []

  def group_by_level(tasks) do
    task_map = Map.new(tasks, fn t -> {t.id, t} end)

    # Compute level for each task
    levels = compute_levels(tasks, task_map)

    # Group tasks by level
    max_level = levels |> Map.values() |> Enum.max(fn -> 0 end)

    for level <- 0..max_level do
      tasks
      |> Enum.filter(fn task -> Map.get(levels, task.id) == level end)
    end
    |> Enum.reject(&Enum.empty?/1)
  end

  # Compute the level (depth) of each task based on dependencies
  defp compute_levels(tasks, task_map) do
    Enum.reduce(tasks, %{}, fn task, levels ->
      {updated_levels, _level} = compute_level(task.id, task_map, levels)
      updated_levels
    end)
  end

  defp compute_level(id, task_map, levels) do
    case Map.fetch(levels, id) do
      {:ok, _level} ->
        # Already computed
        {levels, Map.get(levels, id)}

      :error ->
        task = Map.get(task_map, id)

        if task == nil or task.depends_on == [] do
          # No dependencies = level 0
          {Map.put(levels, id, 0), 0}
        else
          # Level = max(dependency levels) + 1
          {levels, max_dep_level} =
            Enum.reduce(task.depends_on, {levels, -1}, fn dep_id, {lvls, max_lvl} ->
              {lvls, dep_level} = compute_level(dep_id, task_map, lvls)
              {lvls, max(max_lvl, dep_level)}
            end)

          level = max_dep_level + 1
          {Map.put(levels, id, level), level}
        end
    end
  end
end