/*
 *  Copyright (C) 1999 AT&T Laboratories Cambridge.  All Rights Reserved.
 *
 *  This is free software; you can redistribute it and/or modify
 *  it under the terms of the GNU General Public License as published by
 *  the Free Software Foundation; either version 2 of the License, or
 *  (at your option) any later version.
 *
 *  This software is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *  GNU General Public License for more details.
 *
 *  You should have received a copy of the GNU General Public License
 *  along with this software; if not, write to the Free Software
 *  Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA  02111-1307,
 *  USA.
 */

/*
 * rfbproto.c - functions to deal with client side of RFB protocol.
 */

#include <time.h>

#include "oslib/os.h"
#include "oslib/socket.h"

#include "vncviewer.h"
#include "display.h"
#include "vncauth.h"
#include "vnckeys.h"
#include "ip.h"
#include "sockets.h"

#include "zlib.h"


static long ReadCompactLen (void);
static Bool HandleRRE8(int rx, int ry, int rw, int rh);
static Bool HandleCoRRE8(int rx, int ry, int rw, int rh);
static Bool HandleHextile8(int rx, int ry, int rw, int rh);
static Bool HandleRRE16(int rx, int ry, int rw, int rh);
static Bool HandleCoRRE16(int rx, int ry, int rw, int rh);
static Bool HandleHextile16(int rx, int ry, int rw, int rh);
static Bool HandleTight16(int rx, int ry, int rw, int rh);

int rfbsock = -1;
char *desktopName;
rfbPixelFormat myFormat;
rfbServerInitMsg si;
char *serverCutText = NULL;
Bool newServerCutText = False;

int endianTest = 1;

int subrect = 0;

/* note that the CoRRE encoding uses this buffer and assumes it is big enough
   to hold 255 * 255 * 32 bits -> 260100 bytes.  640*480 = 307200 bytes */
/* also hextile assumes it is big enough to hold 16 * 16 * 32 bits */

#define BUFFER_SIZE (640*480)
static char buffer[BUFFER_SIZE];

/*
 * Variables for the ``tight'' encoding implementation.
 */

/* Separate buffer for compressed data. */
#define ZLIB_BUFFER_SIZE 4096
static char zlib_buffer[ZLIB_BUFFER_SIZE];

/* Four independent compression streams for zlib library. */
static z_stream zlibStream[4];
static Bool zlibStreamActive[4] = {
  False, False, False, False
};

/* Filter stuff. Should be initialized by filter initialization code. */
static Bool cutZeros;
static int rectWidth, rectColors;
static char tightPalette[256*4];
static CARD8 tightPrevRow[2048*3*sizeof(CARD16)];


/*
 * ConnectToRFBServer.
 */

Bool ConnectToRFBServer(unsigned int host, int port) {

  rfbsock = ConnectToTcpAddr(host, port);

  if (rfbsock < 0) {
    fprintf(stderr,"Unable to connect to VNC server\n");
    return False;
  }
  if (debug)   fprintf(stderr, "Connected via socket %d\n", rfbsock);

  return SetNonBlocking(rfbsock);
}


/*
 * InitialiseRFBConnection.
 */

