package org.msgpack.jruby;


import java.math.BigInteger;
import java.nio.ByteBuffer;
import java.nio.BufferUnderflowException;
import java.util.Iterator;
import java.util.Arrays;

import org.jruby.Ruby;
import org.jruby.RubyObject;
import org.jruby.RubyClass;
import org.jruby.RubyBignum;
import org.jruby.RubyString;
import org.jruby.RubyArray;
import org.jruby.RubyHash;
import org.jruby.RubyInteger;
import org.jruby.exceptions.RaiseException;
import org.jruby.runtime.builtin.IRubyObject;
import org.jruby.util.ByteList;

import org.jcodings.Encoding;
import org.jcodings.specific.UTF8Encoding;

import static org.msgpack.jruby.Types.*;


public class Decoder implements Iterator<IRubyObject> {
  private final Ruby runtime;
  private final Encoding binaryEncoding;
  private final Encoding utf8Encoding;
  private final RubyClass unpackErrorClass;
  private final RubyClass underflowErrorClass;
  private final RubyClass malformedFormatErrorClass;
  private final RubyClass stackErrorClass;
  private final RubyClass unexpectedTypeErrorClass;
  private final RubyClass unknownExtTypeErrorClass;

  private Unpacker unpacker;
  private ByteBuffer buffer;
  private boolean symbolizeKeys;
  private boolean freeze;
  private boolean allowUnknownExt;

  public Decoder(Ruby runtime) {
    this(runtime, null, new byte[] {}, 0, 0, false, false, false);
  }

  public Decoder(Ruby runtime, Unpacker unpacker) {
    this(runtime, unpacker, new byte[] {}, 0, 0, false, false, false);
  }

  public Decoder(Ruby runtime, byte[] bytes) {
    this(runtime, null, bytes, 0, bytes.length, false, false, false);
  }

  public Decoder(Ruby runtime, Unpacker unpacker, byte[] bytes) {
    this(runtime, unpacker, bytes, 0, bytes.length, false, false, false);
  }

  public Decoder(Ruby runtime, Unpacker unpacker, byte[] bytes, boolean symbolizeKeys, boolean freeze, boolean allowUnknownExt) {
    this(runtime, unpacker, bytes, 0, bytes.length, symbolizeKeys, freeze, allowUnknownExt);
  }

  public Decoder(Ruby runtime, Unpacker unpacker, byte[] bytes, int offset, int length) {
    this(runtime, unpacker, bytes, offset, length, false, false, false);
  }

  public Decoder(Ruby runtime, Unpacker unpacker, byte[] bytes, int offset, int length, boolean symbolizeKeys, boolean freeze, boolean allowUnknownExt) {
    this.runtime = runtime;
    this.unpacker = unpacker;
    this.symbolizeKeys = symbolizeKeys;
    this.freeze = freeze;
    this.allowUnknownExt = allowUnknownExt;
    this.binaryEncoding = runtime.getEncodingService().getAscii8bitEncoding();
    this.utf8Encoding = UTF8Encoding.INSTANCE;
    this.unpackErrorClass = runtime.getModule("MessagePack").getClass("UnpackError");
    this.underflowErrorClass = runtime.getModule("MessagePack").getClass("UnderflowError");
    this.malformedFormatErrorClass = runtime.getModule("MessagePack").getClass("MalformedFormatError");
    this.stackErrorClass = runtime.getModule("MessagePack").getClass("StackError");
    this.unexpectedTypeErrorClass = runtime.getModule("MessagePack").getClass("UnexpectedTypeError");
    this.unknownExtTypeErrorClass = runtime.getModule("MessagePack").getClass("UnknownExtTypeError");
    this.symbolizeKeys = symbolizeKeys;
    this.allowUnknownExt = allowUnknownExt;
    feed(bytes, offset, length);
  }

  public void feed(byte[] bytes) {
    feed(bytes, 0, bytes.length);
  }

