/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.apache.juneau.marshall.cbor;

import static org.apache.juneau.commons.utils.IoUtils.*;
import static org.apache.juneau.commons.utils.Shorts.*;
import static org.apache.juneau.marshall.cbor.DataType.*;

import java.io.*;
import java.math.*;

import org.apache.juneau.marshall.parser.*;

/**
 * Specialized input stream for parsing CBOR streams (RFC 8949).
 *
 * <p>
 * Reads CBOR-encoded binary data, extracting major types and additional information from
 * initial bytes and decoding values according to RFC 8949.
 *
 * <h5 class='section'>Notes:</h5><ul>
 *    <li class='note'>This class is not intended for external use.
 *    <li class='note'>Indefinite-length encoding (RFC 8949 §3.2) is supported: arrays/maps decode via the
 *       BREAK-aware container loop, and chunked text/byte strings (heads <c>0x7F</c>/<c>0x5F</c>) are
 *       reassembled from their definite-length pieces.
 * </ul>
 *
 * <h5 class='section'>See Also:</h5><ul>
 *    <li class='link'><a class="doclink" href="https://juneau.apache.org/docs/topics/Cbor">CBOR</a>
 *    <li class='link'><a class="doclink" href="https://www.rfc-editor.org/rfc/rfc8949.html">RFC 8949</a>
 * </ul>
 */
public class CborInputStream extends ParserInputStream {

   /** Default cap (16 MiB) for wire-declared lengths, mirroring {@code BsonInputStream}. */
   static final int DEFAULT_MAX_LENGTH = 16 * 1024 * 1024;

   // checkLength's "what" label for byte/text-string length checks (definite-length and reassembled indefinite-length).
   private static final String WHAT_STRING = "string";

   private long length;
   private int lastInitialByte;
   private DataType lastDataType;
   // Single-byte pushback for indefinite-length container loops that need to peek for BREAK
   // without consuming the next data item.  -1 means no pushback present.
   private int pushedBackByte = -1;
   // True when the last-read data item carried additional-info 31 (indefinite-length array/map/string).
   // Tracked separately from length so a definite value whose argument happens to be -1 (an 8-byte
   // 0xFFFFFFFFFFFFFFFF) is NOT mistaken for an indefinite item.
   private boolean indefinite;

   /**
    * Constructor.
    *
    * @param pipe The parser input.
    * @throws IOException Thrown by underlying stream.
    */
   protected CborInputStream(ParserPipe pipe) throws IOException {
      super(pipe);
   }

   @Override /* InputStream */
   public int read() throws IOException {
      if (pushedBackByte != -1) {
         int b = pushedBackByte;
         pushedBackByte = -1;
         return b;
      }
      return super.read();
   }

   /**
    * Peeks one byte and returns whether it is the CBOR BREAK marker (0xFF).  Non-BREAK bytes
    * are pushed back so the next {@link #read()} sees them.
    *
    * <p>
    * Used by definite/indefinite-length container loops to detect the end of an indefinite
    * container without consuming the next data item.
    *
    * @return <jk>true</jk> if the next byte is BREAK; <jk>false</jk> otherwise (or at EOF).
    * @throws IOException If the underlying stream fails.
    */
   boolean peekBreak() throws IOException {
      int b = read();
      if (b == -1)
         return false;
      if (b == 0xFF) {
         lastDataType = BREAK;
         length = 0;
         return true;
      }
      pushedBackByte = b;
      return false;
   }

   /**
    * Reads the initial byte and returns it.
    *
    * @return The initial byte (0-255), or -1 on EOF.
    * @throws IOException If read fails.
    */
   int readInitialByte() throws IOException {
      int b = read();
      if (b == -1)
         throw ioex("Unexpected end of CBOR input");
      return b & 0xFF;
   }

   /**
    * Extracts the major type from an initial byte (high 3 bits).
    *
    * @param initialByte The initial byte.
    * @return Major type 0-7.
    */
   static int getMajorType(int initialByte) {
      return (initialByte >> 5) & 0x07;
   }

