tla_connect

defmodule TlaConnect.Validator do
  @moduledoc """
  Approach 3, part 2: Trace validation via Apalache.

  Validates NDJSON traces against a TLA+ spec by generating a TraceData
  module and running Apalache in check mode.
  """

  alias TlaConnect.ValidationError

  @type config :: %{
          trace_spec: Path.t(),
          init: String.t(),
          next: String.t(),
          inv: String.t(),
          cinit: String.t() | nil,
          apalache_bin: String.t(),
          timeout: pos_integer() | nil
        }

  @default_config %{
    cinit: nil,
    apalache_bin: "apalache-mc",
    timeout: nil
  }

  @doc "Validate an NDJSON trace file against a TLA+ spec."
  @spec validate_trace(map(), Path.t()) :: {:ok, %{valid: boolean(), reason: String.t() | nil}}
  def validate_trace(config, trace_path) do
    config = Map.merge(@default_config, config)

    case File.read(trace_path) do
      {:ok, content} ->
        content = String.trim(content)

        if content == "" do
          {:ok, %{valid: false, reason: "Trace file is empty"}}
        else
          do_validate(config, content)
        end

      {:error, reason} ->
        {:ok, %{valid: false, reason: "Cannot read trace file: #{inspect(reason)}"}}
    end
  end

  @doc "Convert a list of maps (parsed NDJSON objects) to a TLA+ TraceData module string."
  @spec ndjson_to_tla_module([map()]) :: String.t()
  def ndjson_to_tla_module([]) do
    raise ValidationError, message: "Cannot generate TLA+ module from empty trace"
  end

  def ndjson_to_tla_module(objects) do
    # Validate schema consistency
    ref_keys = objects |> hd() |> Map.keys() |> Enum.sort()

    objects
    |> Enum.with_index(1)
    |> Enum.each(fn {obj, i} ->
      keys = obj |> Map.keys() |> Enum.sort()

      if keys != ref_keys do
        raise ValidationError,
          message: "Schema inconsistency: line 1 has keys [#{Enum.join(ref_keys, ", ")}] but line #{i} has keys [#{Enum.join(keys, ", ")}]"
      end
    end)

    # Validate no floats
    objects
    |> Enum.with_index(1)
    |> Enum.each(fn {obj, line} ->
      Enum.each(obj, fn {key, val} ->
        validate_no_floats(val, line, key)
      end)
    end)

    # Build TLA+ module
    lines = [
      "---- MODULE TraceData ----",
      "EXTENDS Integers, Sequences",
      "",
      "TraceLen == #{length(objects)}",
      ""
    ]

    # Generate typed trace data
    trace_lines =
      Enum.map(ref_keys, fn key ->
        values = Enum.map(objects, fn obj -> json_to_tla(Map.get(obj, key)) end)
        "trace_#{key} == <<#{Enum.join(values, ", ")}>>"
      end)

    # Type annotations (Snowcat)
    type_lines =
      ["", "\\* Type annotations for Snowcat"] ++
        Enum.flat_map(ref_keys, fn key ->
          typ = infer_tla_type(Map.get(hd(objects), key))

          [
            "\\* @type: Seq(#{typ});",
            "ASSUME trace_#{key} \\in Seq(#{typ})"
          ]
        end)

    all_lines = lines ++ trace_lines ++ type_lines ++ ["", "===="]
    Enum.join(all_lines, "\n")
  end

  # -- Private --

  defp do_validate(config, content) do
    lines = String.split(content, "\n")

    case parse_ndjson_lines(lines) do
      {:ok, objects} ->
        tla_module = ndjson_to_tla_module(objects)

        dir = Path.join(System.tmp_dir!(), "tla-validate-#{:erlang.unique_integer([:positive])}")
        File.mkdir_p!(dir)

        # Copy trace spec and sibling .tla files to work dir so Apalache
        # can resolve EXTENDS references
        spec_basename = Path.basename(config.trace_spec)
        spec_dir = Path.dirname(config.trace_spec)
        copy_tla_files_to_dir(spec_dir, dir)
        spec_in_dir = Path.join(dir, spec_basename)

        # Write the generated TraceData module alongside the spec
        module_path = Path.join(dir, "TraceData.tla")
        File.write!(module_path, tla_module)

        try do
          run_apalache_validation(config, objects, dir, spec_in_dir)
        after
          File.rm_rf(dir)
        end

      {:error, reason} ->
        {:ok, %{valid: false, reason: reason}}
    end
  end

  defp copy_tla_files_to_dir(src_dir, dest_dir) do
    case File.ls(src_dir) do
      {:ok, entries} ->
        Enum.each(entries, fn entry ->
          if String.ends_with?(entry, ".tla") do
            src = Path.join(src_dir, entry)
            dest = Path.join(dest_dir, entry)
            File.cp(src, dest)
          end
        end)

      {:error, _} ->
        :ok
    end
  end

  defp parse_ndjson_lines(lines) do
    Enum.reduce_while(lines, {:ok, []}, fn line, {:ok, acc} ->
      case Jason.decode(line) do
        {:ok, obj} when is_map(obj) ->
          {:cont, {:ok, acc ++ [obj]}}

        {:ok, _other} ->
          {:halt, {:error, "Line #{length(acc) + 1} is not a JSON object"}}

        {:error, _} ->
          {:halt, {:error, "Line #{length(acc) + 1} is not valid JSON"}}
      end
    end)
  end

  defp run_apalache_validation(config, objects, _dir, spec_in_dir) do
    args = [
      "check",
      "--init=#{config.init}",
      "--next=#{config.next}",
      "--inv=#{config.inv}",
      "--length=#{length(objects)}"
    ]

    args = if config.cinit, do: args ++ ["--cinit=#{config.cinit}"], else: args
    args = args ++ [spec_in_dir]

    bin = config.apalache_bin
    port_opts = [:binary, :exit_status, :stderr_to_stdout, args: args]

    try do
      port = Port.open({:spawn_executable, System.find_executable(bin) || bin}, port_opts)
      result = collect_port_output(port, "", config.timeout)

      case result do
        {:ok, 12} ->
          {:ok, %{valid: true, reason: nil}}

        {:ok, 0} ->
          {:ok, %{valid: false, reason: "Apalache found no counterexample: trace does not satisfy spec"}}

        {:ok, code} ->
          {:ok, %{valid: false, reason: "Apalache exited with code #{code}"}}

        {:error, reason} ->
          {:ok, %{valid: false, reason: reason}}
      end
    rescue
      e ->
        {:ok, %{valid: false, reason: Exception.message(e)}}
    end
  end

  defp collect_port_output(port, acc, timeout) do
    deadline = if timeout, do: System.monotonic_time(:millisecond) + timeout, else: nil
    port_receive_loop(port, acc, deadline)
  end

  defp port_receive_loop(port, acc, deadline) do
    remaining =
      if deadline, do: max(deadline - System.monotonic_time(:millisecond), 0), else: :infinity

    receive do
      {^port, {:data, data}} ->
        port_receive_loop(port, acc <> data, deadline)

      {^port, {:exit_status, code}} ->
        {:ok, code}
    after
      remaining ->
        Port.close(port)
        {:error, "Apalache timed out"}
    end
  end

  @doc false
  def json_to_tla(nil), do: "\"null\""
  def json_to_tla(true), do: "TRUE"
  def json_to_tla(false), do: "FALSE"

  def json_to_tla(v) when is_integer(v), do: Integer.to_string(v)

  def json_to_tla(v) when is_float(v) do
    raise ValidationError, message: "Float values not supported in TLA+: #{v}"
  end

  def json_to_tla(v) when is_binary(v), do: "\"#{escape_tla_string(v)}\""

  def json_to_tla(v) when is_list(v) do
    "<<#{Enum.map(v, &json_to_tla/1) |> Enum.join(", ")}>>"
  end

  def json_to_tla(v) when is_map(v) do
    entries = Map.to_list(v)

    if entries == [] do
      "[dummy |-> 0]"
    else
      inner = Enum.map(entries, fn {k, val} -> "#{k} |-> #{json_to_tla(val)}" end) |> Enum.join(", ")
      "[#{inner}]"
    end
  end

  def json_to_tla(v), do: "\"#{v}\""

  defp escape_tla_string(s) do
    s
    |> String.replace("\\", "\\\\")
    |> String.replace("\"", "\\\"")
    |> String.replace("\n", "\\n")
    |> String.replace("\t", "\\t")
    |> String.replace("\r", "\\r")
  end

  @doc false
  def infer_tla_type(nil), do: "Str"
  def infer_tla_type(v) when is_boolean(v), do: "Bool"
  def infer_tla_type(v) when is_integer(v), do: "Int"
  def infer_tla_type(v) when is_binary(v), do: "Str"

  def infer_tla_type(v) when is_list(v) do
    if v == [] do
      "Seq(Int)"
    else
      "Seq(#{infer_tla_type(hd(v))})"
    end
  end

  def infer_tla_type(v) when is_map(v) do
    entries = Map.to_list(v)

    if entries == [] do
      "[dummy: Int]"
    else
      inner = Enum.map(entries, fn {k, val} -> "#{k}: #{infer_tla_type(val)}" end) |> Enum.join(", ")
      "[#{inner}]"
    end
  end

  def infer_tla_type(_), do: "Str"

  defp validate_no_floats(v, line, path) when is_float(v) do
    raise ValidationError,
      message: "Float value at line #{line}, field \"#{path}\": #{v}. TLA+ uses Int."
  end

  defp validate_no_floats(v, line, path) when is_list(v) do
    v
    |> Enum.with_index()
    |> Enum.each(fn {item, i} -> validate_no_floats(item, line, "#{path}[#{i}]") end)
  end

  defp validate_no_floats(v, line, path) when is_map(v) do
    Enum.each(v, fn {k, val} -> validate_no_floats(val, line, "#{path}.#{k}") end)
  end

  defp validate_no_floats(_v, _line, _path), do: :ok
end