zen_quant

defmodule ZenQuant.OrderState do
  @moduledoc """
  Immutable order lifecycle tracking — pure functions for order state transitions.

  Tracks order creation, partial fills, VWAP accumulation, and cancellation
  without any exchange coupling. Agents feed fill events and get back
  updated state maps they can store however they choose.

  ## Terminology

    * **VWAP** - Volume-Weighted Average Price across fills
    * **Terminal** - An order in `:filled` or `:cancelled` status (no further transitions)

  ## Status Flow

      :pending → :partial → :filled
      :pending → :cancelled
      :partial → :cancelled (preserves fill history)

  ## Example

      {:ok, order} = ZenQuant.OrderState.new(%{
        id: "order-123", symbol: "BTC/USDT:USDT",
        side: :buy, type: :limit, price: 67800.0, quantity: 1.0
      })

      {:ok, order} = ZenQuant.OrderState.on_fill(order, %{price: 67800.0, quantity: 0.5})
      order.status
      # => :partial

      ZenQuant.OrderState.summary(order)
      # => %{id: "order-123", status: :partial, fill_pct: 50.0, ...}

  """

  use Descripex, namespace: "/order-state"

  @valid_sides [:buy, :sell]
  @valid_types [:limit, :market, :stop, :stop_limit]
  @terminal_statuses [:filled, :cancelled]
  @fill_tolerance 1.0e-9

  @string_to_type %{
    "limit" => :limit,
    "market" => :market,
    "stop" => :stop,
    "stop_limit" => :stop_limit
  }

  @typedoc "Order state map tracking lifecycle"
  @type order_state :: %{
          id: String.t(),
          symbol: String.t(),
          side: :buy | :sell,
          type: :limit | :market | :stop | :stop_limit,
          price: number() | nil,
          quantity: float(),
          status: :pending | :partial | :filled | :cancelled,
          filled_qty: float(),
          remaining_qty: float(),
          avg_price: float(),
          fills: [map()],
          created_at: DateTime.t(),
          updated_at: DateTime.t()
        }

  @typedoc "Order summary for agent consumption"
  @type summary_result :: %{
          id: String.t(),
          status: atom(),
          fill_pct: float(),
          avg_price: float(),
          fill_count: non_neg_integer(),
          remaining_qty: float()
        }

  api(:new, "Create initial order state from params map.",
    params: [
      params: [
        kind: :value,
        description: "Map with :id, :symbol, :side, :type, :price, :quantity. Sides/types accept strings or atoms."
      ],
      opts: [kind: :value, default: [], description: "Options: now (DateTime override for testing)"]
    ],
    returns: %{
      type: :tuple,
      description: "{:ok, order_state} or {:error, reason} for validation failures"
    },
    returns_example:
      {:ok,
       %{
         id: "order-123",
         symbol: "BTC/USDT:USDT",
         side: :buy,
         type: :limit,
         price: 67_800.0,
         quantity: 1.0,
         status: :pending,
         filled_qty: 0.0,
         remaining_qty: 1.0,
         avg_price: 0.0,
         fills: [],
         created_at: ~U[2026-02-25 12:00:00Z],
         updated_at: ~U[2026-02-25 12:00:00Z]
       }},
    errors: [:missing_required_field, :invalid_side, :invalid_type, :invalid_quantity, :missing_price, :invalid_price]
  )

  @spec new(map(), keyword()) :: {:ok, order_state()} | {:error, atom()}
  def new(params, opts \\ []) when is_map(params) do
    now = Keyword.get(opts, :now, DateTime.utc_now())

    with {:ok, id} <- require_field(params, :id),
         {:ok, symbol} <- require_field(params, :symbol),
         {:ok, side} <- normalize_side(params),
         {:ok, type} <- normalize_type(params),
         {:ok, quantity} <- validate_quantity(params),
         :ok <- validate_price(params, type) do
      price = get_field(params, :price)

      {:ok,
       %{
         id: id,
         symbol: symbol,
         side: side,
         type: type,
         price: price,
         quantity: quantity * 1.0,
         status: :pending,
         filled_qty: 0.0,
         remaining_qty: quantity * 1.0,
         avg_price: 0.0,
         fills: [],
         created_at: now,
         updated_at: now
       }}
    end
  end

  api(:on_fill, "Apply a fill event to an order, returning updated state with VWAP.",
    params: [
      order: [kind: :value, description: "Current order state map from new/2 or previous on_fill/3"],
      fill: [
        kind: :exchange_data,
        source: "WebSocket fill events or fetch_order",
        description: "Map with :price and :quantity for this fill"
      ],
      opts: [kind: :value, default: [], description: "Options: now (DateTime override for testing)"]
    ],
    returns: %{
      type: :tuple,
      description: "{:ok, updated_order} with status :partial or :filled, or {:error, :already_terminal | :invalid_fill}"
    },
    returns_example: {:ok, %{status: :partial, filled_qty: 0.5, remaining_qty: 0.5, avg_price: 67_800.0}},
    errors: [:already_terminal, :invalid_fill]
  )

  @spec on_fill(order_state(), map(), keyword()) :: {:ok, order_state()} | {:error, atom()}
  def on_fill(order, fill, opts \\ []) when is_map(order) and is_map(fill) do
    now = Keyword.get(opts, :now, DateTime.utc_now())

    with :ok <- check_not_terminal(order),
         {:ok, fill_price, fill_qty} <- validate_fill(fill) do
      apply_fill(order, fill_price, fill_qty, now)
    end
  end

  api(:cancel, "Cancel a non-terminal order, preserving any partial fill state.",
    params: [
      order: [kind: :value, description: "Current order state map"],
      opts: [kind: :value, default: [], description: "Options: now (DateTime override for testing)"]
    ],
    returns: %{
      type: :tuple,
      description: "{:ok, cancelled_order} or {:error, :already_terminal}"
    },
    returns_example: {:ok, %{status: :cancelled, remaining_qty: 0.0}},
    errors: [:already_terminal]
  )

  @spec cancel(order_state(), keyword()) :: {:ok, order_state()} | {:error, atom()}
  def cancel(order, opts \\ []) when is_map(order) do
    now = Keyword.get(opts, :now, DateTime.utc_now())

    case check_not_terminal(order) do
      :ok ->
        {:ok,
         %{
           order
           | status: :cancelled,
             remaining_qty: 0.0,
             updated_at: now
         }}

      error ->
        error
    end
  end

  api(:summary, "Structured summary of order state for agent consumption.",
    params: [
      order: [kind: :value, description: "Current order state map"]
    ],
    returns: %{
      type: :map,
      description: "Map with :id, :status, :fill_pct, :avg_price, :fill_count, :remaining_qty"
    },
    returns_example: %{
      id: "order-123",
      status: :partial,
      fill_pct: 50.0,
      avg_price: 67_800.0,
      fill_count: 1,
      remaining_qty: 0.5
    }
  )

  @spec summary(order_state()) :: summary_result()
  def summary(order) when is_map(order) do
    fill_pct =
      if order.quantity > 0 do
        order.filled_qty / order.quantity * 100.0
      else
        0.0
      end

    %{
      id: order.id,
      status: order.status,
      fill_pct: fill_pct,
      avg_price: order.avg_price,
      fill_count: length(order.fills),
      remaining_qty: order.remaining_qty
    }
  end

  # --- Private Helpers ---

  # Reads a field from a map, trying atom key first then string fallback
  defp get_field(map, key) when is_atom(key) do
    case Map.get(map, key) do
      nil -> Map.get(map, Atom.to_string(key))
      val -> val
    end
  end

  # Requires a field to be present and non-nil
  defp require_field(params, key) do
    case get_field(params, key) do
      nil -> {:error, :missing_required_field}
      val -> {:ok, val}
    end
  end

  # Normalizes side field to atom, accepting strings and atoms
  defp normalize_side(params) do
    case get_field(params, :side) do
      side when side in @valid_sides -> {:ok, side}
      "buy" -> {:ok, :buy}
      "sell" -> {:ok, :sell}
      nil -> {:error, :missing_required_field}
      _ -> {:error, :invalid_side}
    end
  end

  # Normalizes type field to atom, accepting strings and atoms
  defp normalize_type(params) do
    case get_field(params, :type) do
      type when type in @valid_types ->
        {:ok, type}

      type when is_binary(type) ->
        case Map.fetch(@string_to_type, type) do
          {:ok, atom} -> {:ok, atom}
          :error -> {:error, :invalid_type}
        end

      nil ->
        {:error, :missing_required_field}

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

  # Validates quantity is positive
  defp validate_quantity(params) do
    case get_field(params, :quantity) do
      q when is_number(q) and q > 0 -> {:ok, q}
      nil -> {:error, :missing_required_field}
      _ -> {:error, :invalid_quantity}
    end
  end

  # Validates price: required for limit/stop_limit, optional for market/stop
  defp validate_price(params, type) when type in [:limit, :stop_limit] do
    case get_field(params, :price) do
      p when is_number(p) and p > 0 -> :ok
      nil -> {:error, :missing_price}
      _ -> {:error, :invalid_price}
    end
  end

  defp validate_price(_params, _type), do: :ok

  # Checks that the order is not in a terminal status
  defp check_not_terminal(%{status: status}) when status in @terminal_statuses do
    {:error, :already_terminal}
  end

  defp check_not_terminal(_order), do: :ok

  # Validates a fill event has positive price and quantity
  defp validate_fill(fill) do
    price = get_field(fill, :price)
    qty = get_field(fill, :quantity)

    if is_number(price) and price > 0 and is_number(qty) and qty > 0 do
      {:ok, price * 1.0, qty * 1.0}
    else
      {:error, :invalid_fill}
    end
  end

  # Applies a validated fill to the order, computing VWAP and new status
  defp apply_fill(order, fill_price, fill_qty, now) do
    capped_qty = min(fill_qty, order.remaining_qty)
    new_filled = order.filled_qty + capped_qty
    new_remaining = order.quantity - new_filled

    new_avg =
      if new_filled > 0 do
        (order.avg_price * order.filled_qty + fill_price * capped_qty) / new_filled
      else
        0.0
      end

    new_status = if order.quantity - new_filled < @fill_tolerance, do: :filled, else: :partial

    fill_record = %{
      price: fill_price,
      quantity: capped_qty,
      timestamp: now
    }

    {:ok,
     %{
       order
       | status: new_status,
         filled_qty: new_filled,
         remaining_qty: max(new_remaining, 0.0),
         avg_price: new_avg,
         fills: order.fills ++ [fill_record],
         updated_at: now
     }}
  end
end