Bool
InitialiseRFBConnection() {
  rfbProtocolVersionMsg pv;
  int major,minor;
  CARD32 authScheme, reasonLen, authResult;
  char *reason;
  CARD8 challenge[CHALLENGESIZE];
  char passwd[9];
  int i;
  rfbClientInitMsg ci;

  if (!ReadFromRFBServer(pv, sz_rfbProtocolVersionMsg)) {
    fprintf(stderr, "Failed to read version\n");
    return False;
  }

  errorMessageOnReadFailure = True;

  pv[sz_rfbProtocolVersionMsg] = 0;

  if (sscanf(pv,rfbProtocolVersionFormat,&major,&minor) != 2) {
    fprintf(stderr,"Not a valid VNC server\n");
    return False;
  }

  if (debug)
    fprintf(stderr,"VNC server supports protocol version %d.%d (viewer %d.%d)\n",
     	  major, minor, rfbProtocolMajorVersion, rfbProtocolMinorVersion);

  major = rfbProtocolMajorVersion;
  minor = rfbProtocolMinorVersion;

  sprintf(pv,rfbProtocolVersionFormat,major,minor);

  if (!WriteExact(rfbsock, pv, sz_rfbProtocolVersionMsg)) return False;

  if (!ReadFromRFBServer((char *)&authScheme, 4)) return False;

  authScheme = Swap32IfLE(authScheme);

  switch (authScheme) {

  case rfbConnFailed:
    if (!ReadFromRFBServer((char *)&reasonLen, 4)) return False;
    reasonLen = Swap32IfLE(reasonLen);

    reason = malloc(reasonLen);

    if (!ReadFromRFBServer(reason, reasonLen)) return False;

    fprintf(stderr,"VNC connection failed: %.*s\n",(int)reasonLen, reason);
    return False;

  case rfbNoAuth:
    if (debug)   fprintf(stderr,"No authentication needed\n");
    break;

  case rfbVncAuth:
    if (!ReadFromRFBServer((char *)challenge, CHALLENGESIZE)) return False;
    strcpy(passwd, password);
    vncEncryptBytes(challenge, passwd);

	/* Lose the password from memory */
    for (i = strlen(passwd); i >= 0; i--) {
      passwd[i] = '\0';
    }

    if (!WriteExact(rfbsock, (void *)challenge, CHALLENGESIZE)) return False;

    if (!ReadFromRFBServer((char *)&authResult, 4)) return False;

    authResult = Swap32IfLE(authResult);

    switch (authResult) {
    case rfbVncAuthOK:
      if (debug)   fprintf(stderr,"VNC authentication succeeded\n");
      break;
    case rfbVncAuthFailed:
      fprintf(stderr,"VNC authentication failed\n");
      return False;
    case rfbVncAuthTooMany:
      fprintf(stderr,"VNC authentication failed - too many tries\n");
      return False;
    default:
      fprintf(stderr,"Unknown VNC authentication result: %d\n",
	      (int)authResult);
      return False;
    }
    break;

  default:
    fprintf(stderr,"Unknown authentication scheme from VNC server: %d\n",
	    (int)authScheme);
    return False;
  }

//  ci.shared = (appData.shareDesktop ? 1 : 0);

  if (!WriteExact(rfbsock, (char *)&ci, sz_rfbClientInitMsg)) return False;

  if (!ReadFromRFBServer((char *)&si, sz_rfbServerInitMsg)) return False;

  si.framebufferWidth = Swap16IfLE(si.framebufferWidth);
  si.framebufferHeight = Swap16IfLE(si.framebufferHeight);
  si.format.redMax = Swap16IfLE(si.format.redMax);
  si.format.greenMax = Swap16IfLE(si.format.greenMax);
  si.format.blueMax = Swap16IfLE(si.format.blueMax);
  si.nameLength = Swap32IfLE(si.nameLength);

  if (debug)
    fprintf(stderr,"Desktop size = %d x %d\n", si.framebufferWidth, si.framebufferHeight);

  desktopName = malloc(si.nameLength + 1);

  if (!ReadFromRFBServer(desktopName, si.nameLength)) return False;

  desktopName[si.nameLength] = 0;

  return True;
}


/*
 * SetFormatAndEncodings.
 */

