require 'eventmachine'


# TODO: Break this out into its own little library, since it's general-purpose.
require 'bindata'
# Implementation of the BitCoin wire protocol, written using bindata.
# Reference: https://en.bitcoin.it/wiki/Protocol_specification
# 
# @author Nick Thomas <nick@lupine.me.uk
module BtcWireProto
  CURRENT_VERSION = 32100
  # Comprehensive list of known networks. The hex values are what you see in
  # MessageHdr#magic and the symbols are their known friendly names.
  NETWORKS = {
    :main      => 0xF9BEB4D9,
    0xF9BEB4D9 => :main,

    :testnet   => 0xFABFB5DA,
    0xFABFB5DA => :testnet
  }
  
  # Comprehensive list of known inventory vector types.
  INV_VEC_TYPES = {
    0 => :error,
    :error => 0,
    
    1 => :msg_tx,
    :msg_tx => 1,
    
    2 => :msg_block,
    :msg_block => 2
  }
  
  ## Components of payloads ##
  
  # Bitmask advertising various capabilities of the node.
  # @author Nick Thomas <nick@lupine.me.uk>  
  class ServicesMask < BinData::Record
    endian :little
    
    bit62 :undefined
    bit1  :node_network
    bit1  :undefined
  end

  # Structure holding an IP address and port in a slightly unusual format.
  # This one is big-endian - everything else is little-endian.
  #
  # @author Nick Thomas <nick@lupine.me.uk>
  class NetAddr < BinData::Record
    endian :big
    services_mask :services
    uint128 :ip # IPv6 address. IPv4 addresses given as IPv6-mapped IPv4
    uint16 :port
  end
  
  # Like a NetAddr but with a timestamp to boot.
  # @author Nick Thomas <nick@lupine.me.uk>
  def TimestampedNetAddr < BinData::Record
    uint32 :timestamp, :endian => :little
    net_addr :net_addr
  end
  
  # Variable-length integer. This is slightly scary.
  # @author Nick Thomas <nick@lupine.me.uk>
  class VarInt < BinData::BasePrimitive
    def value_to_binary_string(val)
      val = val.to_i

      case val
      if val < -0xffffffffffffffff # unrepresentable
        ""
      elsif val < 0  # 64-bit negative integer
        top_32 = (val & 0xffffffff00000000) >> 32
        btm_32 = val & 0x00000000ffffffff
        [0xff, top_32, btm_32].pack("CVV") 
      elsif val <= 0xfc # 8-bit (almost) positive integer
        [val].pack("C")
      elsif val <= 0xffff # 16-bit positive integer
        [0xfd, val].pack("Cv")
      elsif val <= 0xffffffff # 32-bit positive integer
        [0xfe, val].pack("CV")
      else  # We can't represent this, whatever it is
        ""
      end 
    end
    
    def read_and_return_value(io)
      return nil if str.size < 1
      magic = read_uint8(io)
        if magic <= 0xfc # 8-bit (almost) positive integer
          magic
        elsif magic == 0xfd # 16-bit positive integer
          read_uint16(io)
        elsif magic == 0xfe # 32-bit positive integer
          read_uint32(io)
        elsif magic == 0xff # 64-bit negative integer  
          -(read_uint64(io)
        end
      end
    end
    
    def sensible_default
      0
    end
  end

  # Variable-length pascal string with a variable-length int specifying the
  # length. I kid you not.
  # @author Nick Thomas <nick@lupine.me.uk>
  class VarStr < BinData::Primitive
    endian :little

    var_int :len, :value => lambda { data.length }
    string  :data, :read_length => :len
    
    def get    ; self.data     ; end
    def set(v) ; self.data = v ; end
  end
  
  class InventoryVector < BinData::Record
    endian :little

    uint32 :type # For values, see INV_VEC_TYPES
    string :hash, :length => 32
  end
  
  # Simple class wrapping raw SHA256 data. Might have utility methods later.
  # @author Nick Thomas <nick@lupine.me.uk>
  class Sha256 < BinData::Record
    uint256 :data # Raw SHA256 data
  end
  SHA256 = Sha256

  class TransactionIn < BinData::Record
    struct :previous_output do
      sha256 :hash
      uint32 :index
    end
    var_str :signature_script # Script for confirming transaction authorisation
    uint32 :sequence # Version of this record. 
  end
  
  class TranactionOut < BinData::Record
    uint64 :value
    var_str :pk_script # Script containing conditions to claim to transaction
  end


  ## Payloads ##

  # Payload for a version message
  # @author Nick Thomas <nick@lupine.me.uk>
  class Version < BinData::Record
    endian :little

    uint32        :version
    services_mask :services
    uint64        :timestamp
    net_addr      :addr_me
    net_addr      :addr_you,     :only_if => lambda { version >= 106 }
    uint64        :nonce,        :only_if => lambda { version >= 106 }
    var_str       :sub_version,  :only_if => lambda { version >= 106 }
    uint32        :start_height, :only_if => lambda { version >= 209 }
  end
  
  # Payload for an addr message in versions earlier than 31402. These are
  # used to get a list of peers to interact with.
  # @author Nick Thomas <nick@lupine.me.uk>
  class AddrPre31402 < BinData::Record
    endian :little
    
    var_int :count
    array :addrs, :type => :net_addr,
      :read_until => lambda { index == count - 1 }

  end
  
  # Payload for an addr message in versions later than 31402. A timestamp was
  # added to the list of addresses, but otherwise it's the same as AddrPre31402
  # @author Nick Thomas <nick@lupine.me.uk>
  class AddrFrom31402 < BinData::Record
    endian :little
    
    var_int :count
    array :timestamped_addrs, :type => :timestamped_net_addr,
      :read_until => lambda { index == count - 1 }
  end
  
  # Payload for a getdata or inv message. This lets the peer advertise the
  # various objects it has knowledge of.
  # @author Nick Thomas <nick@lupine.me.uk>  
  class Inventory < BinData::Record
    endian :little

    var_int :count
    array :items, :type => :inventory_vector,
      :read_until => lambda { index == count - 1 }
  end
  
  # Payload for a getblocks or getheaders message. Specifies a set of blocks
  # that the sender wants details of.
  # @author Nick thomas <nick@lupine.me.uk>
  class BlockSet < BinData::Record
    endian :little
    
    uint32  :version
    var_int :start_count
    array   :hash_start, :type => :sha256,
            :read_until => lambda { index == start_count - 1 }    
    # Hash of the last desired block, or 0 to get as many as possible (max: 500)
    sha256  :hash_stop, :length => 32
  end
  
  class Transaction < BinData::Record
    endian :little
    
    uint32  :version
    var_int :tx_in_count
    array   :transactions_in, :type => :transaction_in,
            :read_until => lambda { index == tx_in_count - 1 }
    var_int :tx_out_count
    array   :transactions_out, :type => :transaction_out,
            :read_until => lambda { index == tx_in_count - 1 }
    uint32  :lock_time
  end

  ## Top-level message format ##
  
  class MessageHdr < BinData::Record
    endian :little
    uint32 :magic
    string :command, :length => 12
    uint32 :payload_len
    uint32 :checksum, :onlyif => :has_checksum?
    
    protected
    
    # version and verack messages don't have a checksum. The rest do.
    # @return[Boolean] does this message header have a checksum field or not?
    def has_checksum?
      command != "version" && command != "verack"
    end 
  end

  # Everything on the wire is a Message.
  class Message < BinData::Record
  
    # @param[Fixnum,nil] version The protocol version. Setting this affects
    # the layout of various fields.
    def initialize(version = nil)
      @version = version || BtcWireProto::CURRENT_VERSION
    end

    message_hdr :header
    
    choice :payload, :selection => :payload_choice do
      version         "version"
      addr_pre_31402  "addr_pre_31402"
      addr_from_31402 "addr_from_31402"
      inventory       "inv"
      inventory       "getdata"
      block_set       "getblocks"
      block_set       "getheaders"
      transaction     "tx"
    end

    protected

    def payload_choice
      return header.command if %w{
        version inv getdata getblocks getheaders tx
      }.include?(header.command)

      case header.command
      when "verack"  then nil # No payload for a verack message
      when "addr" # two forms, depending on protocol version
        @version < 31402 ? "addr_pre_31402" : "addr_from_31402"
      else
        nil
      end
    end
  end
end


module EventMachine
  module Protocols
    # Implements the TCP protocol that Bitcoin peers speak to each other. This
    # module is mixed into both incoming and outgoing connections.
    #
    # We implement the protocol as a simple(ish!) state machine. When we want
    # something doing, we call state(sym, data) to append that to the 
    # list of things to do. If something is urgent, we can call state! to
    # put it at the beginning of the list.
    #
    # Here is a list of states: 
    #  send_ver, recv_ver, verify_ver
    #  send_verack, recv_verack
    #  wait
    #
    # We must receive a configuration object before we can do much of interest -
    # this is received
    #
    # @author Nick Thomas <nick@lupine.me.uk>
    module BitcoinPeer
      
      protected
      
      # Sets up the variables required to manage the state machine. Should be
      # called before you try to push a state - in post_init, say.
      def init_state!
        @state_m = Mutex.new # Synchronize around @states and @working
        @state_m.synchronize do
          @states = []
          @working = false
        end
      end
      
      # Checks the current configuration object to see if we have a valid config
      # or not.
      # @return[Array[true|false, msg]] Whether the config is valid, and an
      #   optional message specifying why it's invalid, if it is.
      def valid_config?
        [false, "configuration check not implemented yet"]
      end
      
      # Push a state to the end of the state queue.       
      def state(new_state, data = nil)
        @state_m.synchronize { @states.push(new_state, data) }
      end
      
      # Add a state to the start of the state queue.
      def state!(new_state, data = nil)
        @state_m.synchronize { @states.unshift(new_state, data) }
      end
      
      # State machine behaviours now. 
      
      # Send a 'version' message to the peer.
      # Next 
      def send_ver

      end
      
    end
    
    # EventMachine protocol class that handles an *outgoing* connection to
    # another bitcoin peer. Common functionality (p2p!) is held in BitcoinPeer.
    #
    # State machine flow: 
    #   send_ver, recv_verack
    #   recv_ver, verify_ver, send_verack
    # 
    # @author Nick Thomas <nick@lupine.me.uk>
    class BitcoinClient < EM::Connection
    include BitcoinPeer

      # @param[Object] config See the BitcoinPeer#valid_config?
      def initialize(config)
        super
        @config = config
        result, msg = valid_config?
        raise ArgumentError.new("Invalid configuration: #{msg}") unless result
        
        init_state!
      end

      def post_init
        state(:send_ver)
      end
    end
    
    # EventMachine protocol class that handles an *incoming* connection from
    # another bitcoin peer. Common functionality (p2p!) is held in BitcoinPeer
    #
    # State machine flow:
    #   recv_ver, verify_ver, send_verack
    #   send_ver, recv_verack
    #
    # @author Nick Thomas <nick@lupine.me.uk>
    class BitcoinServer < EM::Connection
      include BitcoinPeer

      # @param[Object] config See the BitcoinPeer#valid_config?
      def initialize(config)
        super
        @config = config
        result, msg = valid_config?
        raise ArgumentError.new("Invalid configuration: #{msg}") unless result
        
        init_state!
      end

      def post_init
        state(:recv_ver)
      end
    end
  end
end