using System;
using System.Collections.Generic;
using System.Net;
using System.Net.Sockets;
using System.Net.NetworkInformation;

abstract class Listener
{
  public Socket Socket;
  public virtual  void Close(){ if ( Socket != null ) Socket.Close(); }
  public abstract void SendResponse( ExternalRequest E, Header H );

  protected static readonly Set<TcpListener> TcpListeners = new Set<TcpListener>();
  private static List<Listener> Listeners = new List<Listener>();

  public static void Start()
  {
    // For UDP,  RFC2181 section 4 requires that the response is sent using same IP 
    // address on which request was received.

    NetworkInterface[] Nics = NetworkInterface.GetAllNetworkInterfaces();
    foreach ( NetworkInterface Nic in Nics )
    {
      bool RxOnly = false;  try { RxOnly = Nic.IsReceiveOnly; } catch (Exception e){}
      if ( !RxOnly )
      {
        IPInterfaceProperties Props =  Nic.GetIPProperties();
        UnicastIPAddressInformationCollection U = Props.UnicastAddresses;
        foreach ( UnicastIPAddressInformation UIA in U )
        {      
          Listeners.Add( new UdpListener( ServerPort.UDP, UIA.Address ) );
        }    
      }
    }   

    for ( int i=0; i<2; i+=1 )
    {
      bool IPv6 = i==1;
      Listeners.Add( new TcpListener( ServerPort.TCP, IPv6 ) );
      Listeners.Add( new QrpListener( ServerPort.QRP, IPv6 ) );
    }
  }

  public static void CheckTimeout( long Now )
  {
    if ( TcpListeners.Count > 0 )
    lock( TcpListeners )
    {
      foreach ( TcpListener tc in TcpListeners )
        if ( Now > tc.TimeoutEnd )
          tc.Close();    
    }
  }

  public static void Stop()
  {
    foreach ( Listener x in Listeners ) x.Close();
    foreach ( Listener x in TcpListeners ) x.Close();  
  }
}

/* The Header class is used to record transport information 
   about an incoming DNS request that is needed to send the response.
*/

abstract class Header
{
  public Listener RC;
  public long Version; // Basis of QRP Cookie

  public abstract EndPoint EP { get; }
  public abstract void Read( DnsRx Rx );
  public abstract int Write( DnsTx Tx );
  public abstract TP TP{ get; }  

  public virtual bool HasPayload{ get { return true; } }

  public void SendResponse( ExternalRequest E ){ RC.SendResponse( E, this ); }

  #if (Trace)
  public override String ToString(){ return base.ToString() + " " + EP.ToString();  }
  #endif
}

abstract class Header16 : Header
{
  protected UInt16 ID;

  public override void Read( DnsRx Rx )
  {
    ID = Rx.Read16();
  }

  public override int Write( DnsTx Tx )
  {
    Tx.Put16( ID );
    return 0;
  }
}

////////////////////////////////////////////////////////////// TCP

class TcpHeader : Header16
{
  public override TP TP{ get{ return TP.TCP; } } 
  public override EndPoint EP { get{ return RC.Socket.RemoteEndPoint; } }
  public override int Write( DnsTx Tx )
  {
    Tx.Put16(0 );   // Reserve space for length
    Tx.Put16( ID );
    return 0;
  }
}

sealed class TcpListener : Listener
{  
  private bool connected;  
  private int PacketLength; 
  private DnsRx Rx;
  public long TimeoutEnd;

  public override void SendResponse( ExternalRequest E, Header H )
  {
    DnsTx Tx = new DnsTx( 0x10001 );
    Tx.PutResponse( E );
    byte [] Packet = Tx.TrimBuffer();    

    TimeoutEnd = System.DateTime.Now.Ticks + 10000L * Config.TcpServerTimeout;
    int Length = Packet.Length - 2;  
    Packet[0] = (byte) ( Length / 256 );
    Packet[1] = (byte) ( Length % 256 );
    Socket.BeginSend( Packet, 0, Packet.Length, 0, new AsyncCallback(SendCallback), this );
  }

  public static void SendCallback( IAsyncResult ar )
  {
    TcpListener rc = (TcpListener)ar.AsyncState;
    try
    {
      int N = rc.Socket.EndSend( ar ); // ToDo : Check N is as expected?
    }
    catch ( Exception e )
    {
      #if (Trace)
        Cache.Log( "Tcp Send Exception" + e );
      #endif
    }
  }