Bool SetFormatAndEncodings() {
  rfbSetPixelFormatMsg spf;
  char buf[sz_rfbSetEncodingsMsg + MAX_ENCODINGS * 4];
  rfbSetEncodingsMsg *se;
  CARD32 *encs;
  int len = 0;

  se = (rfbSetEncodingsMsg *)buf;
  encs = (CARD32 *)(buf+sz_rfbSetEncodingsMsg);

  if (bpp == 16) {
    myFormat.bitsPerPixel = 16;
    myFormat.depth = 15;
    myFormat.trueColour = 1;
    myFormat.bigEndian = 0;
    myFormat.redShift = 0;
    myFormat.greenShift = 5;
    myFormat.blueShift = 10;
    myFormat.redMax = 31;
    myFormat.greenMax = 31;
    myFormat.blueMax = 31;
/*
  myFormat.bitsPerPixel = 16;
  myFormat.depth = 16;
  myFormat.trueColour = 1;
  myFormat.bigEndian = 0;
  myFormat.redShift = 11;
  myFormat.greenShift = 5;
  myFormat.blueShift = 0;
  myFormat.redMax = 31;
  myFormat.greenMax = 63;
  myFormat.blueMax = 31;
*/
  } else {
    myFormat.bitsPerPixel = 8;
    myFormat.depth = 8;
    myFormat.trueColour = 0;
    myFormat.bigEndian = 0;
    myFormat.redShift = 0;
    myFormat.greenShift = 3;
    myFormat.blueShift = 6;
    myFormat.redMax = 7;
    myFormat.greenMax = 7;
    myFormat.blueMax = 3;
  }
  spf.type = rfbSetPixelFormat;
  spf.pad1 = spf.pad2 = 0;
  memcpy(&spf.format, &myFormat, sizeof(rfbPixelFormat));
  spf.format.redMax = Swap16IfLE(spf.format.redMax);
  spf.format.greenMax = Swap16IfLE(spf.format.greenMax);
  spf.format.blueMax = Swap16IfLE(spf.format.blueMax);

  if (!WriteExact(rfbsock, (char *)&spf, sz_rfbSetPixelFormatMsg))
    return False;

  se->type = rfbSetEncodings;
  se->nEncodings = 0;

  encs[se->nEncodings++] = Swap32IfLE(rfbEncodingCopyRect);
  if (supporttight) {
    // if 'tight' is supported, give it high priority
    encs[se->nEncodings++] = Swap32IfLE(rfbEncodingTight);
    fprintf(stderr, "Tight requested\n");
  }
  // hextile is usually the highest priority
  encs[se->nEncodings++] = Swap32IfLE(rfbEncodingHextile);
  // coRRE and RRE are old, outdated encodings, so low priority
  encs[se->nEncodings++] = Swap32IfLE(rfbEncodingCoRRE);
  encs[se->nEncodings++] = Swap32IfLE(rfbEncodingRRE);
  // not so much encodings as encoding-controls:
  if (checkclientbitmap)
    encs[se->nEncodings++] = Swap32IfLE(rfbEncodingCheckClientBitmap);
  if ((myFormat.bitsPerPixel == 16) && (dither8))
    encs[se->nEncodings++] = Swap32IfLE(rfbEncodingHextileDither8);
// XCursor not implemented - doesn't work well with hardware-pointer
//  encs[se->nEncodings++] = Swap32IfLE(rfbEncodingXCursor);

  len = sz_rfbSetEncodingsMsg + se->nEncodings * 4;

  se->pad = 0;
  se->nEncodings = Swap16IfLE(se->nEncodings);

  if (!WriteExact(rfbsock, buf, len)) return False;

  return True;
}


/*
 * SendIncrementalFramebufferUpdateRequest.
 */

Bool
SendIncrementalFramebufferUpdateRequest() {
  return SendFramebufferUpdateRequest(0, 0, si.framebufferWidth,
				      si.framebufferHeight, True);
}


/*
 * SendFramebufferUpdateRequest.
 */

Bool
SendFramebufferUpdateRequest(int x, int y, int w, int h, Bool incremental) {
  rfbFramebufferUpdateRequestMsg fur;

  fur.type = rfbFramebufferUpdateRequest;
  fur.incremental = incremental ? 1 : 0;
  fur.x = Swap16IfLE(x);
  fur.y = Swap16IfLE(y);
  fur.w = Swap16IfLE(w);
  fur.h = Swap16IfLE(h);

  if (!WriteExact(rfbsock, (char *)&fur, sz_rfbFramebufferUpdateRequestMsg)) {
    fprintf(stderr, "Failed to send buffer update request\n");
    return False;
  }

  return True;
}


/*
 * SendPointerEvent.
 */

Bool
SendPointerEvent(int x, int y, int buttonMask) {
  rfbPointerEventMsg pe;

  pe.type = rfbPointerEvent;
  pe.buttonMask = buttonMask;
  if (x < 0) x = 0;
  if (y < 0) y = 0;
  pe.x = Swap16IfLE(x);
  pe.y = Swap16IfLE(y);
  return WriteExact(rfbsock, (char *)&pe, sz_rfbPointerEventMsg);
}


/*
 * SendKeyEvent.
 */

