defmodule Mix.Tasks.ExOkf.Context do
  @shortdoc "Assemble AI context for a concept"
  @moduledoc """
  Builds provenance-aware context around a concept and prints a prompt block.

      mix ex_okf.context path/to/bundle tables/orders
      mix ex_okf.context path/to/bundle tables/orders --depth 2 --token-budget 2000
  """

  use Mix.Task

  @impl Mix.Task
  def run(args) do
    {opts, rest, _} =
      OptionParser.parse(args,
        strict: [depth: :integer, token_budget: :integer],
        aliases: [d: :depth]
      )

    [path, concept_id | _] =
      case rest do
        [p, c | _] -> [p, c]
        _ -> Mix.raise("Usage: mix ex_okf.context PATH CONCEPT_ID")
      end

    ctx_opts =
      []
      |> then(fn o -> if opts[:depth], do: Keyword.put(o, :depth, opts[:depth]), else: o end)
      |> then(fn o ->
        if opts[:token_budget], do: Keyword.put(o, :token_budget, opts[:token_budget]), else: o
      end)

    {:ok, bundle} = ExOKF.load(path)

    case ExOKF.context(bundle, concept_id, ctx_opts) do
      {:ok, ctx} ->
        Mix.shell().info(ExOKF.Context.to_prompt(ctx))
        Mix.shell().info("tokens_used=#{ctx.tokens_used} truncated=#{ctx.truncated}")
        :ok

      {:error, reason} ->
        Mix.raise("Failed to build context: #{inspect(reason)}")
    end
  end
end