defmodule Mix.Interactive.SupervisedShell do
  @moduledoc """
  Supervised shell for interactive MCP sessions with automatic restart capability.

  This module provides a robust shell that monitors both the transport and client
  processes, automatically restarting them if they crash due to timeouts or other
  errors.
  """

  alias Hermes.Client
  alias Mix.Interactive.CLI
  alias Mix.Interactive.Commands
  alias Mix.Interactive.UI

  defstruct [
    :transport_module,
    :transport_opts,
    :client_opts,
    :transport_pid,
    :client_pid,
    :restart_count,
    :max_restarts
  ]

  @max_restarts 3
  @restart_delay 1000

  @doc """
  Starts a supervised interactive shell session.

  ## Options
    * `:transport_module` - The transport module (e.g., Hermes.Transport.StreamableHTTP)
    * `:transport_opts` - Options for starting the transport
    * `:client_opts` - Options for starting the client (excluding transport)
    * `:max_restarts` - Maximum number of automatic restarts (default: 3)
  """
  def start(opts) do
    state = %__MODULE__{
      transport_module: Keyword.fetch!(opts, :transport_module),
      transport_opts: Keyword.fetch!(opts, :transport_opts),
      client_opts: Keyword.fetch!(opts, :client_opts),
      restart_count: 0,
      max_restarts: Keyword.get(opts, :max_restarts, @max_restarts)
    }

    Process.flag(:trap_exit, true)

    case start_processes(state) do
      {:ok, state} ->
        IO.puts("\nType #{UI.colors().command}help#{UI.colors().reset} for available commands\n")
        supervised_loop(state)

      {:error, reason} ->
        IO.puts("#{UI.colors().error}Failed to start processes: #{inspect(reason)}#{UI.colors().reset}")
        {:error, reason}
    end
  end

  defp supervised_loop(%{client_pid: client} = state) do
    if Process.alive?(client) do
      IO.write("#{UI.colors().prompt}mcp> #{UI.colors().reset}")

      parent = self()

      input_pid =
        spawn_link(fn ->
          line = IO.gets("")
          send(parent, {:input, line})
        end)

      receive do
        {:EXIT, pid, reason} when pid == state.client_pid or pid == state.transport_pid ->
          Process.exit(input_pid, :kill)
          handle_process_exit(pid, reason, state)

        {:EXIT, ^input_pid, _reason} ->
          supervised_loop(state)

        {:input, :eof} ->
          :ok

        {:input, line} ->
          line
          |> String.trim()
          |> Commands.process_command(client, fn -> supervised_loop(state) end)
      end
    else
      handle_process_exit(state.client_pid, :noproc, state)
    end
  end

  defp handle_process_exit(pid, reason, state) do
    cond do
      pid == state.client_pid ->
        IO.puts("\n#{UI.colors().error}✗ Client process crashed: #{format_exit_reason(reason)}#{UI.colors().reset}")
        handle_restart(state)

      pid == state.transport_pid ->
        IO.puts("\n#{UI.colors().error}✗ Transport process crashed: #{format_exit_reason(reason)}#{UI.colors().reset}")
        handle_restart(state)

      true ->
        supervised_loop(state)
    end
  end

  defp format_exit_reason({:timeout, _}), do: "Request timeout"
  defp format_exit_reason({:error, %{reason: reason}}), do: "Error: #{reason}"
  defp format_exit_reason(:normal), do: "Normal termination"
  defp format_exit_reason(:shutdown), do: "Shutdown"
  defp format_exit_reason(reason), do: inspect(reason)

  defp handle_restart(state) do
    if state.restart_count < state.max_restarts do
      IO.puts(
        "#{UI.colors().info}→ Attempting automatic restart (#{state.restart_count + 1}/#{state.max_restarts})...#{UI.colors().reset}"
      )

      Process.sleep(@restart_delay)

      cleanup_processes(state)

      new_state = %{state | restart_count: state.restart_count + 1}

      case start_processes(new_state) do
        {:ok, restarted_state} ->
          IO.puts("#{UI.colors().success}✓ Successfully restarted#{UI.colors().reset}")
          IO.write("\n#{UI.colors().prompt}mcp> #{UI.colors().reset}")
          supervised_loop(restarted_state)

        {:error, reason} ->
          IO.puts("#{UI.colors().error}✗ Restart failed: #{inspect(reason)}#{UI.colors().reset}")
          offer_manual_restart(new_state)
      end
    else
      IO.puts("#{UI.colors().error}✗ Maximum restart attempts reached#{UI.colors().reset}")
      offer_manual_restart(state)
    end
  end

  defp offer_manual_restart(state) do
    IO.puts("\n#{UI.colors().info}Options:#{UI.colors().reset}")
    IO.puts("  #{UI.colors().command}r#{UI.colors().reset} - Retry connection")
    IO.puts("  #{UI.colors().command}q#{UI.colors().reset} - Quit")
    IO.write("\n#{UI.colors().prompt}Choice: #{UI.colors().reset}")

    case "" |> IO.gets() |> String.trim() |> String.downcase() do
      "r" ->
        IO.puts("#{UI.colors().info}→ Retrying connection...#{UI.colors().reset}")
        cleanup_processes(state)
        new_state = %{state | restart_count: 0}

        case start_processes(new_state) do
          {:ok, restarted_state} ->
            IO.puts("#{UI.colors().success}✓ Successfully reconnected#{UI.colors().reset}")
            supervised_loop(restarted_state)

          {:error, reason} ->
            IO.puts("#{UI.colors().error}✗ Retry failed: #{inspect(reason)}#{UI.colors().reset}")
            offer_manual_restart(new_state)
        end

      "q" ->
        IO.puts("#{UI.colors().info}Exiting...#{UI.colors().reset}")
        cleanup_processes(state)
        :ok

      _ ->
        offer_manual_restart(state)
    end
  end

  defp start_processes(state) do
    # Start client first - it will hibernate waiting for transport's :initialize message
    with {:ok, client_pid} <- start_client(state),
         {:ok, transport_pid} <- start_transport(state) do
      Process.monitor(transport_pid)
      Process.monitor(client_pid)

      IO.puts("#{UI.colors().info}• Checking connection...#{UI.colors().reset}")
      CLI.check_client_connection(client_pid)

      {:ok, %{state | transport_pid: transport_pid, client_pid: client_pid}}
    end
  end

  defp start_transport(%{transport_module: module, transport_opts: opts}) do
    IO.puts("#{UI.colors().info}• Starting transport...#{UI.colors().reset}")

    case module.start_link(opts) do
      {:ok, pid} ->
        IO.puts("#{UI.colors().success}✓ Transport started#{UI.colors().reset}")
        {:ok, pid}

      {:error, {:already_started, pid}} ->
        IO.puts("#{UI.colors().info}• Transport already running#{UI.colors().reset}")
        {:ok, pid}

      error ->
        error
    end
  end

  defp start_client(%{client_opts: opts}) do
    IO.puts("#{UI.colors().info}• Starting client...#{UI.colors().reset}")

    case Client.start_link(opts) do
      {:ok, pid} ->
        IO.puts("#{UI.colors().success}✓ Client started#{UI.colors().reset}")
        {:ok, pid}

      {:error, {:already_started, pid}} ->
        IO.puts("#{UI.colors().info}• Client already running#{UI.colors().reset}")
        {:ok, pid}

      error ->
        error
    end
  end

  defp cleanup_processes(state) do
    if state.client_pid && Process.alive?(state.client_pid) do
      Process.exit(state.client_pid, :shutdown)
    end

    if state.transport_pid && Process.alive?(state.transport_pid) do
      Process.exit(state.transport_pid, :shutdown)
    end

    Process.sleep(100)
  end
end