defmodule Bond.Compiler.AnnotatedFunction do
  @moduledoc internal: true
  @moduledoc """
  Internal model of a user function plus everything attached to it at compile time: its clauses,
  its preconditions, its postconditions, and any `@doc` attributes that precede it.

  All clauses of a function (with the same name and arity) are gathered into one
  `AnnotatedFunction` struct. The `:clauses` field holds an ordered list of
  `Bond.Compiler.AnnotatedFunction.Clause` structs, one per `def`/`defp` clause. The
  `:preconditions`, `:postconditions`, and `:doc_attributes` fields apply to all clauses.

  An `AnnotatedFunction` is produced by `Bond.Compiler.CompileStateFSM` from the
  `Bond.Compiler.FunctionDefinition` events emitted by `@on_definition`, and consumed by
  `apply_contract/1` in this module to generate the override that wraps the original function in
  pre/post evaluation.
  """

  alias Bond.Compiler.Assertion
  alias Bond.Compiler.ContractDocs
  alias Bond.Compiler.FunctionDefinition
  alias Bond.Compiler.Invariants
  alias Bond.Compiler.OldExpression

  defstruct kind: nil,
            module: nil,
            fun: nil,
            arity: nil,
            clauses: [],
            preconditions: [],
            postconditions: [],
            invariants: [],
            doc_attributes: []

  @type t :: %__MODULE__{
          kind: :def | :defp | nil,
          module: module() | nil,
          fun: atom() | nil,
          arity: non_neg_integer() | nil,
          clauses: [__MODULE__.Clause.t()],
          preconditions: [Bond.Compiler.Assertion.t()],
          postconditions: [Bond.Compiler.Assertion.t()],
          invariants: [Bond.Compiler.Assertion.t()],
          doc_attributes: [FunctionDefinition.doc_attribute()]
        }

  defmodule Clause do
    @moduledoc internal: true
    @moduledoc """
    Struct to represent an individual clause of a function.
    """

    alias Bond.Compiler.FunctionDefinition

    defstruct [:env, :params, :guards, :body]

    @type t :: %__MODULE__{
            env: Macro.Env.t() | nil,
            params: list() | nil,
            guards: list() | nil,
            body: keyword() | nil
          }

    def new(%FunctionDefinition{} = function_def) do
      struct(__MODULE__, Map.take(function_def, [:env, :params, :guards, :body]))
    end
  end

  def new(%FunctionDefinition{} = function_def) do
    %__MODULE__{
      kind: function_def.kind,
      module: function_def.module,
      fun: function_def.fun,
      arity: function_def.arity,
      clauses: [Clause.new(function_def)]
    }
  end

  def mfa(%__MODULE__{module: module, fun: function, arity: arity}), do: {module, function, arity}

  def add_clause(
        %__MODULE__{module: module, fun: function, arity: arity, clauses: clauses} = function_def,
        %FunctionDefinition{module: module, fun: function, arity: arity} = clause_def
      ) do
    %{function_def | clauses: clauses ++ [Clause.new(clause_def)]}
  end

  def put_preconditions(
        %__MODULE__{preconditions: existing_preconditions} = annotated_function,
        preconditions
      )
      when is_list(preconditions) do
    # Check to make sure each element in the list is a `Bond.Assertion` struct.
    Enum.each(preconditions, fn %Assertion{} -> :ok end)

    %{annotated_function | preconditions: existing_preconditions ++ preconditions}
  end

  def put_postconditions(
        %__MODULE__{postconditions: existing_postconditions} = annotated_function,
        postconditions
      )
      when is_list(postconditions) do
    # Check to make sure each element in the list is a `Bond.Assertion` struct.
    Enum.each(postconditions, fn %Assertion{} -> :ok end)

    %{annotated_function | postconditions: existing_postconditions ++ postconditions}
  end

  def put_invariants(
        %__MODULE__{invariants: existing_invariants} = annotated_function,
        invariants
      )
      when is_list(invariants) do
    Enum.each(invariants, fn %Assertion{kind: :invariant} -> :ok end)

    %{annotated_function | invariants: existing_invariants ++ invariants}
  end

  def put_doc_attributes(
        %__MODULE__{doc_attributes: existing_doc_attributes} = annotated_function,
        doc_attributes
      )
      when is_list(doc_attributes) do
    %{annotated_function | doc_attributes: existing_doc_attributes ++ doc_attributes}
  end

  def has_preconditions?(%__MODULE__{preconditions: preconditions}),
    do: not Enum.empty?(preconditions)

  def has_postconditions?(%__MODULE__{postconditions: postconditions}),
    do: not Enum.empty?(postconditions)

  def has_invariants?(%__MODULE__{invariants: invariants}),
    do: not Enum.empty?(invariants)

  def has_doc_attributes?(%__MODULE__{doc_attributes: doc_attributes}),
    do: not Enum.empty?(doc_attributes)

  def override?(%__MODULE__{} = annotated_function) do
    has_preconditions?(annotated_function) or
      has_postconditions?(annotated_function) or
      (annotated_function.kind == :def and has_invariants?(annotated_function))
  end

  @typedoc """
  Per-kind configuration mode controlling how `apply_contract/2` emits each contract kind:

    * `true` — emit the override with a runtime guard that defaults to "evaluate."
    * `false` — emit the override with a runtime guard that defaults to "do not evaluate."
    * `:purge` — emit no code for this kind. The override may still be emitted for the
      *other* kind, in which case the doc section for the purged kind is also omitted.

  `true` and `false` both produce code that reads `Application.get_env(:bond, <kind>, default)`
  at every call, where `default` is the compile-time value. Anything other than `false` at
  runtime causes evaluation to occur, so the toggle is "set to `false` to disable."
  """
  @type mode :: true | false | :purge

  @typedoc """
  Configuration controlling which kinds of contracts `apply_contract/2` emits and how.

  Used by `Bond.Compiler.__before_compile__/1` to honour the `:bond` application config keys
  `:preconditions` and `:postconditions` (and any per-module `:overrides`).
  """
  @type contract_config :: %{
          required(:preconditions) => mode(),
          required(:postconditions) => mode(),
          optional(:invariants) => mode()
        }

  @doc """
  Returns a quoted expression that wraps the annotated function with its contract, or `nil`
  when nothing needs to be emitted.

  Per-kind mode (see `t:mode/0`):

    * `true` / `false` — the override IS emitted; `Bond.Runtime.Eval` performs a runtime
      `Application.get_env(:bond, <kind>, <compile_time_value>)` check and skips evaluation
      when the result is exactly `false`. The auto-generated doc section for that kind
      appears.
    * `:purge` — no code is emitted for that kind. The doc section is omitted.

  When both `:preconditions` and `:postconditions` resolve to `:purge` (or when the function
  has no contracts of either kind), `apply_contract/2` returns `nil`. The caller filters
  `nil`s out and the user's `def`/`defp` runs as written, with zero per-call overhead.

  When an override IS emitted, the expression contains:

    1. A `defoverridable` declaration making the function overridable.
    2. Zero or more `@doc` clauses re-emitting the user's `@doc` attributes, with the
       auto-generated `#### Preconditions` / `#### Postconditions` sections appended (filtered
       by the per-kind mode).
    3. A single override clause for the function that:

         * (when `preconditions != :purge`) calls `Bond.Runtime.Eval.evaluate_preconditions/2`
           with a thin closure that delegates to the lifted precondition defp;
         * resolves any `old(...)` expressions found in the postconditions into local bindings;
         * delegates to the original implementation via `super(...)`, capturing the result;
         * (when `postconditions != :purge`) calls `Bond.Runtime.Eval.evaluate_postconditions/2`
           with a thin closure that delegates to the lifted postcondition defp (passing the
           function params, the captured result, and the resolved old-value bindings);
         * returns the captured result.

    4. One private `defp` per non-purged kind containing the assertion-evaluation block
       produced by `Bond.Compiler.Assertion.assertions_body/2`. Naming convention:
       `:"__bond_preconditions__\#{fun}__\#{arity}"` /
       `:"__bond_postconditions__\#{fun}__\#{arity}"`. Lifting the closure body into a named
       defp keeps the override clause itself tiny and avoids re-emitting the full assertion
       AST inline.

  The override clause uses the parameter names from the function's first clause. For
  multi-clause functions Elixir's normal pattern matching applies inside `super(...)`, so a
  single wrapper clause covers every original clause.
  """
  @spec apply_contract(t(), contract_config()) :: Macro.t() | nil
  def apply_contract(annotated_function, config \\ %{preconditions: true, postconditions: true})

  def apply_contract(%__MODULE__{} = annotated_function, config) do
    pre_mode =
      resolve_mode(Map.fetch!(config, :preconditions), has_preconditions?(annotated_function))

    post_mode =
      resolve_mode(Map.fetch!(config, :postconditions), has_postconditions?(annotated_function))

    inv_mode =
      Invariants.resolve_mode(
        Map.get(config, :invariants, true),
        annotated_function.kind,
        annotated_function.invariants
      )

    if pre_mode != :purge or post_mode != :purge or inv_mode != :purge do
      build_contract_override(annotated_function, pre_mode, post_mode, inv_mode)
    end
  end

  # A kind is effectively purged if either the user purged it OR the function has no
  # assertions of that kind — there's nothing to evaluate or document either way.
  defp resolve_mode(:purge, _has_assertions?), do: :purge
  defp resolve_mode(_value, false), do: :purge
  defp resolve_mode(value, true) when value in [true, false], do: value

  defp build_contract_override(
         %__MODULE__{kind: kind, fun: fun, arity: arity, module: struct_module} =
           annotated_function,
         pre_mode,
         post_mode,
         inv_mode
       ) do
    first_clause = List.first(annotated_function.clauses)
    function_info = {fun, arity}

    # `strip_default_args` removes `\\` defaults so the override is a plain arity-N def;
    # `params_split/3` then partitions into head_params (for def/defp heads) and
    # super_args (for super/eval call sites). See Invariants.rewrite_call_params/2.
    clean_params = strip_default_args(first_clause.params)

    {struct_params, head_params, super_args} =
      Invariants.params_split(clean_params, first_clause, inv_mode)

    {postconditions, old_context} =
      if post_mode != :purge do
        OldExpression.precompile(annotated_function.postconditions)
      else
        {[], %{}}
      end

    old_pairs = OldExpression.pairs(old_context)
    old_assignments = OldExpression.resolve(old_context)

    pre_fn_name = lifted_fn_name(:preconditions, fun, arity)
    post_fn_name = lifted_fn_name(:postconditions, fun, arity)
    inv_fn_name = lifted_fn_name(:invariants, fun, arity)

    doc_asts = ContractDocs.doc_clauses(annotated_function, first_clause.env, pre_mode, post_mode)

    pre_invariant_stmts =
      Invariants.all_pre_invariant_stmts(
        inv_fn_name,
        struct_params,
        inv_mode,
        pre_mode,
        post_mode
      )

    post_invariant_stmts =
      Invariants.post_invariant_stmts(inv_fn_name, inv_mode, struct_module, pre_mode, post_mode)

    body_stmts =
      build_override_body(
        super_args,
        pre_fn_name,
        post_fn_name,
        old_assignments,
        old_pairs,
        pre_mode,
        post_mode,
        pre_invariant_stmts,
        post_invariant_stmts
      )

    assertion_defs =
      Enum.reject(
        [
          maybe_build_assertion_defp(
            pre_fn_name,
            head_params,
            [],
            annotated_function.preconditions,
            function_info,
            first_clause.env,
            pre_mode
          ),
          maybe_build_assertion_defp(
            post_fn_name,
            head_params,
            postcondition_extra_params(old_pairs),
            postconditions,
            function_info,
            first_clause.env,
            post_mode
          ),
          Invariants.build_lifted_defp(
            inv_fn_name,
            annotated_function.invariants,
            function_info,
            first_clause.env,
            inv_mode
          )
        ],
        &is_nil/1
      )

    quote file: first_clause.env.file, line: first_clause.env.line do
      defoverridable([{unquote(fun), unquote(arity)}])

      unquote_splicing(doc_asts)

      unquote(kind)(unquote(fun)(unquote_splicing(head_params))) do
        (unquote_splicing(body_stmts))
      end

      unquote_splicing(assertion_defs)
    end
  end

  defp lifted_fn_name(kind, fun, arity) do
    :"__bond_#{kind}__#{fun}__#{arity}"
  end

  # Extra parameters that come after the original function's params in the lifted
  # postcondition defp: the captured `result` and one parameter per resolved `old(...)` value,
  # in the order produced by `OldExpression.pairs/1`. All wrapped in `var!/1` so they bind
  # with the calling function's hygiene context (matching how the assertion expressions and
  # `OldExpression.resolve/1` reference them).
  defp postcondition_extra_params(old_pairs) do
    result_param = quote(do: var!(result))
    old_params = for {var, _expression} <- old_pairs, do: quote(do: var!(unquote(var)))
    [result_param | old_params]
  end

  defp build_override_body(
         call_params,
         pre_fn_name,
         post_fn_name,
         old_assignments,
         old_pairs,
         pre_mode,
         post_mode,
         pre_invariant_stmts,
         post_invariant_stmts
       ) do
    pre_stmts = pre_eval_stmts(call_params, pre_fn_name, pre_mode)
    super_call = quote(do: var!(result) = super(unquote_splicing(call_params)))
    post_stmts = post_eval_stmts(call_params, post_fn_name, old_pairs, post_mode, pre_mode)
    return_stmt = quote(do: var!(result))

    pre_invariant_stmts ++
      pre_stmts ++
      old_assignments ++
      [super_call] ++ post_stmts ++ post_invariant_stmts ++ [return_stmt]
  end

  defp pre_eval_stmts(_call_params, _name, :purge), do: []

  defp pre_eval_stmts(call_params, name, mode) do
    [
      quote do
        if Bond.Runtime.Eval.should_evaluate?(:preconditions, unquote(mode)) do
          Bond.Runtime.Eval.evaluate_preconditions(fn ->
            unquote(name)(unquote_splicing(call_params))
          end)
        end
      end
    ]
  end

  defp post_eval_stmts(_call_params, _name, _old_pairs, :purge, _pre_mode), do: []

  defp post_eval_stmts(call_params, name, old_pairs, mode, pre_mode) do
    old_args = for {var, _expression} <- old_pairs, do: quote(do: var!(unquote(var)))
    args = call_params ++ [quote(do: var!(result)) | old_args]
    chain = %{preconditions: pre_mode}

    [
      quote do
        if Bond.Runtime.Eval.should_evaluate?(
             :postconditions,
             unquote(mode),
             unquote(Macro.escape(chain))
           ) do
          Bond.Runtime.Eval.evaluate_postconditions(fn ->
            unquote(name)(unquote_splicing(args))
          end)
        end
      end
    ]
  end

  defp maybe_build_assertion_defp(_name, _params, _extra, _assertions, _info, _env, :purge),
    do: nil

  defp maybe_build_assertion_defp(
         name,
         call_params,
         extra_params,
         assertions,
         function_info,
         env,
         _mode
       ) do
    body = Assertion.assertions_body(assertions, function_info)
    params = call_params ++ extra_params

    quote file: env.file, line: env.line do
      defp unquote(name)(unquote_splicing(params)) do
        unquote(body)
      end
    end
  end

  defp strip_default_args(params) do
    Enum.map(params, fn
      {:\\, _meta, [param, _default]} -> param
      other -> other
    end)
  end
end