  private void Listen()
  {
    Socket.Listen(1000); // Queue size : not sure what is appropriate (ToDo)
    Socket.BeginAccept( new AsyncCallback(AcceptCallback), this );
  }

  public static void AcceptCallback( IAsyncResult ar )
  {  
    if ( Cache.Stopping ) return;  
    TcpListener rc = (TcpListener)ar.AsyncState;
    try
    {      
      Socket s = rc.Socket.EndAccept( ar );
      new TcpListener( s );  
    }
    catch ( Exception e )
    {
      #if (Trace)
        Cache.Log( "Tcp End Accept Exception " + e.ToString() );
      #endif
    }
    if ( rc.Socket != null ) // Not shutting down
      rc.Socket.BeginAccept( new AsyncCallback(AcceptCallback), rc );     
  }

  public TcpListener( Socket s )
  {
    Socket = s;
    Rx = new DnsRx(2);
    TimeoutEnd = System.DateTime.Now.Ticks + 10000L * Config.TcpServerTimeout;
    lock(TcpListeners)
    {
      TcpListener x = TcpListeners[ this ];
    }
    Receive();
  }

  private void Receive()
  {    
    int Length = PacketLength == 0 ? 2 : Rx.Buffer.Length - Rx.N;
    Socket.BeginReceive( Rx.Buffer, Rx.N, Length,  0/* Socket flags */, new AsyncCallback(ReceiveCallback), this );      
  }

  private static void ReceiveCallback( IAsyncResult ar )
  {    
    TcpListener rc = (TcpListener)ar.AsyncState;
    lock(TcpListeners)
    {
      rc.Respond( ar );
    }
  }

  private void Respond( IAsyncResult ar )
  {
    if ( Socket == null ) return;
    try
    {
      Rx.N += Socket.EndReceive( ar );

      #if (Trace)
        // Cache.Log( "TCP Rx.N=" + Rx.N + " PacketLength=" + PacketLength );
      #endif

      if ( Rx.N == 0 )
      {
        // Do nothing
      }
      else if ( PacketLength == 0 )
      {      
        Rx.Ix = 0;
        PacketLength = Rx.Read16(); 
        // ToDo : should check PacketLength not excessive ( DoS attack )
        Rx.Buffer = new byte[ PacketLength ];
        Rx.N = 0;
        Receive();
      }
      else if ( Rx.N == PacketLength )
      {    
        TcpHeader H = new TcpHeader();
        H.RC = this;
        Rx.H = H;
        Cache.Respond( Rx );
        PacketLength = 0;
        Rx = new DnsRx(2);
        Receive();
      }
      else
      {
        Receive();
      }   
    }
    catch ( Exception e )
    {
      #if (Trace)
        Cache.Log( "Tcp Respond exception" + e );
      #endif
      Close();
      return;
    }
  }

  public override void Close()
  {     
    lock(TcpListeners)
    {
      if ( Socket != null )
      {
        Socket.Close();
        Socket = null;
        if ( TimeoutEnd != 0 && !Cache.Stopping )
          TcpListeners.Remove(this);
      }
    }
  }

  public TcpListener( int Port, bool IPv6 )
  {
    try
    {
      Socket = new Socket
      ( 
        IPv6 ? AddressFamily.InterNetworkV6 : AddressFamily.InterNetwork,
        SocketType.Stream,
        ProtocolType.Tcp
      );
      EndPoint Local = new IPEndPoint( IPv6 ? IPAddress.IPv6Any : IPAddress.Any, Port);
      Socket.Bind(Local);
      Listen();
    }
    catch ( Exception e )
    {
      #if (Trace)
        Cache.Log( IPv6 ? "IPv6 not available" : "IPv4 not available" );
      #endif
      Cache.Any_Address = IPv6 ? QT.A : QT.AAAA;
    }    
  }

}

////////////////////////////////////////////////////////////// UDP

class UdpHeader : Header16
{
  public EndPoint REP;
  public override TP TP{ get{ return TP.UDP; } } 
  public override EndPoint EP { get{ return REP; } }
  public UdpHeader( bool IPv6 )
  {
    REP = new IPEndPoint( IPv6 ? IPAddress.IPv6Any : IPAddress.Any, 0 );
  }
}


class UdpListener : Listener
{
  protected bool IPv6;

  protected virtual UdpHeader GetHeader()
  {
    return new UdpHeader( IPv6 );
  }

