pi_bridge

defmodule Pi.Transport.Stdio do
  @moduledoc "Line-delimited JSON transport for extension-owned BEAM sessions."

  require Logger
  require Pi.Features

  alias Pi.Bridge.Info
  alias Pi.LLM.Broker
  alias Pi.MCP.Tools
  alias Pi.Plugin.Event
  alias Pi.Plugin.Manager
  alias Pi.Protocol.Call
  alias Pi.Protocol.LLM.Chunk
  alias Pi.Protocol.LLM.Done
  alias Pi.Protocol.LLM.Error
  alias Pi.Protocol.PluginHook
  alias Pi.Protocol.PluginHookResponse
  alias Pi.Protocol.Ready
  alias Pi.Protocol.Response
  alias Pi.Protocol.Result
  alias Pi.Skill.Loader
  alias Pi.Transport
  alias Pi.Transport.Payload

  defdelegate emit_request(id, op, payload), to: Transport, as: :request
  defdelegate emit(payload), to: Transport

  @doc false
  def __test_payload__(payload), do: Payload.encode(payload)

  @doc false
  def __test_dispatch__(name, args), do: dispatch(name, args)

  @doc false
  def __test_handle_line__(line), do: handle_line(line)

  @doc false
  def __test_safe_call__(name, fun) when is_binary(name) and is_function(fun, 0) do
    safe_call("pi stdio test call failed: #{name}", fun)
  end

  def start(input \\ IO.stream(:stdio, :line)) do
    Transport.register(self())

    try do
      Event.install()
      if Pi.Features.plugins?(), do: Manager.install()
      Broker.install()
      Pi.Eval.Supervisor.install()
      ready()
      start_input_reader(input)
      loop()
    after
      clear_transport_owner()
    end
  end

  defp start_input_reader(input) do
    owner = self()

    spawn_link(fn ->
      Enum.each(input, &send(owner, {:stdin, &1}))
      send(owner, :stdin_closed)
    end)
  end

  defp clear_transport_owner, do: Transport.unregister(self())

  defp loop do
    receive do
      {:stdin, line} ->
        handle_line(line)
        loop()

      {:pi_transport_emit, payload} ->
        write(payload)
        loop()

      :stdin_closed ->
        :ok
    end
  end

  defp handle_line(line) do
    case decode_line(line) do
      {:ok, payload} -> handle_payload(payload)
      :ignore -> :ok
    end
  end

  defp decode_line(line) do
    case Jason.decode(line) do
      {:ok, %{"type" => "call"} = payload} -> decode_payload(Call, payload)
      {:ok, %{"type" => "response"} = payload} -> decode_payload(Response, payload)
      {:ok, %{"type" => "llm_chunk"} = payload} -> decode_payload(Chunk, payload)
      {:ok, %{"type" => "llm_done"} = payload} -> decode_payload(Done, payload)
      {:ok, %{"type" => "llm_error"} = payload} -> decode_payload(Error, payload)
      _ -> :ignore
    end
  end

  defp decode_payload(module, payload) do
    case module.from_map(payload) do
      {:ok, payload} -> {:ok, payload}
      {:error, _reason} -> :ignore
    end
  end

  defp handle_payload(%Call{} = call) do
    Task.Supervisor.start_child(Pi.Transport.TaskSupervisor, fn ->
      safe_respond(call.id, call.name, safe_dispatch(call))
    end)
  end

  defp handle_payload(%Response{} = response), do: Broker.deliver(response.id, response)
  defp handle_payload(%Chunk{} = chunk), do: Broker.deliver_stream(chunk.id, :chunk, chunk.delta)
  defp handle_payload(%Done{} = done), do: Broker.deliver_stream(done.id, :done, done.result)
  defp handle_payload(%Error{} = error), do: Broker.deliver_stream(error.id, :error, error.error)

  defp safe_dispatch(%Call{} = call) do
    safe_call("pi stdio tool call failed: #{call.name}", fn ->
      dispatch(call.name, call.arguments)
    end)
  end

  defp safe_call(log_prefix, fun) when is_binary(log_prefix) and is_function(fun, 0) do
    fun.()
  catch
    kind, reason ->
      message = Exception.format(kind, reason, __STACKTRACE__)
      Logger.error("#{log_prefix}\n#{message}")
      {:error, message}
  end

  defp dispatch("pi_skills_list", _args) do
    if Pi.Features.skills?(),
      do: {:ok, encode_structs(Loader.serializable())},
      else: {:ok, "[]"}
  end

  defp dispatch("pi_event", args) do
    Event.push(args)
    if Pi.Features.plugins?(), do: Manager.dispatch_event(args)
    {:ok, "ok"}
  end

  defp dispatch("pi_plugin_command", %{"name" => name, "args" => args}) when is_binary(name) do
    Pi.Features.gate :plugins do
      name
      |> run_plugin_command(to_string(args || ""))
      |> encode_reply()
    end
    |> encode_plugin_command_reply()
  end

  defp dispatch("pi_plugin_tool_call", args) do
    Pi.Features.gate :plugins do
      case PluginHook.from_wire(args) do
        {:ok, hook} ->
          hook
          |> plugin_hook_payload()
          |> Manager.tool_call(hook.context)
          |> encode_hook_reply()

        {:error, _reason} ->
          encode_hook_reply({:error, "Invalid plugin hook payload"})
      end
    end
    |> encode_plugin_hook_reply()
  end

  defp dispatch("pi_plugin_tool_result", args) do
    Pi.Features.gate :plugins do
      case PluginHook.from_wire(args) do
        {:ok, hook} ->
          hook
          |> plugin_hook_payload()
          |> Manager.tool_result(hook.context)
          |> encode_hook_reply()

        {:error, _reason} ->
          encode_hook_reply({:error, "Invalid plugin hook payload"})
      end
    end
    |> encode_plugin_hook_reply()
  end

  defp dispatch("pi_bridge_info", _args) do
    {:ok, Jason.encode!(Info.snapshot(:stdio))}
  end

  defp dispatch("pi_bridge_apis", _args) do
    {:ok, encode_structs(Info.apis())}
  end

  defp dispatch("pi_apis", _args) do
    {:ok, encode_structs(Info.runtime_apis())}
  end

  defp dispatch(name, args), do: Tools.dispatch(name, args)

  defp run_plugin_command("quack" <> _rest = name, args) do
    case Manager.run_command(name, args) do
      {:error, "Unknown plugin command: " <> _command} ->
        _ = Manager.load(Pi.Mirror.QuackDB)
        Manager.run_command(name, args)

      reply ->
        reply
    end
  end

  defp run_plugin_command(name, args), do: Manager.run_command(name, args)

  defp plugin_hook_payload(%PluginHook{} = hook) do
    %{
      "toolName" => hook.tool_name,
      "toolCallId" => hook.tool_call_id,
      "input" => hook.input,
      "content" => hook.content,
      "isError" => hook.is_error
    }
  end

  defp encode_plugin_command_reply({:error, message}), do: encode_reply({:error, message})
  defp encode_plugin_command_reply(reply), do: reply

  defp encode_plugin_hook_reply({:error, message}), do: encode_hook_reply({:error, message})
  defp encode_plugin_hook_reply(reply), do: reply

  defp encode_hook_reply({:ok, value}) when is_map(value) do
    encode_reply(PluginHookResponse.ok(value))
  end

  defp encode_hook_reply({:block, reason}), do: encode_reply(PluginHookResponse.block(reason))
  defp encode_hook_reply({:error, reason}), do: encode_reply(PluginHookResponse.error(reason))
  defp encode_hook_reply(:ok), do: encode_reply(PluginHookResponse.ok())
  defp encode_hook_reply(value), do: encode_reply(value)

  defp encode_reply(reply), do: {:ok, Jason.encode!(reply |> reply_payload() |> Payload.encode())}

  defp reply_payload(%_module{} = value), do: Payload.encode(value)
  defp reply_payload({:ok, value}) when is_map(value), do: %{ok: value}
  defp reply_payload({:ok, value}) when is_binary(value), do: %{ok: value}
  defp reply_payload({:error, value}), do: %{error: value}
  defp reply_payload({:block, value}), do: %{block: value}
  defp reply_payload(:ok), do: %{ok: %{}}
  defp reply_payload(value) when is_binary(value), do: %{ok: value}
  defp reply_payload(value) when is_map(value), do: value

  defp safe_respond(id, name, reply) do
    respond(id, reply)
  catch
    kind, reason ->
      message = Exception.format(kind, reason, __STACKTRACE__)
      Logger.error("pi stdio tool response failed: #{name}\n#{message}")
      respond(id, {:error, message})
  end

  defp respond(id, {:ok, text}) do
    write(%Result{type: :result, id: id, text: text, is_error: false})
  end

  defp respond(id, {:error, message}) do
    write(%Result{type: :result, id: id, text: message, is_error: true})
  end

  defp ready, do: write(%Ready{type: :ready, info: Info.snapshot(:stdio)})

  defp encode_structs(values), do: values |> Enum.map(&Payload.encode/1) |> Jason.encode!()

  defp write(payload), do: IO.write([Jason.encode!(Payload.encode(payload)), ?\n])
end