   /**
    * Extracts the additional information from an initial byte (low 5 bits).
    *
    * @param initialByte The initial byte.
    * @return Additional info 0-31.
    */
   static int getAdditionalInfo(int initialByte) {
      return initialByte & 0x1F;
   }

   /**
    * Reads the argument based on additional info.
    * For 0-23, returns additionalInfo; for 24-27, reads 1/2/4/8 bytes big-endian.
    *
    * @param additionalInfo The additional info from the initial byte.
    * @return The argument value.
    * @throws IOException If read fails.
    */
   long readArgument(int additionalInfo) throws IOException {
      if (additionalInfo <= 23)
         return additionalInfo;
      if (additionalInfo == 24)
         return readUInt1();
      if (additionalInfo == 25)
         return readUInt2();
      if (additionalInfo == 26)
         return readUInt4();
      if (additionalInfo == 27)
         return readUInt8();
      if (additionalInfo == 31)
         throw ioex("Indefinite-length CBOR encoding not supported");
      throw ioex("Reserved additional info value: %s", additionalInfo);
   }

   /**
    * Reads 1 unsigned byte.
    */
   @SuppressWarnings({
      "java:S4144" // Identical to readInitialByte by design; each method has distinct semantic purpose
   })
   int readUInt1() throws IOException {
      int b = read();
      if (b == -1)
         throw ioex("Unexpected end of CBOR input");
      return b & 0xFF;
   }

   /**
    * Reads 2 bytes big-endian as unsigned integer.
    */
   int readUInt2() throws IOException {
      return (readUInt1() << 8) | readUInt1();
   }

   /**
    * Reads 4 bytes big-endian as unsigned long.
    */
   long readUInt4() throws IOException {
      long l = readUInt1();
      l = (l << 8) | readUInt1();
      l = (l << 8) | readUInt1();
      l = (l << 8) | readUInt1();
      return l;
   }

   /**
    * Reads 8 bytes big-endian as unsigned long.
    */
   long readUInt8() throws IOException {
      long l = readUInt1();
      l = (l << 8) | readUInt1();
      l = (l << 8) | readUInt1();
      l = (l << 8) | readUInt1();
      l = (l << 8) | readUInt1();
      l = (l << 8) | readUInt1();
      l = (l << 8) | readUInt1();
      l = (l << 8) | readUInt1();
      return l;
   }

   /**
    * Reads the next CBOR data item's type and length/argument, returning the DataType.
    * The argument is stored and can be retrieved via {@link #readLength()}.
    *
    * @return The DataType of the next item.
    * @throws IOException If read fails.
    */
   @SuppressWarnings({
      "java:S3776" // Cognitive complexity acceptable for CBOR data type reading
   })
   DataType readDataType() throws IOException {
      int ib = readInitialByte();
      lastInitialByte = ib;
      int majorType = getMajorType(ib);
      int additionalInfo = getAdditionalInfo(ib);

      DataType dt;
      switch (majorType) {
         case 0 -> dt = UINT;
         case 1 -> dt = NINT;
         case 2 -> dt = BINARY;
         case 3 -> dt = STRING;
         case 4 -> dt = ARRAY;
         case 5 -> dt = MAP;
         case 6 -> dt = TAG;
         case 7 -> {
            switch (additionalInfo) {
               case 20 -> dt = BOOLEAN;
               case 21 -> dt = BOOLEAN;
               case 22 -> dt = NULL;
               case 23 -> dt = UNDEFINED;
               case 25, 26, 27 -> dt = FLOAT;
               case 31 -> dt = BREAK;
               default -> dt = SIMPLE;
            }
         }
         default -> throw ioex("Invalid major type: %s", majorType);
      }

      lastDataType = dt;
      indefinite = false;
      if (dt == FLOAT) {
         if (additionalInfo == 25)
            length = 2;
         else if (additionalInfo == 26)
            length = 4;
         else
            length = 8;
      } else if ((dt == ARRAY || dt == MAP || dt == STRING || dt == BINARY) && additionalInfo == 31) {
         // Indefinite-length container or string (RFC 8949 §3.2).  Flagged via 'indefinite' rather than
         // overloading length so callers distinguish it from a definite item whose 8-byte argument is
         // 0xFFFFFFFFFFFFFFFF.  Containers are driven by the BREAK-aware loop (shouldContinueContainer /
         // peekBreak); indefinite text/byte strings are reassembled from their chunks in readBinary.
         indefinite = true;
         length = 0;
      } else if (dt != BOOLEAN && dt != NULL && dt != UNDEFINED && dt != BREAK) {
         // NOTE: element-COUNT caps for arrays/maps are intentionally NOT applied here.  This method is
         // shared by the O(1)-memory streaming cursor (CborTokenReader), which must accept large
         // definite-length containers.  The count cap is enforced only on the databind path
         // (CborParserSession), where elements actually accumulate into a collection.  Byte/text-string
         // LENGTH caps remain active on every path via readBinary/readIndefiniteString.
         length = readArgument(additionalInfo);
      } else
         length = dt == BOOLEAN ? -1 : 0;

      return dt;
   }