  public void feed(byte[] bytes, int offset, int length) {
    if (buffer == null) {
      buffer = ByteBuffer.wrap(bytes, offset, length);
    } else {
      ByteBuffer newBuffer = ByteBuffer.allocate(buffer.remaining() + length);
      newBuffer.put(buffer);
      newBuffer.put(bytes, offset, length);
      newBuffer.flip();
      buffer = newBuffer;
    }
  }

  public void reset() {
    buffer = null;
  }

  public int offset() {
    return buffer.position();
  }

  private IRubyObject consumeUnsignedLong() {
    long value = buffer.getLong();
    if (value < 0) {
      return RubyBignum.newBignum(runtime, BigInteger.valueOf(value & ((1L<<63)-1)).setBit(63));
    } else {
      return runtime.newFixnum(value);
    }
  }

  private IRubyObject consumeString(int size, Encoding encoding) {
    byte[] bytes = readBytes(size);
    ByteList byteList = new ByteList(bytes, encoding);
    RubyString string = runtime.newString(byteList);
    if (this.freeze) {
      string = runtime.freezeAndDedupString(string);
    }
    return string;
  }

  private IRubyObject consumeArray(int size) {
    IRubyObject[] elements = new IRubyObject[size];
    for (int i = 0; i < size; i++) {
      elements[i] = next();
    }
    return runtime.newArray(elements);
  }

  private IRubyObject consumeHash(int size) {
    RubyHash hash = RubyHash.newHash(runtime);
    for (int i = 0; i < size; i++) {
      IRubyObject key = next();
      if (key instanceof RubyString) {
        if (this.symbolizeKeys) {
          key = ((RubyString) key).intern();
        } else {
          key = runtime.freezeAndDedupString((RubyString) key);
        }
      }

      hash.fastASet(key, next());
    }
    return hash;
  }

  private IRubyObject consumeExtension(int size) {
    int type = buffer.get();
    if (unpacker != null) {
      ExtensionRegistry.ExtensionEntry entry = unpacker.lookupExtensionByTypeId(type);
      if (entry != null) {
        IRubyObject proc = entry.getUnpackerProc();
        if (entry.isRecursive()) {
          return proc.callMethod(runtime.getCurrentContext(), "call", unpacker);
        } else {
          ByteList byteList = new ByteList(readBytes(size), runtime.getEncodingService().getAscii8bitEncoding());
          return proc.callMethod(runtime.getCurrentContext(), "call", runtime.newString(byteList));
        }
      }
    }

    if (this.allowUnknownExt) {
      return ExtensionValue.newExtensionValue(runtime, type, readBytes(size));
    }

    throw runtime.newRaiseException(unknownExtTypeErrorClass, "unexpected extension type");
  }

  private byte[] readBytes(int size) {
    byte[] payload = new byte[size];
    buffer.get(payload);
    return payload;
  }

  @Override
  public void remove() {
    throw new UnsupportedOperationException();
  }

  @Override
  public boolean hasNext() {
    return buffer.remaining() > 0;
  }

  public IRubyObject read_array_header() {
    int position = buffer.position();
    try {
      byte b = buffer.get();
      if ((b & 0xf0) == 0x90) {
        return runtime.newFixnum(b & 0x0f);
      } else if (b == ARY16) {
        return runtime.newFixnum(buffer.getShort() & 0xffff);
      } else if (b == ARY32) {
        return runtime.newFixnum(buffer.getInt());
      }
      throw runtime.newRaiseException(unexpectedTypeErrorClass, "unexpected type");
    } catch (RaiseException re) {
      buffer.position(position);
      throw re;
    } catch (BufferUnderflowException bue) {
      buffer.position(position);
      throw runtime.newRaiseException(underflowErrorClass, "Not enough bytes available");
    }
  }

