gust_web

defmodule GustWeb.MCP.Tools.Call do
  @moduledoc false

  alias Gust.DAG.{Adapter, Definition, Loader, TaskWaiter, Terminator}
  alias Gust.DAG.Run.Trigger
  alias Gust.Flows
  alias GustWeb.MCP.{Content, Tool, Tools}

  def handle(%Tool{name: :list_dags}, _args) do
    {false,
     for {id, {:ok, %Definition{} = dag_def}} <-
           Loader.get_definitions() do
       dag_def_text(id, dag_def) |> content()
     end}
  end

  def handle(%Tool{name: :list_secrets}, _args) do
    {false,
     for %Flows.Secret{name: name, id: id, value_type: type} <- Flows.list_secrets() do
       content("Name: #{name}; ID: #{id}; Type: #{type}")
     end}
  end

  def handle(%Tool{name: :query_dag_run}, %{
        "dag_name" => dag_name,
        "limit" => limit,
        "offset" => offset
      }) do
    dag = Flows.get_dag_by_name(dag_name)

    if dag do
      dag = Flows.get_dag_by_name_with_runs!(dag_name, limit: limit, offset: offset)

      {false,
       for %Flows.Run{id: id, inserted_at: inserted_at, updated_at: updated, status: status} <-
             dag.runs do
         content(
           "ID: #{id}; Inserted at: #{inserted_at}; Updated at: #{updated}; Status: #{status}"
         )
       end}
    else
      dag_not_found(dag_name)
    end
  end

  def handle(%Tool{name: :get_dag_def}, %{"dag_name" => dag_name}) do
    dag = Flows.get_dag_by_name(dag_name)

    if dag do
      dag.id |> dag_definition_reply()
    else
      dag_not_found(dag_name)
    end
  end

  def handle(%Tool{name: :get_dag_def}, %{"dag_id" => dag_id}) do
    dag_id |> dag_definition_reply()
  end

  def handle(%Tool{name: :toggle_enabled_dag}, %{"dag_id" => dag_id}) do
    {:ok, dag} = Flows.get_dag!(dag_id) |> Flows.toggle_enabled()

    if dag.enabled do
      Trigger.dispatch_all_runs(dag.id)
    end

    dag_id |> dag_definition_reply()
  end

  def handle(%Tool{name: :get_tasks_on_run}, %{"run_id" => run_id}) do
    run = Flows.get_run_with_tasks!(run_id)

    {false,
     for %Flows.Task{id: id, name: name, status: status, error: e, result: res} <-
           run.tasks do
       content(
         "ID: #{id}; Name: #{name}, Status: #{status}; Error: #{inspect(e)}, Result: #{inspect(res)}"
       )
     end}
  end

  def handle(%Tool{name: :get_logs_on_task}, %{"task_id" => task_id}) do
    logs = Flows.get_logs(task_id)

    {false,
     for %Flows.Log{id: id, level: lvl, inserted_at: inserted_at, content: content} <-
           logs do
       content(
         "ID: #{id}; level: #{lvl}, inserted_at: #{inspect(inserted_at)}; Content: #{content}"
       )
     end}
  end

  def handle(%Tool{name: :restart_run}, %{"run_id" => run_id}) do
    run = Flows.get_run!(run_id) |> Trigger.reset_run()

    {false, [content("Run: #{run.id} was restarted")]}
  end

  def handle(%Tool{name: :restart_task}, %{"task_id" => task_id}) do
    task = Flows.get_task!(task_id)
    {:ok, dag_def} = get_def_by_task(task)
    tasks_graph = dag_def.tasks

    Trigger.reset_task(tasks_graph, task)

    {false, [content("Task: #{task.name} was restarted")]}
  end

  def handle(%Tool{name: :cancel_task}, %{"task_id" => task_id}) do
    task = Flows.get_task!(task_id)
    {:ok, dag_def} = get_def_by_task(task)

    text =
      case task.status do
        :running ->
          runtime = Adapter.impl!(dag_def.adapter, :runtime)
          Terminator.kill_task(task, :cancelled, runtime)
          "Task: #{task.name} was cancelled"

        :retrying ->
          Terminator.cancel_timer(task, :cancelled)
          "Task: #{task.name} retrying cancelled"

        :waiting ->
          Terminator.cancel_waiting(task)
          "Task: #{task.name} waiting cancelled"

        status ->
          "Task: #{task.name} cannot be cancelled from status #{inspect(status)}. Only :running, :retrying, and :waiting tasks can be cancelled."
      end

    {false, [content(text)]}
  end

  def handle(%Tool{name: :resume_task}, %{"waiting_for" => waiting_for} = args) do
    opts =
      []
      |> maybe_put_resume_run_id(args)
      |> Keyword.put(:payload, Map.get(args, "payload", %{}))

    {:ok, tasks} = TaskWaiter.resume(waiting_for, opts)

    task_ids = Enum.map(tasks, & &1.id)
    run_ids = tasks |> Enum.map(& &1.run_id) |> Enum.uniq()

    {false,
     [
       content(
         "Resumed #{length(tasks)} task(s) waiting for #{inspect(waiting_for)}. Task IDs: #{inspect(task_ids)}; Run IDs: #{inspect(run_ids)}"
       )
     ]}
  end

  def handle(%Tool{name: :trigger_dag_run}, %{"dag_id" => dag_id} = args) do
    trigger_dag_run_reply(dag_id, Map.get(args, "params", %{}))
  end

  def handle(%Tool{name: :trigger_dag_run}, %{"dag_name" => dag_name} = args) do
    dag = Flows.get_dag_by_name(dag_name)

    if dag do
      trigger_dag_run_reply(dag.id, Map.get(args, "params", %{}))
    else
      dag_not_found(dag_name)
    end
  end

  def handle(%Tool{name: name} = tool, _args) do
    tool = Tools.find(to_string(name)) || tool

    {true, [content(invalid_properties_text(tool))]}
  end

  defp get_def_by_task(task) do
    run = Flows.get_run!(task.run_id)
    Loader.get_definition(run.dag_id)
  end

  defp dag_definition_reply(id) do
    case Loader.get_definition(id) do
      {:ok, dag_def} ->
        {false, [dag_def_text(id, dag_def) |> content()]}

      {:error, error} ->
        {false, [content("DAG with ID #{id} cannot be parsed, error: #{inspect(error)}")]}
    end
  end

  defp dag_def_text(
         dag_id,
         %Definition{
           name: name,
           adapter: adapter,
           options: opts,
           tasks: tasks,
           stages: stages,
           messages: warnings,
           error: e,
           mod: mod,
           file_path: fp
         }
       ) do
    dag = Flows.get_dag!(dag_id)

    """
    Name: #{name}
    ID: #{dag_id}
    Enabled: #{dag.enabled}
    File Path: #{fp}
    Options: #{inspect(opts)}
    Stages: #{inspect(stages)}
    Module: #{mod}
    Adapter: #{adapter}
    Tasks: #{inspect(tasks)}
    Error: #{inspect(e)}
    Warnings: #{inspect(warnings)}
    """
  end

  defp trigger_dag_run_reply(dag_id, params) do
    {:ok, run} = Flows.create_run(%{dag_id: dag_id, params: params})

    run = Flows.get_run_with_tasks!(run.id) |> Trigger.dispatch_run()

    {false, [content("Run #{run.id} triggered")]}
  end

  defp maybe_put_resume_run_id(opts, %{"run_id" => run_id}),
    do: Keyword.put(opts, :run_id, run_id)

  defp maybe_put_resume_run_id(opts, _args), do: opts

  defp invalid_properties_text(%Tool{name: name, props: []}) do
    "Tool #{name} supports no properties."
  end

  defp invalid_properties_text(%Tool{name: name, props: props}) do
    properties =
      Enum.map_join(props, ", ", fn {prop_name, _required, %{"description" => description}} ->
        "#{prop_name}: #{description}"
      end)

    "Tool #{name} supports the following properties: #{properties}"
  end

  defp content(txt), do: Content.new(txt)

  defp dag_not_found(name) do
    {true,
     [
       content("DAG with name #{name} does not exist. Use list_dags to find available DAG names")
     ]}
  end
end