Bool
SendKeyEvent(CARD32 key, Bool down) {
  rfbKeyEventMsg ke;

  // HBP - it looks as if it is necessary to send a UpdateRequest
  // before a TAB, otherwise the server seems to get stuck
  if (key == VNCKEY_TAB)
    if (SendIncrementalFramebufferUpdateRequest() == False)
      return False;

  ke.type = rfbKeyEvent;
  ke.down = down ? 1 : 0;
  ke.key = Swap32IfLE(key);
  return WriteExact(rfbsock, (char *)&ke, sz_rfbKeyEventMsg);
}


/*
 * SendClientCutText.
 */

Bool
SendClientCutText(char *str, int len) {
  rfbClientCutTextMsg cct;

  if (serverCutText)
    free(serverCutText);
  serverCutText = NULL;

  cct.type = rfbClientCutText;
  cct.length = Swap32IfLE(len);
  return  (WriteExact(rfbsock, (char *)&cct, sz_rfbClientCutTextMsg) &&
	   WriteExact(rfbsock, str, len));
}

/*
 * CloseConnection
 */
void CloseConnection() {
  CloseSocket(rfbsock);
}

/*
 * HandleRFBServerMessage.
 */

Bool
HandleRFBServerMessage() {
  rfbServerToClientMsg msg;

  if (buffered == 0 && ip_ready(rfbsock) < 1)
    return True;
  if (!ReadFromRFBServer((char *)&msg, 1))
    return False;

  switch (msg.type) {
  case rfbSetColourMapEntries:
  {
    int i;
    CARD16 rgb[3];

    if (!ReadFromRFBServer(((char *)&msg) + 1,
			   sz_rfbSetColourMapEntriesMsg - 1))
      return False;

    msg.scme.firstColour = Swap16IfLE(msg.scme.firstColour);
    msg.scme.nColours = Swap16IfLE(msg.scme.nColours);

    for (i = 0; i < msg.scme.nColours; i++) {
      if (!ReadFromRFBServer((char *)rgb, 6))
	return False;
    }

    break;
  }

  case rfbFramebufferUpdate:
  {
    rfbFramebufferUpdateRectHeader rect;
    int linesToRead;
    int bytesPerLine;
    int i;

    if (!ReadFromRFBServer(((char *)&msg.fu) + 1,
			   sz_rfbFramebufferUpdateMsg - 1))
      return False;

    msg.fu.nRects = Swap16IfLE(msg.fu.nRects);
    for (i = 0; i < msg.fu.nRects; i++) {
      if (!ReadFromRFBServer((char *)&rect, sz_rfbFramebufferUpdateRectHeader))
	return False;

      rect.r.x = Swap16IfLE(rect.r.x);
      rect.r.y = Swap16IfLE(rect.r.y);
      rect.r.w = Swap16IfLE(rect.r.w);
      rect.r.h = Swap16IfLE(rect.r.h);
      rect.encoding = Swap32IfLE(rect.encoding);

      if (rect.encoding == rfbEncodingXCursor) {
	if (!HandleXCursor(rect.r.x, rect.r.y, rect.r.w, rect.r.h)) {
	  return False;
	}
	continue;
      }

      if ((rect.r.x + rect.r.w > si.framebufferWidth) ||
	  (rect.r.y + rect.r.h > si.framebufferHeight))	{
	  fprintf(stderr,"Rect too large: %dx%d at (%d, %d)\n",
		  rect.r.w, rect.r.h, rect.r.x, rect.r.y);
	  return False;
	}

      if ((rect.r.h * rect.r.w) == 0) {
	fprintf(stderr,"Zero size rect - ignoring\n");
	continue;
      }

      if (rect.encoding != rfbEncodingCheckClientBitmap &&
          rect.encoding != rfbEncodingCopyRect)
        display_add_rectangle(rect.r.x, rect.r.y, rect.r.w, rect.r.h);

      switch (rect.encoding) {

      case rfbEncodingRaw:
	bytesPerLine = rect.r.w * bpp / 8;
	linesToRead = BUFFER_SIZE / bytesPerLine;

	while (rect.r.h > 0) {
	  if (linesToRead > rect.r.h)
	    linesToRead = rect.r.h;
	  if (linesToRead * bytesPerLine > 36*1024)
	    linesToRead = 36*1024/bytesPerLine;

	  if (!ReadFromRFBServer(buffer,bytesPerLine * linesToRead))
	    return False;

	  display_raw(rect.r.x, rect.r.y, rect.r.w, linesToRead, buffer, 0);

	  rect.r.h -= linesToRead;
	  rect.r.y += linesToRead;

	}
	break;

      case rfbEncodingCopyRect:
      {
	rfbCopyRect cr;

	if (!ReadFromRFBServer((char *)&cr, sz_rfbCopyRect))
	  return False;

	cr.srcX = Swap16IfLE(cr.srcX);
	cr.srcY = Swap16IfLE(cr.srcY);

        if (fast_copy_ok())
	  BlockCopy(rect.r.x, rect.r.y, rect.r.w, rect.r.h, cr.srcX, cr.srcY);
	else
	  display_add_rectangle(rect.r.x, rect.r.y, rect.r.w, rect.r.h);

        display_copy(rect.r.x, rect.r.y, rect.r.w, rect.r.h, cr.srcX, cr.srcY);
	break;
      }

      case rfbEncodingCheckClientBitmap:
        // no data
        break;

      case rfbEncodingRRE:
        if (bpp == 8) {
	  if (!HandleRRE8(rect.r.x,rect.r.y,rect.r.w,rect.r.h))     return False;
	} else {
          if (!HandleRRE16(rect.r.x,rect.r.y,rect.r.w,rect.r.h))     return False;
	}
	break;

      case rfbEncodingCoRRE:
        if (bpp == 8) {
          if (!HandleCoRRE8(rect.r.x,rect.r.y,rect.r.w,rect.r.h))   return False;
        } else {
          if (!HandleCoRRE16(rect.r.x,rect.r.y,rect.r.w,rect.r.h))   return False;
        }
        break;

      case rfbEncodingHextile:
        if (bpp == 8) {
          if (!HandleHextile8(rect.r.x,rect.r.y,rect.r.w,rect.r.h))   return False;
        } else {
          if (!HandleHextile16(rect.r.x,rect.r.y,rect.r.w,rect.r.h))   return False;
        }
        break;

      case rfbEncodingTight:
        if (bpp == 8) {
        } else {
          if (!HandleTight16(rect.r.x,rect.r.y,rect.r.w,rect.r.h))    return False;
        }
        break;

      default:
	fprintf(stderr,"Unknown rect encoding %d\n",
		(int)rect.encoding);
	return False;
      }
    }

    if (!SendIncrementalFramebufferUpdateRequest())   return False;

    break;
  }

  case rfbBell:
    xos_bell();
    break;

  case rfbServerCutText:
  {
    if (!ReadFromRFBServer(((char *)&msg) + 1,
			   sz_rfbServerCutTextMsg - 1))
      return False;
    msg.sct.length = Swap32IfLE(msg.sct.length);

    if (serverCutText)
      free(serverCutText);

    serverCutText = malloc(msg.sct.length+1);

    if (!ReadFromRFBServer(serverCutText, msg.sct.length))
      return False;

    serverCutText[msg.sct.length] = 0;

    newServerCutText = True;

    // HBP - it seems it is necessary to send an UpdateRequest
    // after a cut text has been received, otherwise the server
    // gets stuck
    if (!SendIncrementalFramebufferUpdateRequest())   return False;

    break;
  }

  default:
    fprintf(stderr,"Unknown message type %d from VNC server\n",msg.type);
    return False;
  }

  return True;
}


//#define GET_PIXEL16(pix, ptr) (((CARD8*)&(pix))[0] = *(ptr)++, \
//			       ((CARD8*)&(pix))[1] = *(ptr)++)

#define BPP 16
#include "rre16.c"
#include "corre16.c"
#include "hextile16.c"
#include "tight16.c"
#undef BPP

#define BPP 8
#include "rre8.c"
#include "corre8.c"
#include "hextile8.c"


static long ReadCompactLen (void)
{
  long len;
  CARD8 b;

  if (!ReadFromRFBServer((char *)&b, 1))
    return -1;
  len = (long)b & 0x7F;
  if (b & 0x80) {
    if (!ReadFromRFBServer((char *)&b, 1))
      return -1;
    len |= ((long)b & 0x7F) << 7;
    if (b & 0x80) {
      if (!ReadFromRFBServer((char *)&b, 1))
        return -1;
      len |= ((long)b & 0xFF) << 14;
    }
  }
  return len;
}