   /**
    * Returns whether the last-read data item used indefinite-length encoding (additional-info 31).
    *
    * @return <jk>true</jk> if the last {@link #readDataType()} returned an indefinite-length array, map, or string.
    */
   boolean isIndefinite() {
      return indefinite;
   }

   /**
    * Returns the length/argument for the last read data type.
    * For strings/binary: byte count. For arrays: element count. For maps: pair count. For TAG: tag number.
    */
   long readLength() {
      return length;
   }

   /**
    * Reads a boolean (must have been identified as BOOLEAN by readDataType).
    */
   boolean readBoolean() {
      return lastInitialByte == 0xF5;
   }

   /**
    * Reads an unsigned integer value (UINT) as long.
    * The value was already consumed by readDataType/readArgument.
    */
   long readUnsignedLong() {
      return length;
   }

   /**
    * Reads a signed integer. For UINT, returns the value. For NINT, returns -1 - value.
    */
   long readSignedLong() {
      if (lastDataType == UINT)
         return readUnsignedLong();
      if (lastDataType == NINT)
         return -1 - readUnsignedLong();
      throw isex("Expected integer type, got %s", lastDataType);
   }

   /**
    * Reads the last integer (UINT/NINT) as a {@link BigInteger}, preserving the full unsigned
    * 64-bit magnitude.
    *
    * <p>
    * Used when the target Java type is {@link BigInteger} so that values beyond the signed
    * {@code long} range (CBOR major type 0 values {@code >= 2^63}, and major type 1 values
    * {@code < -2^63}) widen losslessly instead of surfacing as a sign-wrapped {@code long}.
    *
    * @return The integer value as a {@link BigInteger}.
    */
   BigInteger readBigInteger() {
      var magnitude = new BigInteger(Long.toUnsignedString(length));
      if (lastDataType == UINT)
         return magnitude;
      if (lastDataType == NINT)
         return BigInteger.valueOf(-1).subtract(magnitude);
      throw isex("Expected integer type, got %s", lastDataType);
   }

   /**
    * Reads an int (for UINT/NINT when value fits in int).
    */
   int readInt() {
      return (int)readSignedLong();
   }

   /**
    * Reads a long (for UINT/NINT).
    */
   long readLong() {
      return readSignedLong();
   }

   /**
    * Reads a float. Handles half (0xF9), single (0xFA), and double (0xFB) precision.
    */
   float readFloat() throws IOException {
      if (lastDataType == FLOAT) {
         if (lastInitialByte == 0xF9) {
            int half = (readUInt1() << 8) | readUInt1();
            return halfFloatToFloat(half);
         }
         if (lastInitialByte == 0xFA)
            return Float.intBitsToFloat((int)readUInt4());
         if (lastInitialByte == 0xFB)
            return (float)Double.longBitsToDouble(readUInt8());
      }
      throw isex("Expected float type, got %s", lastDataType);
   }

   /**
    * Reads a double. Handles half, single, and double precision.
    */
   double readDouble() throws IOException {
      if (lastDataType == FLOAT) {
         if (lastInitialByte == 0xF9)
            return halfFloatToFloat((readUInt1() << 8) | readUInt1());
         if (lastInitialByte == 0xFA)
            return Float.intBitsToFloat((int)readUInt4());
         if (lastInitialByte == 0xFB)
            return Double.longBitsToDouble(readUInt8());
      }
      throw isex("Expected float type, got %s", lastDataType);
   }

