defmodule CentralcloudCore.WazuhMcp do
  @moduledoc """
  MCP (Model Context Protocol) server implementation for Wazuh endpoint security.

  Implements the JSON-RPC 2.0 / MCP streamable-HTTP protocol so that the
  Hermes security-analyst agent can call Wazuh tools directly from the
  centralcloud_staff app — no separate Python MCP sidecar required.

  Supported methods:
    initialize   — handshake, returns capabilities
    tools/list   — enumerates available tools + JSON schemas
    tools/call   — dispatches a named tool call and returns the result

  All tools are read-only wrappers around CentralcloudCore.Security.
  """

  alias CentralcloudCore.Security

  @protocol_version "2024-11-05"
  @server_info %{"name" => "wazuh-mcp", "version" => "1.0.0"}

  @tools [
    %{
      "name" => "list_agents",
      "description" =>
        "List monitored endpoints, optionally scoped to a customer. " <>
          "Pass customer_id (HostBill integer) to scope to cc_client_{id}.",
      "inputSchema" => %{
        "type" => "object",
        "properties" => %{
          "customer_id" => %{"type" => "integer", "description" => "HostBill customer ID"},
          "status" => %{
            "type" => "string",
            "enum" => ["active", "disconnected", "pending", "never_connected"],
            "description" => "Filter by agent status"
          },
          "limit" => %{"type" => "integer", "description" => "Max agents (default 100, max 500)"}
        }
      }
    },
    %{
      "name" => "get_vulnerabilities",
      "description" =>
        "Return CVE vulnerabilities for a specific Wazuh agent, filterable by severity.",
      "inputSchema" => %{
        "type" => "object",
        "required" => ["agent_id"],
        "properties" => %{
          "agent_id" => %{"type" => "string", "description" => "Wazuh agent ID, e.g. \"001\""},
          "severity" => %{
            "type" => "string",
            "enum" => ["Critical", "High", "Medium", "Low"],
            "description" => "Filter by severity"
          },
          "limit" => %{"type" => "integer", "description" => "Max CVEs (default 100)"}
        }
      }
    },
    %{
      "name" => "get_alerts",
      "description" =>
        "Return recent SIEM alerts at or above a rule level, optionally scoped to a customer or agent.",
      "inputSchema" => %{
        "type" => "object",
        "properties" => %{
          "customer_id" => %{"type" => "integer", "description" => "Scope to a customer's agents"},
          "agent_id" => %{"type" => "string", "description" => "Scope to a single agent"},
          "level" => %{
            "type" => "integer",
            "description" => "Minimum rule level 0-15 (default 7)"
          },
          "limit" => %{"type" => "integer", "description" => "Max alerts (default 50)"}
        }
      }
    },
    %{
      "name" => "get_sca",
      "description" =>
        "Return SCA/CIS benchmark compliance results for an agent. " <>
          "Omit policy_id to list all policies; provide it for per-check detail.",
      "inputSchema" => %{
        "type" => "object",
        "required" => ["agent_id"],
        "properties" => %{
          "agent_id" => %{"type" => "string", "description" => "Wazuh agent ID"},
          "policy_id" => %{
            "type" => "string",
            "description" => "SCA policy ID for per-check results"
          }
        }
      }
    },
    %{
      "name" => "agent_summary",
      "description" =>
        "Connectivity health: counts of active / disconnected / never-connected agents.",
      "inputSchema" => %{
        "type" => "object",
        "properties" => %{
          "customer_id" => %{"type" => "integer", "description" => "Scope to a customer. Omit for all."}
        }
      }
    },
    %{
      "name" => "top_critical_cves",
      "description" =>
        "Top N Critical CVEs ranked by affected-host count + CVSS score across all active agents.",
      "inputSchema" => %{
        "type" => "object",
        "properties" => %{
          "customer_id" => %{"type" => "integer", "description" => "Scope to a customer. Omit for all."},
          "top_n" => %{"type" => "integer", "description" => "Number of CVEs to return (default 20)"}
        }
      }
    }
  ]

  # ---------------------------------------------------------------------------
  # Dispatch
  # ---------------------------------------------------------------------------

  @doc "Handle a single JSON-RPC request map. Returns a JSON-RPC result map."
  def handle(%{"method" => "initialize", "id" => id}) do
    result(%{
      "protocolVersion" => @protocol_version,
      "capabilities" => %{"tools" => %{}},
      "serverInfo" => @server_info
    }, id)
  end

  def handle(%{"method" => "tools/list", "id" => id}) do
    result(%{"tools" => @tools}, id)
  end

  def handle(%{"method" => "tools/call", "id" => id, "params" => %{"name" => name} = params}) do
    args = Map.get(params, "arguments", %{})

    case call_tool(name, args) do
      {:ok, content} ->
        result(%{"content" => [%{"type" => "text", "text" => Jason.encode!(content)}]}, id)

      {:error, reason} ->
        error_result(-32_000, format_error(reason), id)
    end
  end

  def handle(%{"method" => method, "id" => id}) do
    error_result(-32_601, "Method not found: #{method}", id)
  end

  def handle(_), do: error_result(-32_600, "Invalid request", nil)

  # ---------------------------------------------------------------------------
  # Tool implementations
  # ---------------------------------------------------------------------------

  defp call_tool("list_agents", args) do
    opts =
      []
      |> maybe_opt(:customer_id, args["customer_id"])
      |> maybe_opt(:status, args["status"])
      |> maybe_opt(:limit, args["limit"])

    Security.list_agents(opts)
  end

  defp call_tool("get_vulnerabilities", %{"agent_id" => agent_id} = args) do
    opts =
      []
      |> maybe_opt(:severity, args["severity"])
      |> maybe_opt(:limit, args["limit"])

    Security.get_vulnerabilities(agent_id, opts)
  end

  defp call_tool("get_alerts", args) do
    opts =
      []
      |> maybe_opt(:customer_id, args["customer_id"])
      |> maybe_opt(:agent_id, args["agent_id"])
      |> maybe_opt(:level, args["level"])
      |> maybe_opt(:limit, args["limit"])

    Security.get_alerts(opts)
  end

  defp call_tool("get_sca", %{"agent_id" => agent_id} = args) do
    opts = maybe_opt([], :policy_id, args["policy_id"])
    Security.get_sca(agent_id, opts)
  end

  defp call_tool("agent_summary", args) do
    opts = maybe_opt([], :customer_id, args["customer_id"])
    Security.agent_summary(opts)
  end

  defp call_tool("top_critical_cves", args) do
    opts =
      []
      |> maybe_opt(:customer_id, args["customer_id"])
      |> maybe_opt(:top_n, args["top_n"])

    Security.top_critical_cves(opts)
  end

  defp call_tool(name, _args), do: {:error, "Unknown tool: #{name}"}

  # ---------------------------------------------------------------------------
  # Helpers
  # ---------------------------------------------------------------------------

  defp result(data, id), do: %{"jsonrpc" => "2.0", "id" => id, "result" => data}

  defp error_result(code, message, id),
    do: %{"jsonrpc" => "2.0", "id" => id, "error" => %{"code" => code, "message" => message}}

  defp maybe_opt(opts, _key, nil), do: opts
  defp maybe_opt(opts, key, value), do: Keyword.put(opts, key, value)

  defp format_error({:http, status}), do: "Wazuh API returned HTTP #{status}"
  defp format_error({:auth_failed, status}), do: "Wazuh authentication failed (#{status})"
  defp format_error({:connect_error, reason}), do: "Connection error: #{inspect(reason)}"
  defp format_error(reason) when is_binary(reason), do: reason
  defp format_error(reason), do: inspect(reason)
end
