gen_server_virtual_time

defmodule ActorSimulation do
  @moduledoc """
  A DSL for simulating actor systems with message rates and statistics.

  This module provides a way to define actors, their message sending patterns,
  and simulate their interactions using virtual time.

  ## Example

      simulation =
        ActorSimulation.new()
        |> ActorSimulation.add_actor(:producer,
            send_pattern: {:periodic, 100, {:data, :id}},
            targets: [:consumer])
        |> ActorSimulation.add_actor(:consumer,
            on_receive: fn msg, state ->
              # Process message and maybe send response
              {:ok, state}
            end)
        |> ActorSimulation.run(duration: 5000)

      stats = ActorSimulation.get_stats(simulation)
      IO.inspect(stats)
  """

  alias ActorSimulation.{Actor, Definition, MermaidReportGenerator, Stats}

  defstruct [
    :clock,
    :actors,
    :stats,
    :running,
    :trace,
    :trace_enabled,
    :actual_duration,
    :terminated_early,
    :termination_reason
  ]

  @doc """
  Creates a new actor simulation.

  Options:
  - `:trace` - Enable message tracing for sequence diagrams (default: false)

  ## Example

      iex> simulation = ActorSimulation.new()
      iex> is_pid(simulation.clock)
      true
      iex> simulation.actors
      %{}

  """
  def new(opts \\ []) do
    {:ok, clock} = VirtualClock.start_link()
    trace_enabled = Keyword.get(opts, :trace, false)

    %__MODULE__{
      clock: clock,
      actors: %{},
      stats: Stats.new(),
      running: false,
      trace: [],
      trace_enabled: trace_enabled,
      actual_duration: 0
    }
  end

  @doc """
  Helper to collect current stats during simulation (for termination conditions).
  Can be called from terminate_when functions.

  Note: This is a live snapshot and may be called multiple times during run/2.
  """
  def collect_current_stats(simulation) do
    Enum.reduce(simulation.actors, Stats.new(), fn {name, actor_info}, stats ->
      case actor_info.type do
        :simulated ->
          actor_stats = Actor.get_stats(actor_info.pid)
          Stats.add_actor_stats(stats, name, actor_stats)

        :real_process ->
          # Get stats from VirtualTimeGenServer's built-in tracking
          real_stats = GenServer.call(actor_info.pid, :__vtgs_get_stats__)
          Stats.add_actor_stats(stats, name, real_stats)
      end
    end)
  end

  @doc """
  Enables message tracing for the simulation.
  """
  def enable_trace(simulation) do
    %{simulation | trace_enabled: true}
  end

  @doc """
  Adds an actor to the simulation.

  Options:
  - `:send_pattern` - How this actor sends messages:
    - `{:periodic, interval, message}` - Send message every interval ms
    - `{:rate, messages_per_second, message}` - Send at a specific rate
    - `{:burst, count, interval, message}` - Send count messages every interval
  - `:targets` - List of actor names to send messages to
  - `:on_receive` - Function called when receiving a message: `fn msg, state -> {:ok, new_state} | {:send, msgs, new_state} end`
  - `:on_match` - Pattern matching responses: `[{pattern, response_fn}]`
  - `:initial_state` - Initial state for the actor (default: %{})
  """
  def add_actor(simulation, name, opts \\ []) do
    actor_def = Definition.new(name, opts)
    {:ok, pid} = Actor.start_link(actor_def, simulation.clock)

    actors =
      Map.put(simulation.actors, name, %{pid: pid, definition: actor_def, type: :simulated})

    %{simulation | actors: actors}
  end

  @doc """
  Adds a real GenServerVirtualTime process to the simulation ("Process in the Loop").

  This allows you to test real GenServer implementations alongside simulated actors.

  Options:
  - `:module` - The GenServer module to start (required)
  - `:args` - Arguments to pass to the module's init/1
  - `:targets` - List of actor names this process can send to (optional)

  ## Example

      defmodule MyRealServer do
        use VirtualTimeGenServer

        def init(args), do: {:ok, args}
        def handle_call(:ping, _from, state), do: {:reply, :pong, state}
      end

      simulation =
        ActorSimulation.new()
        |> ActorSimulation.add_process(:my_server, module: MyRealServer, args: %{})
        |> ActorSimulation.add_actor(:pinger,
            send_pattern: {:periodic, 100, {:call, :my_server, :ping}})
  """
  def add_process(simulation, name, opts \\ []) do
    module = Keyword.fetch!(opts, :module)
    args = Keyword.get(opts, :args, nil)
    targets = Keyword.get(opts, :targets, [])

    # Start the real GenServer with virtual time
    VirtualTimeGenServer.set_virtual_clock(simulation.clock)

    # Enable stats tracking before starting the process
    VirtualTimeGenServer.enable_stats_tracking()

    {:ok, pid} = VirtualTimeGenServer.start_link(module, args, name: name)

    # Set up actor name context for trace generation
    Process.put(:__vtgs_actor_name__, name)

    actors =
      Map.put(simulation.actors, name, %{
        pid: pid,
        type: :real_process,
        module: module,
        targets: targets
      })

    %{simulation | actors: actors}
  end

  @doc """
  Runs the simulation for the specified duration (in milliseconds).

  Options:
  - `:duration` - Maximum duration in milliseconds (default: 10,000)
  - `:terminate_when` - Function that takes simulation and returns true to stop
  - `:check_interval` - How often to check termination condition in ms (default: 100)

  ## Example

      # Run for fixed duration (backward compatible)
      simulation = ActorSimulation.run(simulation, duration: 5000)

      # Run until condition met (new feature)
      simulation = ActorSimulation.run(simulation,
        max_duration: 10_000,
        terminate_when: fn sim ->
          # Stop when all actors have sent at least 10 messages
          stats = ActorSimulation.collect_current_stats(sim)
          Enum.all?(stats.actors, fn {_name, s} -> s.sent_count >= 10 end)
        end
      )
  """
  def run(simulation, opts \\ []) do
    duration = Keyword.get(opts, :duration) || Keyword.get(opts, :max_duration, 10_000)
    terminate_when = Keyword.get(opts, :terminate_when)
    check_interval = Keyword.get(opts, :check_interval, 100)

    # Register trace collector if tracing enabled
    if simulation.trace_enabled do
      # Unregister if already exists
      case Process.whereis(:trace_collector) do
        nil ->
          Process.register(self(), :trace_collector)

        # Already registered to us
        pid when pid == self() ->
          :ok

        _other ->
          Process.unregister(:trace_collector)
          Process.register(self(), :trace_collector)
      end
    end

    # Start all actors (only simulated actors need setup)
    Enum.each(simulation.actors, fn {_name, actor_info} ->
      case actor_info.type do
        :simulated ->
          Actor.start_sending(actor_info.pid, simulation.actors, simulation.trace_enabled)

        :real_process ->
          # Real processes are already started and don't need actor map
          :ok
      end
    end)

    # Measure real time for simulation execution
    start_time = System.monotonic_time(:millisecond)

    # Advance virtual time with optional termination check
    {actual_duration, accumulated_trace, termination_reason} =
      cond do
        terminate_when == :quiescence ->
          {dur, tr} = advance_until_quiescence(simulation, duration)
          {dur, tr, :quiescence}

        is_function(terminate_when) ->
          {dur, tr} = advance_with_condition(simulation, duration, terminate_when, check_interval)
          {dur, tr, :condition}

        true ->
          VirtualClock.advance(simulation.clock, duration)
          {duration, [], :max_time}
      end

    end_time = System.monotonic_time(:millisecond)
    real_elapsed = end_time - start_time

    # Mark if terminated early due to condition
    terminated_early =
      case termination_reason do
        :condition -> actual_duration < duration
        _ -> false
      end

    # Update simulation with timing info before collecting stats
    # (stats need actual_duration for rate calculations)
    simulation_with_timing =
      simulation
      |> Map.put(:actual_duration, actual_duration)
      |> Map.put(:max_duration, duration)
      |> Map.put(:terminated_early, terminated_early)
      |> Map.put(:real_time_elapsed, real_elapsed)
      |> Map.put(:termination_reason, termination_reason)

    # Collect statistics and trace (now that actual_duration is set)
    stats = collect_stats(simulation_with_timing)
    # If we accumulated trace during termination checks, use it and collect any remaining messages
    # Otherwise collect fresh trace
    trace =
      if simulation.trace_enabled do
        remaining_trace = collect_trace()
        accumulated_trace ++ remaining_trace
      else
        []
      end

    # Return enhanced simulation state with timing information
    %{simulation_with_timing | stats: stats, trace: trace, running: false}
  end

  @doc """
  Gets statistics from the simulation.
  """
  def get_stats(simulation) do
    simulation.stats
  end

  @doc """
  Gets the message trace from the simulation.
  Returns a list of trace events for building sequence diagrams.

  Each event is a map with:
  - `:timestamp` - Virtual time when message was sent
  - `:from` - Sender actor name
  - `:to` - Receiver actor name
  - `:message` - The message sent
  - `:type` - `:cast`, `:call`, or `:send`
  """
  def get_trace(simulation) do
    simulation.trace
  end

  @doc """
  Formats the trace as a Mermaid sequence diagram with enhanced styling.

  Mermaid is widely supported in GitHub, GitLab, and many markdown viewers.
  Uses features from [Mermaid sequence diagrams](https://docs.mermaidchart.com/mermaid-oss/syntax/sequenceDiagram.html):
  - Different arrow types for call/cast/send
  - Activation boxes for processing
  - Notes with timestamps
  - Background highlighting for grouped interactions
  - Termination indicators showing when simulation stopped

  ## Example

      iex> simulation = %ActorSimulation{trace: [
      ...>   %{from: :alice, to: :bob, message: :hello, type: :send, timestamp: 100},
      ...>   %{from: :bob, to: :alice, message: :hi, type: :send, timestamp: 200}
      ...> ]}
      iex> mermaid = ActorSimulation.trace_to_mermaid(simulation)
      iex> String.contains?(mermaid, "sequenceDiagram")
      true
      iex> String.contains?(mermaid, "alice->>bob")
      true

  """
  def trace_to_mermaid(simulation, opts \\ []) do
    enhanced = Keyword.get(opts, :enhanced, true)
    show_timestamps = Keyword.get(opts, :timestamps, false)
    show_termination = Keyword.get(opts, :show_termination, true)

    lines = ["sequenceDiagram"]

    # Group events by time ranges for background highlighting
    message_lines =
      if enhanced do
        generate_enhanced_mermaid(simulation.trace, show_timestamps)
      else
        generate_simple_mermaid(simulation.trace)
      end

    # Add termination indicator if simulation ended early
    final_lines =
      if show_termination && Map.get(simulation, :terminated_early, false) do
        # Get all unique actors from trace
        actors =
          simulation.trace
          |> Enum.flat_map(fn event -> [event.from, event.to] end)
          |> Enum.uniq()

        termination_note =
          case {actors, Map.get(simulation, :actual_duration)} do
            {[], _} ->
              []

            {[single], duration} ->
              ["    Note over #{single}: ⚡ Terminated at t=#{duration}ms (goal achieved)"]

            {[first, second | _], duration} ->
              [
                "    Note over #{first},#{second}: ⚡ Terminated at t=#{duration}ms (goal achieved)"
              ]
          end

        message_lines ++ termination_note
      else
        message_lines
      end

    (lines ++ final_lines)
    |> Enum.join("\n")
  end

  defp generate_enhanced_mermaid(trace, show_timestamps) do
    trace
    |> Enum.map(fn event ->
      # Different arrow styles based on message type
      # See: https://docs.mermaidchart.com/mermaid-oss/syntax/sequenceDiagram.html#messages
      arrow =
        case event.type do
          # Solid line with arrowhead (synchronous)
          :call -> "->>"
          # Dotted line with arrowhead (asynchronous)
          :cast -> "-->>"
          # Solid line with arrowhead
          :send -> "->>"
        end

      msg = inspect(event.message)

      # Add activation for calls (shows processing)
      lines =
        if event.type == :call do
          [
            "    activate #{event.to}",
            "    #{event.from}#{arrow}#{event.to}: #{msg}",
            "    deactivate #{event.to}"
          ]
        else
          ["    #{event.from}#{arrow}#{event.to}: #{msg}"]
        end

      # Add timestamp note if requested
      if show_timestamps && event.timestamp do
        lines ++ ["    Note over #{event.from},#{event.to}: t=#{event.timestamp}ms"]
      else
        lines
      end
    end)
    |> List.flatten()
  end

  defp generate_simple_mermaid(trace) do
    Enum.map(trace, fn event ->
      arrow =
        case event.type do
          :call -> "->>"
          :cast -> "->>"
          :send -> "->>"
        end

      msg = inspect(event.message)
      "    #{event.from}#{arrow}#{event.to}: #{msg}"
    end)
  end

  @doc """
  Generates a Mermaid flowchart report showing actor topology and statistics.

  This creates a visual report with:
  - Actor topology as a flowchart
  - Statistics embedded in nodes (message counts, rates)
  - Activity-based styling (color coding by traffic)
  - Complete HTML page with stats table

  ## Options

  - `:title` - Report title (default: "Simulation Report")
  - `:show_stats_on_nodes` - Show stats on nodes (default: true)
  - `:show_message_labels` - Show message types on edges (default: true)
  - `:layout` - Direction: "TB", "LR", "RL", "BT" (default: "TB")
  - `:style_by_activity` - Color by activity level (default: true)

  ## Example

      simulation = ActorSimulation.new()
        |> add_actor(:producer, send_pattern: {:periodic, 100, :data}, targets: [:consumer])
        |> add_actor(:consumer)
        |> run(duration: 1000)

      html = ActorSimulation.generate_flowchart_report(simulation,
        title: "My System",
        layout: "LR")

      File.write!("report.html", html)

  Based on [Mermaid Flowchart Syntax](https://mermaid.js.org/syntax/flowchart.html)
  """
  def generate_flowchart_report(simulation, opts \\ []) do
    MermaidReportGenerator.generate_report(simulation, opts)
  end

  @doc """
  Writes a flowchart report directly to a file.

  Same options as `generate_flowchart_report/2`.

  ## Example

      simulation = ActorSimulation.new()
        |> add_actor(:source, send_pattern: {:periodic, 100, :msg}, targets: [:sink])
        |> add_actor(:sink)
        |> run(duration: 500)

      {:ok, path} = ActorSimulation.write_flowchart_report(
        simulation,
        "my_report.html",
        title: "Production System")
  """
  def write_flowchart_report(simulation, filename, opts \\ []) do
    MermaidReportGenerator.write_report(simulation, filename, opts)
  end

  @doc """
  Stops the simulation and cleans up resources.
  """
  def stop(simulation) do
    Enum.each(simulation.actors, fn {_name, %{pid: pid}} ->
      GenServer.stop(pid, :normal, :infinity)
    end)

    GenServer.stop(simulation.clock)
    :ok
  end

  # Private functions

  defp advance_with_condition(simulation, max_duration, condition_fn, check_interval) do
    # Start with empty accumulated trace
    advance_with_condition_loop(simulation, max_duration, condition_fn, check_interval, 0, [])
  end

  defp advance_with_condition_loop(
         simulation,
         max_duration,
         condition_fn,
         check_interval,
         elapsed,
         accumulated_trace
       ) do
    if elapsed >= max_duration do
      {max_duration, accumulated_trace}
    else
      # Advance by check_interval
      step = min(check_interval, max_duration - elapsed)
      VirtualClock.advance(simulation.clock, step)
      new_elapsed = elapsed + step

      # Collect new trace messages and accumulate them
      # This allows terminate_when to access the trace accumulated so far
      {current_simulation, new_accumulated_trace} =
        if simulation.trace_enabled do
          new_trace_messages = collect_trace()
          full_trace = accumulated_trace ++ new_trace_messages
          {%{simulation | trace: full_trace}, full_trace}
        else
          {simulation, accumulated_trace}
        end

      # Check termination condition
      if condition_fn.(current_simulation) do
        {new_elapsed, new_accumulated_trace}
      else
        advance_with_condition_loop(
          simulation,
          max_duration,
          condition_fn,
          check_interval,
          new_elapsed,
          new_accumulated_trace
        )
      end
    end
  end

  # Advance until quiescence or max_duration
  # Quiescence: no scheduled timers and mailbox drains to empty across all simulated actors
  defp advance_until_quiescence(simulation, max_duration) do
    advance_until_quiescence_loop(simulation, max_duration, 0, [])
  end

  defp advance_until_quiescence_loop(simulation, max_duration, elapsed, accumulated_trace) do
    if elapsed >= max_duration do
      {max_duration, accumulated_trace}
    else
      # If no timers scheduled, we are quiescent
      if VirtualClock.scheduled_count(simulation.clock) == 0 do
        {elapsed, accumulated_trace}
      else
        # Advance to next scheduled event to let messages flow
        advance = VirtualClock.advance_to_next(simulation.clock)
        new_elapsed = elapsed + advance

        {acc_sim, acc_trace} =
          if simulation.trace_enabled do
            new_msgs = collect_trace()
            {%{simulation | trace: accumulated_trace ++ new_msgs}, accumulated_trace ++ new_msgs}
          else
            {simulation, accumulated_trace}
          end

        advance_until_quiescence_loop(acc_sim, max_duration, new_elapsed, acc_trace)
      end
    end
  end

  defp collect_stats(simulation) do
    # Get the actual duration that was simulated
    actual_duration = Map.get(simulation, :actual_duration, 0)

    stats =
      Enum.reduce(simulation.actors, simulation.stats, fn {name, actor_info}, stats ->
        case actor_info.type do
          :simulated ->
            actor_stats = Actor.get_stats(actor_info.pid)
            Stats.add_actor_stats(stats, name, actor_stats)

          :real_process ->
            # Get stats from VirtualTimeGenServer's built-in tracking
            real_stats = GenServer.call(actor_info.pid, :__vtgs_get_stats__)
            Stats.add_actor_stats(stats, name, real_stats)
        end
      end)

    # Set the time range for rate calculations
    %{stats | start_time: 0, end_time: actual_duration}
  end

  defp collect_trace do
    # Collect all trace messages
    collect_trace_messages([])
  end

  defp collect_trace_messages(acc) do
    receive do
      {:trace, event} -> collect_trace_messages([event | acc])
    after
      0 -> Enum.reverse(acc)
    end
  end
end