  public IRubyObject read_map_header() {
    int position = buffer.position();
    try {
      byte b = buffer.get();
      if ((b & 0xf0) == 0x80) {
        return runtime.newFixnum(b & 0x0f);
      } else if (b == MAP16) {
        return runtime.newFixnum(buffer.getShort() & 0xffff);
      } else if (b == MAP32) {
        return runtime.newFixnum(buffer.getInt());
      }
      throw runtime.newRaiseException(unexpectedTypeErrorClass, "unexpected type");
    } catch (RaiseException re) {
      buffer.position(position);
      throw re;
    } catch (BufferUnderflowException bue) {
      buffer.position(position);
      throw runtime.newRaiseException(underflowErrorClass, "Not enough bytes available");
    }
  }

  @Override
  public IRubyObject next() {
    IRubyObject next = consumeNext();
    if (freeze) {
      next.setFrozen(true);
    }
    return next;
  }

  private IRubyObject consumeNext() {
    int position = buffer.position();
    try {
      byte b = buffer.get();
      outer: switch ((b >> 4) & 0xf) {
      case 0x8: return consumeHash(b & 0x0f);
      case 0x9: return consumeArray(b & 0x0f);
      case 0xa:
      case 0xb: return consumeString(b & 0x1f, utf8Encoding);
      case 0xc:
        switch (b) {
        case NIL:      return runtime.getNil();
        case FALSE:    return runtime.getFalse();
        case TRUE:     return runtime.getTrue();
        case BIN8:     return consumeString(buffer.get() & 0xff, binaryEncoding);
        case BIN16:    return consumeString(buffer.getShort() & 0xffff, binaryEncoding);
        case BIN32:    return consumeString(buffer.getInt(), binaryEncoding);
        case VAREXT8:  return consumeExtension(buffer.get() & 0xff);
        case VAREXT16: return consumeExtension(buffer.getShort() & 0xffff);
        case VAREXT32: return consumeExtension(buffer.getInt());
        case FLOAT32:  return runtime.newFloat(buffer.getFloat());
        case FLOAT64:  return runtime.newFloat(buffer.getDouble());
        case UINT8:    return runtime.newFixnum(buffer.get() & 0xffL);
        case UINT16:   return runtime.newFixnum(buffer.getShort() & 0xffffL);
        case UINT32:   return runtime.newFixnum(buffer.getInt() & 0xffffffffL);
        case UINT64:   return consumeUnsignedLong();
        default: break outer;
        }
      case 0xd:
        switch (b) {
        case INT8:     return runtime.newFixnum(buffer.get());
        case INT16:    return runtime.newFixnum(buffer.getShort());
        case INT32:    return runtime.newFixnum(buffer.getInt());
        case INT64:    return runtime.newFixnum(buffer.getLong());
        case FIXEXT1:  return consumeExtension(1);
        case FIXEXT2:  return consumeExtension(2);
        case FIXEXT4:  return consumeExtension(4);
        case FIXEXT8:  return consumeExtension(8);
        case FIXEXT16: return consumeExtension(16);
        case STR8:     return consumeString(buffer.get() & 0xff, utf8Encoding);
        case STR16:    return consumeString(buffer.getShort() & 0xffff, utf8Encoding);
        case STR32:    return consumeString(buffer.getInt(), utf8Encoding);
        case ARY16:    return consumeArray(buffer.getShort() & 0xffff);
        case ARY32:    return consumeArray(buffer.getInt());
        case MAP16:    return consumeHash(buffer.getShort() & 0xffff);
        case MAP32:    return consumeHash(buffer.getInt());
        default: break outer;
        }
      case 0xe:
      case 0xf: return runtime.newFixnum((0x1f & b) - 0x20);
      default: return runtime.newFixnum(b);
      }
      buffer.position(position);
      throw runtime.newRaiseException(malformedFormatErrorClass, "Illegal byte sequence");
    } catch (RaiseException re) {
      buffer.position(position);
      throw re;
    } catch (BufferUnderflowException bue) {
      buffer.position(position);
      throw runtime.newRaiseException(underflowErrorClass, "Not enough bytes available");
    }
  }
}