  public override void SendResponse( ExternalRequest E, Header _H )
  {  
    int UdpLimit = E.UdpLimit;
    if ( UdpLimit < 512 ) UdpLimit = 512;
    DnsTx Tx = new DnsTx( UdpLimit );
    Tx.PutResponse( E );
    byte [] Packet = Tx.TrimBuffer();
    UdpHeader H = (UdpHeader) _H;
    Socket.SendTo( Packet, H.REP );
  }

  protected void Listen()
  {
    while (true) 
    {
      DnsRx Rx = new DnsRx(1024);
      try
      {
        try
        {
          UdpHeader H = GetHeader();
          H.RC = this;
          Rx.H = H;
          Rx.N = Socket.ReceiveFrom ( Rx.Buffer, ref H.REP );         
        }
        catch ( SocketException se )
        {
          if ( Cache.Stopping ) break;  
          #if (Trace)
            Cache.Log( "Udp Rx Socket exception " + se );
          #endif
        }
        Cache.Respond( Rx );
      }
      catch ( Exception e )
      {
        #if (Trace)
          Cache.Log( "Udp Listener Exception " + e );
        #endif
      }
    }    
  }

  public UdpListener( int Port, IPAddress Address )
  {    
    this.IPv6 = ( Address.AddressFamily == AddressFamily.InterNetworkV6 );
    try
    {
      Socket = new Socket
      ( 
        IPv6 ? AddressFamily.InterNetworkV6 : AddressFamily.InterNetwork,
        SocketType.Dgram,
        ProtocolType.Udp
      );
      EndPoint Local = new IPEndPoint( Address, Port);
      Socket.Bind(Local);    
      new System.Threading.Thread( new System.Threading.ThreadStart(Listen) ).Start(); 
    }
    catch ( Exception e )
    {
      #if (Trace)
        if ( Cache.Any_Address == QT.ANY_ADDRESS )
          Cache.Log( "Remote Client Udp failed to start IPv6=" + IPv6 );
      #endif
    }
  }
}

////////////////////////////////////////////////////////////// QRP

sealed class QrpHeader : UdpHeader
{
  private bool IPv6;

  public QRM Opcode;
  public QRS Status;
  public UInt64 ID;
  public UInt64 Cookie;
  public byte Count;
  public int PageNum;
  public int PageSize;
  public SecurityArg Sa;

  public bool Encrypt { get{ return Sa != null; } }
  public override TP TP{ get{ return TP.QRP; } } 
  public override bool HasPayload { get  { return Opcode == QRM.Single || Opcode == QRM.Multi; } }

  public override void Read( DnsRx Rx )
  {
    Opcode = (QRM)Rx.Read16();
    if ( Opcode == QRM.Encrypt )
    {
      Sa = Crypt.DecodeQuery( Rx );
      Opcode = (QRM)Rx.Read16();      
    }
    else
    {
      ID = Rx.Read64();
    }
    if ( Opcode == QRM.Setup )
    {
      // No defined params yet
    }
    else if ( Opcode == QRM.Single || Opcode == QRM.Multi )
    {
      byte [] ServerToken = Rx.ReadBytes(4);
      byte [] Expected  = Random.GetServerToken( REP );

      int IPoverhead = IPv6 ? 40 : 20; // ToDo : what if IPv6 is tunnelled ?
      // UDP overhead = 8
      // Multi-page overhead = 20 ( Opcode(2) + Total(4)  + Cookie(8) + Page(4) + Pagesize(2)
      // Sub-Total = 28
      // Encrypt overhead = 42 ( Opcode(2), Nonce(24), MAC(16) )
      // Non-encrypt overhead = 8 ( ID(8) )
      // Hence
      int Overhead = IPoverhead + ( Encrypt ? 70 : 36 );

      if ( Opcode == QRM.Single )
      {
        Count = Rx.Read8();
        byte Reserved = Rx.Read8();
        int MTU = Rx.Read16();
        PageSize = MTU - Overhead;
      }
      else // Multi
      {
        Cookie = Rx.Read64();
        Count = Rx.Read8();
        PageNum = Rx.Read24();
        PageSize = Rx.Read16();       
      }
      if ( !Data.BytesEqual(ServerToken,Expected) ) { Opcode = QRM.Setup; Status = QRS.ServerToken; }
      if ( PageSize > 1500 - Overhead ) PageSize = 1500 - Overhead; 
      if ( PageSize < 510 ) PageSize = 510; // ToDo : could use larger number for IPv6 ?
      if ( Count > 8 ) Count = 8; // Maximum server will send at one go
    }   
    else Opcode = QRM.Setup; // (Error)
  }

