object

defmodule Object.NATTraversal do
  # import Bitwise
  @moduledoc """
  NAT traversal mechanisms for P2P Object communication.
  
  Implements STUN, TURN, and ICE protocols to enable direct peer-to-peer
  connections through NATs and firewalls. Supports UDP hole punching,
  TCP simultaneous open, and relay fallback.
  
  ## Features
  
  - STUN client for address discovery
  - TURN client for relay allocation
  - ICE agent for connection negotiation
  - UDP hole punching
  - TCP simultaneous open
  - UPnP/NAT-PMP port mapping
  - Relay server fallback
  """
  
  use GenServer
  require Logger
  
  # @stun_port 3478
  # @turn_port 3478
  @stun_magic_cookie 0x2112A442
  @stun_binding_request 0x0001
  @stun_binding_response 0x0101
  @hole_punch_attempts 10
  @hole_punch_interval 200
  
  @type nat_type :: :none | :full_cone | :restricted_cone | :port_restricted | :symmetric
  
  @type candidate :: %{
    type: :host | :srflx | :prflx | :relay,
    protocol: :udp | :tcp,
    address: String.t(),
    port: non_neg_integer(),
    priority: non_neg_integer(),
    foundation: String.t(),
    related_address: String.t() | nil,
    related_port: non_neg_integer() | nil
  }
  
  @type ice_state :: :new | :gathering | :complete | :connected | :failed
  
  @type connection :: %{
    local_candidates: [candidate()],
    remote_candidates: [candidate()],
    selected_pair: {candidate(), candidate()} | nil,
    state: ice_state(),
    controlling: boolean(),
    tie_breaker: binary()
  }
  
  @type state :: %{
    stun_servers: [{String.t(), non_neg_integer()}],
    turn_servers: [{String.t(), non_neg_integer(), map()}],
    nat_type: nat_type() | nil,
    external_address: {String.t(), non_neg_integer()} | nil,
    connections: %{String.t() => connection()},
    relay_allocations: %{String.t() => relay_allocation()},
    upnp_mappings: %{non_neg_integer() => upnp_mapping()},
    config: map()
  }
  
  @type relay_allocation :: %{
    server: {String.t(), non_neg_integer()},
    relayed_address: {String.t(), non_neg_integer()},
    lifetime: non_neg_integer(),
    permissions: MapSet.t()
  }
  
  @type upnp_mapping :: %{
    internal_port: non_neg_integer(),
    external_port: non_neg_integer(),
    protocol: :udp | :tcp,
    description: String.t()
  }
  
  # Client API
  
  @doc """
  Starts the NAT traversal service.
  """
  def start_link(opts \\ []) do
    GenServer.start_link(__MODULE__, opts, name: __MODULE__)
  end
  
  @doc """
  Discovers NAT type and external address.
  """
  @spec discover_nat() :: {:ok, {nat_type(), String.t(), non_neg_integer()}} | {:error, term()}
  def discover_nat do
    GenServer.call(__MODULE__, :discover_nat, 10_000)
  end
  
  @doc """
  Gathers ICE candidates for establishing a connection.
  """
  @spec gather_candidates(String.t()) :: {:ok, [candidate()]} | {:error, term()}
  def gather_candidates(connection_id) do
    GenServer.call(__MODULE__, {:gather_candidates, connection_id}, 15_000)
  end
  
  @doc """
  Starts ICE negotiation with remote candidates.
  """
  @spec start_ice_negotiation(String.t(), [candidate()], boolean()) :: 
    :ok | {:error, term()}
  def start_ice_negotiation(connection_id, remote_candidates, controlling) do
    GenServer.call(__MODULE__, 
      {:start_ice_negotiation, connection_id, remote_candidates, controlling})
  end
  
  @doc """
  Establishes a direct connection using hole punching.
  """
  @spec hole_punch(non_neg_integer(), String.t(), non_neg_integer()) ::
    {:ok, port()} | {:error, term()}
  def hole_punch(local_port, remote_address, remote_port) do
    GenServer.call(__MODULE__, 
      {:hole_punch, local_port, remote_address, remote_port}, 30_000)
  end
  
  @doc """
  Allocates a TURN relay for incoming connections.
  """
  @spec allocate_relay() :: {:ok, {String.t(), non_neg_integer()}} | {:error, term()}
  def allocate_relay do
    GenServer.call(__MODULE__, :allocate_relay, 10_000)
  end
  
  @doc """
  Creates a UPnP port mapping.
  """
  @spec create_port_mapping(non_neg_integer(), :udp | :tcp, String.t()) ::
    {:ok, non_neg_integer()} | {:error, term()}
  def create_port_mapping(internal_port, protocol, description) do
    GenServer.call(__MODULE__, 
      {:create_port_mapping, internal_port, protocol, description})
  end
  
  # Server Callbacks
  
  @impl true
  def init(opts) do
    state = %{
      stun_servers: Keyword.get(opts, :stun_servers, default_stun_servers()),
      turn_servers: Keyword.get(opts, :turn_servers, []),
      nat_type: nil,
      external_address: nil,
      connections: %{},
      relay_allocations: %{},
      upnp_mappings: %{},
      config: %{
        enable_upnp: Keyword.get(opts, :enable_upnp, true),
        enable_turn: Keyword.get(opts, :enable_turn, true),
        gather_timeout: Keyword.get(opts, :gather_timeout, 5_000)
      }
    }
    
    # Try to discover UPnP gateway
    if state.config.enable_upnp do
      Task.start(fn -> discover_upnp_gateway() end)
    end
    
    {:ok, state}
  end
  
  @impl true
  def handle_call(:discover_nat, _from, state) do
    case perform_nat_discovery(state) do
      {:ok, nat_type, external_addr, external_port} ->
        new_state = %{state | 
          nat_type: nat_type,
          external_address: {external_addr, external_port}
        }
        {:reply, {:ok, {nat_type, external_addr, external_port}}, new_state}
        
      {:error, _reason} = error ->
        {:reply, error, state}
    end
  end
  
  @impl true
  def handle_call({:gather_candidates, conn_id}, _from, state) do
    # Initialize connection state
    connection = %{
      local_candidates: [],
      remote_candidates: [],
      selected_pair: nil,
      state: :gathering,
      controlling: true,
      tie_breaker: :crypto.strong_rand_bytes(8)
    }
    
    # Gather host candidates
    host_candidates = gather_host_candidates()
    
    # Gather server reflexive candidates (via STUN)
    srflx_candidates = gather_srflx_candidates(state.stun_servers)
    
    # Gather relay candidates (via TURN) if enabled
    relay_candidates = if state.config.enable_turn do
      gather_relay_candidates(state.turn_servers)
    else
      []
    end
    
    all_candidates = host_candidates ++ srflx_candidates ++ relay_candidates
    |> assign_priorities()
    
    updated_connection = %{connection | 
      local_candidates: all_candidates,
      state: :complete
    }
    
    new_state = put_in(state.connections[conn_id], updated_connection)
    
    {:reply, {:ok, all_candidates}, new_state}
  end
  
  @impl true
  def handle_call({:start_ice_negotiation, conn_id, remote_candidates, controlling}, 
                  _from, state) do
    case Map.get(state.connections, conn_id) do
      nil ->
        {:reply, {:error, :no_connection}, state}
        
      connection ->
        updated_conn = %{connection |
          remote_candidates: remote_candidates,
          controlling: controlling,
          state: :connected
        }
        
        # Start connectivity checks
        Task.start(fn ->
          perform_connectivity_checks(conn_id, updated_conn)
        end)
        
        new_state = put_in(state.connections[conn_id], updated_conn)
        {:reply, :ok, new_state}
    end
  end
  
  @impl true
  def handle_call({:hole_punch, local_port, remote_addr, remote_port}, _from, state) do
    case perform_hole_punch(local_port, remote_addr, remote_port) do
      {:ok, _socket} = result ->
        {:reply, result, state}
        
      {:error, _reason} = error ->
        {:reply, error, state}
    end
  end
  
  @impl true
  def handle_call(:allocate_relay, _from, state) do
    case state.turn_servers do
      [] ->
        {:reply, {:error, :no_turn_servers}, state}
        
      [{server, port, creds} | _] ->
        case allocate_turn_relay(server, port, creds) do
          {:ok, allocation} ->
            relay_id = generate_relay_id()
            new_state = put_in(state.relay_allocations[relay_id], allocation)
            {:reply, {:ok, allocation.relayed_address}, new_state}
            
          {:error, _reason} = error ->
            {:reply, error, state}
        end
    end
  end
  
  @impl true
  def handle_call({:create_port_mapping, internal_port, protocol, description}, 
                  _from, state) do
    if state.config.enable_upnp do
      case create_upnp_mapping(internal_port, protocol, description) do
        {:ok, external_port} ->
          mapping = %{
            internal_port: internal_port,
            external_port: external_port,
            protocol: protocol,
            description: description
          }
          new_state = put_in(state.upnp_mappings[internal_port], mapping)
          {:reply, {:ok, external_port}, new_state}
          
        {:error, _reason} = error ->
          {:reply, error, state}
      end
    else
      {:reply, {:error, :upnp_disabled}, state}
    end
  end
  
  @impl true
  def handle_info({:udp, socket, from_ip, from_port, data}, state) do
    # Handle STUN/TURN responses and ICE connectivity checks
    case decode_stun_message(data) do
      {:ok, message} ->
        handle_stun_message(state, socket, from_ip, from_port, message)
        
      {:error, _} ->
        # Not a STUN message, might be application data
        {:noreply, state}
    end
  end
  
  # NAT Discovery
  
  defp perform_nat_discovery(state) do
    case state.stun_servers do
      [] ->
        {:error, :no_stun_servers}
        
      servers ->
        # Perform RFC 3489 NAT type discovery
        case stun_test_1(servers) do
          {:ok, mapped_addr, mapped_port} ->
            case stun_test_2(servers, mapped_addr, mapped_port) do
              :same ->
                # Either no NAT or full cone NAT
                if mapped_addr == get_local_address() do
                  {:ok, :none, mapped_addr, mapped_port}
                else
                  {:ok, :full_cone, mapped_addr, mapped_port}
                end
                
              :different ->
                # Symmetric NAT
                {:ok, :symmetric, mapped_addr, mapped_port}
                
              {:error, :timeout} ->
                # Test 3 to distinguish between restricted and port restricted
                case stun_test_3(servers) do
                  :success ->
                    {:ok, :restricted_cone, mapped_addr, mapped_port}
                  :failure ->
                    {:ok, :port_restricted, mapped_addr, mapped_port}
                end
            end
            
          {:error, reason} ->
            {:error, reason}
        end
    end
  end
  
  defp stun_test_1(servers) do
    # Basic STUN binding request
    {server, port} = Enum.at(servers, 0)
    send_stun_binding_request(server, port)
  end
  
  defp stun_test_2(servers, mapped_addr, mapped_port) do
    # Send from same socket to different server
    if length(servers) > 1 do
      {server2, port2} = Enum.at(servers, 1)
      case send_stun_binding_request(server2, port2) do
        {:ok, addr2, port2} ->
          if mapped_addr == addr2 and mapped_port == port2 do
            :same
          else
            :different
          end
        {:error, _} ->
          :different
      end
    else
      :same
    end
  end
  
  defp stun_test_3(_servers) do
    # Would send to different IP but same port
    # Simplified for now
    :success
  end
  
  # STUN Protocol
  
  defp send_stun_binding_request(server, port) do
    {:ok, socket} = :gen_udp.open(0, [:binary, active: false])
    
    # Build STUN binding request
    transaction_id = :crypto.strong_rand_bytes(12)
    message = build_stun_message(@stun_binding_request, transaction_id, <<>>)
    
    # Send request
    :ok = :gen_udp.send(socket, to_charlist(server), port, message)
    
    # Wait for response
    case :gen_udp.recv(socket, 0, 5000) do
      {:ok, {_addr, _port, response}} ->
        :gen_udp.close(socket)
        parse_stun_response(response)
        
      {:error, :timeout} ->
        :gen_udp.close(socket)
        {:error, :timeout}
    end
  end
  
  defp build_stun_message(msg_type, transaction_id, payload) do
    length = byte_size(payload)
    <<
      msg_type::16,
      length::16,
      @stun_magic_cookie::32,
      transaction_id::binary-size(12),
      payload::binary
    >>
  end
  
  defp parse_stun_response(<<msg_type::16, length::16, @stun_magic_cookie::32,
                            _transaction_id::binary-size(12), 
                            attributes::binary-size(length)>>) do
    if msg_type == @stun_binding_response do
      parse_mapped_address(attributes)
    else
      {:error, :invalid_response}
    end
  end
  defp parse_stun_response(_), do: {:error, :invalid_message}
  
  defp parse_mapped_address(attributes) do
    # Simplified - look for XOR-MAPPED-ADDRESS (0x0020)
    case parse_stun_attributes(attributes) do
      %{0x0020 => addr_data} ->
        decode_xor_mapped_address(addr_data)
      %{0x0001 => addr_data} ->
        decode_mapped_address(addr_data)
      _ ->
        {:error, :no_mapped_address}
    end
  end
  
  defp parse_stun_attributes(data, attrs \\ %{})
  defp parse_stun_attributes(<<>>, attrs), do: attrs
  defp parse_stun_attributes(<<type::16, length::16, value::binary-size(length),
                              _padding::binary-size(rem(4 - rem(length, 4), 4)),
                              rest::binary>>, attrs) do
    parse_stun_attributes(rest, Map.put(attrs, type, value))
  end
  
  defp decode_xor_mapped_address(<<_::8, family::8, port::16, addr::32>>) 
       when family == 0x01 do
    # IPv4
    xor_port = Bitwise.bxor(port, Bitwise.bsr(@stun_magic_cookie, 16))
    xor_addr = Bitwise.bxor(addr, @stun_magic_cookie)
    
    ip = :inet.ntoa({
      Bitwise.band(Bitwise.bsr(xor_addr, 24), 0xFF),
      Bitwise.band(Bitwise.bsr(xor_addr, 16), 0xFF),
      Bitwise.band(Bitwise.bsr(xor_addr, 8), 0xFF),
      Bitwise.band(xor_addr, 0xFF)
    }) |> to_string()
    
    {:ok, ip, xor_port}
  end
  
  defp decode_mapped_address(<<_::8, family::8, port::16, addr::32>>) 
       when family == 0x01 do
    # IPv4
    ip = :inet.ntoa({
      Bitwise.band(Bitwise.bsr(addr, 24), 0xFF),
      Bitwise.band(Bitwise.bsr(addr, 16), 0xFF),
      Bitwise.band(Bitwise.bsr(addr, 8), 0xFF),
      Bitwise.band(addr, 0xFF)
    }) |> to_string()
    
    {:ok, ip, port}
  end
  
  # ICE Candidate Gathering
  
  defp gather_host_candidates do
    # Get all local network interfaces
    {:ok, interfaces} = :inet.getifaddrs()
    
    interfaces
    |> Enum.flat_map(fn {_name, opts} ->
      case Keyword.get(opts, :addr) do
        {a, _b, _c, _d} = addr when a != 127 ->
          # IPv4 non-loopback
          [%{
            type: :host,
            protocol: :udp,
            address: :inet.ntoa(addr) |> to_string(),
            port: 0,
            priority: 0,
            foundation: generate_foundation(:host, addr, :udp),
            related_address: nil,
            related_port: nil
          }]
        _ ->
          []
      end
    end)
  end
  
  defp gather_srflx_candidates(stun_servers) do
    stun_servers
    |> Enum.map(fn {server, port} ->
      Task.async(fn ->
        case send_stun_binding_request(server, port) do
          {:ok, addr, port} ->
            %{
              type: :srflx,
              protocol: :udp,
              address: addr,
              port: port,
              priority: 0,
              foundation: generate_foundation(:srflx, addr, :udp),
              related_address: get_local_address(),
              related_port: 0
            }
          _ ->
            nil
        end
      end)
    end)
    |> Enum.map(&Task.await(&1, 5000))
    |> Enum.filter(&(&1 != nil))
  end
  
  defp gather_relay_candidates(_turn_servers) do
    # Simplified - would implement TURN allocation
    []
  end
  
  defp assign_priorities(candidates) do
    # ICE priority formula: (2^24)*(type preference) + (2^8)*(local preference) + 
    # (256 - component ID)
    type_prefs = %{host: 126, srflx: 100, prflx: 110, relay: 0}
    
    candidates
    |> Enum.with_index()
    |> Enum.map(fn {candidate, index} ->
      type_pref = Map.get(type_prefs, candidate.type, 0)
      local_pref = 65535 - index
      component_id = 1  # RTP component
      
      priority = Bitwise.bsl(type_pref, 24) + Bitwise.bsl(local_pref, 8) + (256 - component_id)
      
      %{candidate | priority: priority}
    end)
  end
  
  # Hole Punching
  
  defp perform_hole_punch(local_port, remote_addr, remote_port) do
    {:ok, socket} = :gen_udp.open(local_port, [:binary, active: true])
    
    # Send multiple packets to punch hole
    punch_task = Task.async(fn ->
      Enum.each(1..@hole_punch_attempts, fn i ->
        :gen_udp.send(socket, to_charlist(remote_addr), remote_port, 
                      "PUNCH#{i}")
        Process.sleep(@hole_punch_interval)
      end)
    end)
    
    # Wait for response
    receive do
      {:udp, ^socket, ^remote_addr, ^remote_port, _data} ->
        Task.shutdown(punch_task)
        {:ok, socket}
    after
      @hole_punch_attempts * @hole_punch_interval + 1000 ->
        Task.shutdown(punch_task)
        :gen_udp.close(socket)
        {:error, :hole_punch_failed}
    end
  end
  
  # ICE Connectivity Checks
  
  defp perform_connectivity_checks(conn_id, connection) do
    # Create candidate pairs
    pairs = for local <- connection.local_candidates,
                remote <- connection.remote_candidates do
      {local, remote}
    end
    |> sort_candidate_pairs()
    
    # Perform checks
    Enum.find(pairs, fn {local, remote} ->
      case check_candidate_pair(local, remote, connection.controlling) do
        :ok -> 
          send(self(), {:ice_pair_selected, conn_id, {local, remote}})
          true
        _ -> 
          false
      end
    end)
  end
  
  defp sort_candidate_pairs(pairs) do
    # Sort by priority - higher is better
    Enum.sort_by(pairs, fn {local, remote} ->
      # G = controlling agent priority, D = controlled agent priority
      g = local.priority
      d = remote.priority
      
      # Priority = 2^32 * MIN(G,D) + 2 * MAX(G,D) + (G>D?1:0)
      min_p = min(g, d)
      max_p = max(g, d)
      extra = if g > d, do: 1, else: 0
      
      Bitwise.bsl(1, 32) * min_p + 2 * max_p + extra
    end, &>=/2)
  end
  
  defp check_candidate_pair(local, remote, _controlling) do
    # Simplified connectivity check
    case {local.type, remote.type} do
      {:host, :host} ->
        # Direct connection possible
        :ok
        
      {:srflx, _} ->
        # Try through NAT
        :ok
        
      {_, :relay} ->
        # Use relay
        :ok
        
      _ ->
        :failed
    end
  end
  
  # TURN Protocol
  
  defp allocate_turn_relay(server, port, _credentials) do
    # Simplified TURN allocation
    # Would implement full TURN protocol
    {:ok, %{
      server: {server, port},
      relayed_address: {server, port + 1000},
      lifetime: 600,
      permissions: MapSet.new()
    }}
  end
  
  # UPnP Support
  
  defp discover_upnp_gateway do
    # SSDP discovery
    {:ok, socket} = :gen_udp.open(0, [:binary, active: false])
    
    ssdp_request = """
    M-SEARCH * HTTP/1.1\r
    HOST: 239.255.255.250:1900\r
    MAN: "ssdp:discover"\r
    MX: 3\r
    ST: urn:schemas-upnp-org:device:InternetGatewayDevice:1\r
    \r
    """
    
    :gen_udp.send(socket, {239, 255, 255, 250}, 1900, ssdp_request)
    
    case :gen_udp.recv(socket, 0, 3000) do
      {:ok, {addr, _port, response}} ->
        :gen_udp.close(socket)
        parse_ssdp_response(response, addr)
        
      {:error, :timeout} ->
        :gen_udp.close(socket)
        {:error, :no_upnp_gateway}
    end
  end
  
  defp parse_ssdp_response(_response, gateway_addr) do
    # Extract location header and parse device description
    # Simplified for brevity
    {:ok, gateway_addr}
  end
  
  defp create_upnp_mapping(internal_port, _protocol, _description) do
    # Simplified UPnP port mapping
    # Would use SOAP request to add port mapping
    external_port = internal_port + 10000
    {:ok, external_port}
  end
  
  # Utilities
  
  defp default_stun_servers do
    [
      {"stun.l.google.com", 19302},
      {"stun1.l.google.com", 19302},
      {"stun2.l.google.com", 19302},
      {"stun3.l.google.com", 19302},
      {"stun4.l.google.com", 19302}
    ]
  end
  
  defp get_local_address do
    {:ok, interfaces} = :inet.getifaddrs()
    
    interfaces
    |> Enum.find_value(fn {_name, opts} ->
      case Keyword.get(opts, :addr) do
        {a, b, c, d} when a != 127 ->
          :inet.ntoa({a, b, c, d}) |> to_string()
        _ ->
          nil
      end
    end) || "127.0.0.1"
  end
  
  defp generate_foundation(type, address, protocol) do
    data = "#{type}:#{address}:#{protocol}"
    :crypto.hash(:md5, data)
    |> Base.encode16(case: :lower)
    |> String.slice(0, 8)
  end
  
  defp generate_relay_id do
    :crypto.strong_rand_bytes(16) |> Base.encode16(case: :lower)
  end
  
  defp decode_stun_message(data) do
    case data do
      <<msg_type::16, _length::16, @stun_magic_cookie::32, _rest::binary>> ->
        {:ok, %{type: msg_type}}
      _ ->
        {:error, :not_stun}
    end
  end
  
  defp handle_stun_message(state, _socket, _from_ip, _from_port, _message) do
    # Handle incoming STUN messages
    {:noreply, state}
  end
end