   private static float halfFloatToFloat(int half) {
      int exp = (half >> 10) & 0x1F;
      int mant = half & 0x3FF;
      int sign = (half & 0x8000) << 16;
      if (exp == 0) {
         // Zero (mant == 0) or subnormal (mant != 0).  A half subnormal is mant * 2^-24; the
         // previous (mant << 13) form reinterpreted the bits as a float32 subnormal (mant * 2^-149),
         // which is effectively zero.  Compute mant * 2^-24 directly (exact for mant <= 1023) and
         // re-apply the sign bit.
         var v = mant / 16777216f;  // 16777216 == 2^24
         return Float.intBitsToFloat(Float.floatToRawIntBits(v) | sign);
      }
      if (exp == 31)
         return mant == 0 ? Float.intBitsToFloat(sign | 0x7F800000)
            : Float.NaN;
      return Float.intBitsToFloat(sign | ((exp + 112) << 23) | (mant << 13));
   }

   /**
    * Reads a UTF-8 text string.
    */
   String readString() throws IOException {
      byte[] b = readBinary();
      return new String(b, UTF8);
   }

   /**
    * Reads a byte string (or text string's raw bytes).
    *
    * <p>
    * Handles both definite-length strings and indefinite-length strings (RFC 8949 §3.2.3): when the
    * head carried additional-info 31 ({@link #isIndefinite()}) the chunked definite-length pieces are
    * reassembled until the BREAK marker.  A definite-length string whose declared length is out of
    * bounds (negative, e.g. an 8-byte 0xFFFFFFFFFFFFFFFF argument, or beyond the configured maximum)
    * is rejected here via {@link #checkLength(long, String)}.
    */
   byte[] readBinary() throws IOException {
      if (indefinite)
         return readIndefiniteString();
      var b = new byte[checkLength(length, WHAT_STRING)];
      var bytesRead = read(b);
      if (bytesRead != b.length)
         throw ioex("Expected to read %s bytes but only read %s", b.length, bytesRead);
      return b;
   }

   /**
    * Reassembles an indefinite-length text/byte string (RFC 8949 §3.2.3) from its definite-length
    * chunks, terminated by the BREAK marker (0xFF).
    *
    * <p>
    * Per §3.2.3 every chunk must be a definite-length string of the same major type as the
    * enclosing indefinite string, and nested indefinite chunks are not permitted.  Malformed input
    * surfaces as a clean {@link IOException} rather than a {@link NegativeArraySizeException}.
    */
   private byte[] readIndefiniteString() throws IOException {
      var expectedMajorType = getMajorType(lastInitialByte);
      var baos = new ByteArrayOutputStream();
      var total = 0L;
      while (true) {
         int ib = read();
         if (ib == -1)
            throw ioex("Unexpected end of CBOR input in indefinite-length string");
         ib &= 0xFF;
         if (ib == 0xFF)  // BREAK
            break;
         var chunkMajorType = getMajorType(ib);
         var chunkInfo = getAdditionalInfo(ib);
         if (chunkMajorType != expectedMajorType)
            throw ioex("Invalid chunk major type %s in indefinite-length string (expected %s)", chunkMajorType, expectedMajorType);
         if (chunkInfo == 31)
            throw ioex("Nested indefinite-length string chunk not allowed");
         var chunkLen = checkLength(readArgument(chunkInfo), WHAT_STRING);
         // Bound the reassembled aggregate, not just each chunk, so the returned byte[] cannot exceed
         // the configured maximum regardless of how the payload is split across chunks.
         total += chunkLen;
         checkLength(total, WHAT_STRING);
         var chunk = new byte[chunkLen];
         var bytesRead = read(chunk);
         if (bytesRead != chunk.length)
            throw ioex("Expected to read %s bytes but only read %s", chunk.length, bytesRead);
         baos.write(chunk, 0, chunk.length);
      }
      return baos.toByteArray();
   }
}