  public override int Write( DnsTx Tx )
  {
    Tx.CanOverflow = false;
    if ( Opcode == QRM.Single )
    {
      Tx.Put16( (int)Opcode );
      if ( !Encrypt ) Tx.Put64( ID );
      return 0;
    }
    else // Multi
    {
      return PageNum * PageSize;
    }
  }  

  public QrpHeader( bool IPv6 ) : base( IPv6 )
  {
    this.IPv6 = IPv6;
  }
}

sealed class QrpListener : UdpListener
{ 
  private static UInt64 CookieBase = Random.Number64();

  protected override UdpHeader GetHeader()
  {
    return new QrpHeader( IPv6 );
  }

  public override void SendResponse( ExternalRequest E, Header _H )
  {
    QrpHeader H = (QrpHeader) _H;

    if ( H.Opcode == QRM.Setup )
    {
      SendSetupResponse(H);
      return;   
    }

    int PID = H.Encrypt ? 0 : 8;
    int HeaderSize =  H.Opcode == QRM.Single ? 2 + PID : 0;
    DnsTx Tx0 = new DnsTx( HeaderSize + H.Count * H.PageSize );     

    Tx0.PutResponse( E );
    byte [] Packet = Tx0.Buffer;
    int PacketLength = Tx0.Ix;
    int Total = H.PageNum*H.PageSize + PacketLength - HeaderSize;

    if ( H.Opcode == QRM.Single && Total <= H.PageSize + 20 - 2 ) // 20,2 is Multi/Single page overhead resp.
    {
      EncryptAndSendResponse( Tx0.TrimBuffer(), H );    
      return;
    }      
   
    UInt64 Cookie = CookieBase + (UInt64) H.Version;
    if ( H.Opcode == QRM.Multi && Cookie != H.Cookie ) // Cookie Error
    {
      H.Status = QRS.Cookie;
      SendSetupResponse(H);
      return;   
    }

    // Multi-Page response
   
    H.Opcode = QRM.Multi;
    int N = ( 1 + (Total-1) / H.PageSize ) - H.PageNum;
    DnsTx Tx = new DnsTx( 20 + PID + H.PageSize ); // 20 is Multi-page overhead ( 2 + 4 + 8 + 4 + 2 )
    if ( N > H.Count ) N = H.Count;
    for ( int i = 0; i < N; i += 1)
    {      
      Tx.Ix = 0;
      Tx.Put16( (int)H.Opcode );
      if ( !H.Encrypt ) Tx.Put64( H.ID );
      int Offset = HeaderSize + i*H.PageSize;
      int Amount = PacketLength - Offset;
      if ( Amount > H.PageSize ) Amount = H.PageSize;
      Tx.PutSlice( Packet, Offset, Amount );
      Tx.Put32( Total );
      Tx.Put64( Cookie ); // Cookie
      Tx.Put8( N );  // Number of Pages Sent
      Tx.Put24( H.PageNum + i );
      Tx.Put16( H.PageSize );
      EncryptAndSendResponse( Tx.TrimBuffer(), H );
    }

  }

  private void SendSetupResponse( QrpHeader H )
  {
    H.Opcode = QRM.Setup;
    int PID = H.Encrypt ? 0 : 8;
    DnsTx Tx = new DnsTx( 7 + PID ); // Opcode(2) + ServerToken(4) + Status(1)
    Tx.Put16( (int)H.Opcode );
    if ( !H.Encrypt ) Tx.Put64( H.ID );
    byte [] ServerToken = Random.GetServerToken( H.REP );
    Tx.PutBytes( ServerToken );
    Tx.Put8( (int)H.Status );
    EncryptAndSendResponse( Tx.Buffer, H );
  }

  private void EncryptAndSendResponse( byte [] Response, QrpHeader H )
  {
    if ( H.Encrypt )
    {
      Response =  Crypt.EncodeResponse( Response, H.Sa );
    }
    Socket.SendTo( Response, H.REP );
  }

  public QrpListener( int Port, bool IPv6 ) : base( Port, IPv6 ? IPAddress.IPv6Any : IPAddress.Any  )
  {
  }

}