/*
 * Copyright 2016 Google Inc. All Rights Reserved.
 *
 * Licensed 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 com.google.devrel.gmscore.tools.apk.arsc;

import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableList.Builder;
import com.google.common.io.LittleEndianDataOutputStream;

import java.io.ByteArrayOutputStream;
import java.io.DataOutput;
import java.io.IOException;
import java.nio.ByteBuffer;
import java.nio.ByteOrder;
import java.util.*;

import javax.annotation.Nullable;

/** Represents a string pool structure. */
public final class StringPoolChunk extends Chunk {

  // These are the defined flags for the "flags" field of ResourceStringPoolHeader
  private static final int SORTED_FLAG = 1 << 0;
  private static final int UTF8_FLAG   = 1 << 8;

  /** The offset from the start of the header that the stylesStart field is at. */
  private static final int STYLE_START_OFFSET = 24;

  /** Flags. */
  private final int flags;

  /** Index from header of the string data. */
  private final int stringsStart;

  /** Index from header of the style data. */
  private final int stylesStart;

  /**
   * Number of strings in the original buffer. This is not necessarily the number of strings in
   * {@code strings}.
   */
  private final int stringCount;

  /**
   * Number of styles in the original buffer. This is not necessarily the number of styles in
   * {@code styles}.
   */
  private final int styleCount;

  /**
   * The strings ordered as they appear in the arsc file. e.g. strings.get(1234) gets the 1235th
   * string in the arsc file.
   */
  private final List<String> strings = new ArrayList<>();

  /**
   * These styles have a 1:1 relationship with the strings. For example, styles.get(3) refers to
   * the string at location strings.get(3). There are never more styles than strings (though there
   * may be less). Inside of that are all of the styles referenced by that string.
   */
  private final List<StringPoolStyle> styles = new ArrayList<>();

  /**
   * True if the original {@link StringPoolChunk} shows signs of being deduped. Specifically, this
   * is set to true if there exists a string whose offset is <= the previous offset. This is used to
   * preserve the deduping of strings for pools that have been deduped.
   */
  private boolean isOriginalDeduped = false;

  protected StringPoolChunk(ByteBuffer buffer, @Nullable Chunk parent) {
    super(buffer, parent);
    stringCount = buffer.getInt();
    styleCount = buffer.getInt();
    flags        = buffer.getInt();
    stringsStart = buffer.getInt();
    stylesStart  = buffer.getInt();
  }

  @Override
  protected void init(ByteBuffer buffer) {
    super.init(buffer);
    strings.addAll(readStrings(buffer, offset + stringsStart, stringCount));
    styles.addAll(readStyles(buffer, offset + stylesStart, styleCount));
  }

  /**
   * Returns the 0-based index of the first occurrence of the given string, or -1 if the string is
   * not in the pool. This runs in O(n) time.
   *
   * @param string The string to check the pool for.
   * @return Index of the string, or -1 if not found.
   */
  public int indexOf(String string) {
    return strings.indexOf(string);
  }

  /**
   * Returns a string at the given (0-based) index.
   *
   * @param index The (0-based) index of the string to return.
   * @throws IndexOutOfBoundsException If the index is out of range (index < 0 || index >= size()).
   */
  public String getString(int index) {
    return strings.get(index);
  }

  /** Returns the number of strings in this pool. */
  public int getStringCount() {
    return strings.size();
  }

  /**
   * Returns a style at the given (0-based) index.
   *
   * @param index The (0-based) index of the style to return.
   * @throws IndexOutOfBoundsException If the index is out of range (index < 0 || index >= size()).
   */
  public StringPoolStyle getStyle(int index) {
    return styles.get(index);
  }

  /** Returns the number of styles in this pool. */
  public int getStyleCount() {
    return styles.size();
  }

  /** Returns the type of strings in this pool. */
  public BinaryResourceString.Type getStringType() {
    return isUTF8() ? BinaryResourceString.Type.UTF8 : BinaryResourceString.Type.UTF16;
  }

  @Override
  protected Type getType() {
    return Chunk.Type.STRING_POOL;
  }

  /** Returns the number of bytes needed for offsets based on {@code strings} and {@code styles}. */
  private int getOffsetSize() {
    return (strings.size() + styles.size()) * 4;
  }

  /**
   * True if this string pool contains strings in UTF-8 format. Otherwise, strings are in UTF-16.
   *
   * @return true if @{code strings} are in UTF-8; false if they're in UTF-16.
   */
  public boolean isUTF8() {
    return (flags & UTF8_FLAG) != 0;
  }

  /**
   * True if this string pool contains already-sorted strings.
   *
   * @return true if @{code strings} are sorted.
   */
  public boolean isSorted() {
    return (flags & SORTED_FLAG) != 0;
  }

  private List<String> readStrings(ByteBuffer buffer, int offset, int count) {
    List<String> result = new ArrayList<>();
    int previousOffset = -1;
    // After the header, we now have an array of offsets for the strings in this pool.
    for (int i = 0; i < count; ++i) {
      int stringOffset = offset + buffer.getInt();
      result.add(BinaryResourceString.decodeString(buffer, stringOffset, getStringType()));
      if (stringOffset <= previousOffset) {
        isOriginalDeduped = true;
      }
      previousOffset = stringOffset;
    }
    return result;
  }

  private List<StringPoolStyle> readStyles(ByteBuffer buffer, int offset, int count) {
    List<StringPoolStyle> result = new ArrayList<>();
    // After the array of offsets for the strings in the pool, we have an offset for the styles
    // in this pool.
    for (int i = 0; i < count; ++i) {
      int styleOffset = offset + buffer.getInt();
      result.add(StringPoolStyle.create(buffer, styleOffset, this));
    }
    return result;
  }

  private int writeStrings(DataOutput payload, ByteBuffer offsets, boolean shrink)
      throws IOException {
    int stringOffset = 0;
    Map<String, Integer> used = new HashMap<>();  // Keeps track of strings already written
    for (String string : strings) {
      // Dedupe everything except stylized strings, unless shrink is true (then dedupe everything)
      if (used.containsKey(string) && (shrink || isOriginalDeduped)) {
        Integer offset = used.get(string);
        offsets.putInt(offset == null ? 0 : offset);
      } else {
        byte[] encodedString = BinaryResourceString.encodeString(string, getStringType());
        payload.write(encodedString);
        used.put(string, stringOffset);
        offsets.putInt(stringOffset);
        stringOffset += encodedString.length;
      }
    }

    // ARSC files pad to a 4-byte boundary. We should do so too.
    stringOffset = writePad(payload, stringOffset);
    return stringOffset;
  }

  private int writeStyles(DataOutput payload, ByteBuffer offsets, boolean shrink)
      throws IOException {
    int styleOffset = 0;
    if (styles.size() > 0) {
      Map<StringPoolStyle, Integer> used = new HashMap<>();  // Keeps track of bytes already written
      for (StringPoolStyle style : styles) {
        if (!used.containsKey(style) || !shrink) {
          byte[] encodedStyle = style.toByteArray(shrink);
          payload.write(encodedStyle);
          used.put(style, styleOffset);
          offsets.putInt(styleOffset);
          styleOffset += encodedStyle.length;
        } else {  // contains key and shrink is true
          Integer offset = used.get(style);
          offsets.putInt(offset == null ? 0 : offset);
        }
      }
      // The end of the spans are terminated with another sentinel value
      payload.writeInt(StringPoolStyle.RES_STRING_POOL_SPAN_END);
      styleOffset += 4;
      // TODO(acornwall): There appears to be an extra SPAN_END here... why?
      payload.writeInt(StringPoolStyle.RES_STRING_POOL_SPAN_END);
      styleOffset += 4;

      styleOffset = writePad(payload, styleOffset);
    }
    return styleOffset;
  }

  @Override
  protected void writeHeader(ByteBuffer output) {
    int stringsStart = getHeaderSize() + getOffsetSize();
    output.putInt(strings.size());
    output.putInt(styles.size());
    output.putInt(flags);
    output.putInt(strings.isEmpty() ? 0 : stringsStart);
    output.putInt(0);  // Placeholder. The styles starting offset cannot be computed at this point.
  }

  @Override
  protected void writePayload(DataOutput output, ByteBuffer header, boolean shrink)
      throws IOException {
    ByteArrayOutputStream baos = new ByteArrayOutputStream();
    int stringOffset = 0;
    ByteBuffer offsets = ByteBuffer.allocate(getOffsetSize());
    offsets.order(ByteOrder.LITTLE_ENDIAN);

    // Write to a temporary payload so we can rearrange this and put the offsets first
    try (LittleEndianDataOutputStream payload = new LittleEndianDataOutputStream(baos)) {
      stringOffset = writeStrings(payload, offsets, shrink);
      writeStyles(payload, offsets, shrink);
    }

    output.write(offsets.array());
    output.write(baos.toByteArray());
    if (!styles.isEmpty()) {
      header.putInt(STYLE_START_OFFSET, getHeaderSize() + getOffsetSize() + stringOffset);
    }
  }

  /**
   * Represents all of the styles for a particular string. The string is determined by its index
   * in {@link StringPoolChunk}.
   */
  public static class StringPoolStyle implements SerializableResource {

    // Styles are a list of integers with 0xFFFFFFFF serving as a sentinel value.
    static final int RES_STRING_POOL_SPAN_END = 0xFFFFFFFF;
    private final List<StringPoolSpan> spans;

    static StringPoolStyle create(ByteBuffer buffer, int offset, StringPoolChunk parent) {
      Builder<StringPoolSpan> spans = ImmutableList.builder();
      int nameIndex = buffer.getInt(offset);
      while (nameIndex != RES_STRING_POOL_SPAN_END) {
        spans.add(StringPoolSpan.create(buffer, offset, parent));
        offset += StringPoolSpan.SPAN_LENGTH;
        nameIndex = buffer.getInt(offset);
      }
      return new StringPoolStyle(spans.build());
    }

    private StringPoolStyle(List<StringPoolSpan> spans) {
      this.spans = spans;
    }

    @Override
    public byte[] toByteArray() throws IOException {
      return toByteArray(false);
    }

    @Override
    public byte[] toByteArray(boolean shrink) throws IOException {
      ByteArrayOutputStream baos = new ByteArrayOutputStream();

      try (LittleEndianDataOutputStream payload = new LittleEndianDataOutputStream(baos)) {
        for (StringPoolSpan span : spans) {
          byte[] encodedSpan = span.toByteArray(shrink);
          if (encodedSpan.length != StringPoolSpan.SPAN_LENGTH) {
            throw new IllegalStateException("Encountered a span of invalid length.");
          }
          payload.write(encodedSpan);
        }
        payload.writeInt(RES_STRING_POOL_SPAN_END);
      }

      return baos.toByteArray();
    }

    @Override
    public boolean equals(Object o) {
      if (this == o) return true;
      if (o == null || getClass() != o.getClass()) return false;
      StringPoolStyle that = (StringPoolStyle)o;
      return Objects.equals(spans, that.spans);
    }

    @Override
    public int hashCode() {
      return Objects.hash(spans);
    }

    /**
     * Returns a brief description of the contents of this style. The representation of this
     * information is subject to change, but below is a typical example:
     *
     * <pre>"StringPoolStyle{spans=[StringPoolSpan{foo, start=0, stop=5}, ...]}"</pre>
     */
    @Override
    public String toString() {
      return String.format("StringPoolStyle{spans=%s}", spans);
    }
  }

  /** Represents a styled span associated with a specific string. */
  private static class StringPoolSpan implements SerializableResource {
    static final int SPAN_LENGTH = 12;

    private final int nameIndex;
    private final int start;
    private final int stop;
    private final StringPoolChunk parent;

    static StringPoolSpan create(ByteBuffer buffer, int offset, StringPoolChunk parent) {
      int nameIndex = buffer.getInt(offset);
      int start = buffer.getInt(offset + 4);
      int stop = buffer.getInt(offset + 8);
      return new StringPoolSpan(nameIndex, start, stop, parent);
    }

    private StringPoolSpan(int nameIndex, int start, int stop, StringPoolChunk parent) {
      this.nameIndex = nameIndex;
      this.start = start;
      this.stop = stop;
      this.parent = parent;
    }

    @Override
    public final byte[] toByteArray() {
      return toByteArray(false);
    }

    @Override
    public final byte[] toByteArray(boolean shrink) {
      ByteBuffer buffer = ByteBuffer.allocate(SPAN_LENGTH).order(ByteOrder.LITTLE_ENDIAN);
      buffer.putInt(nameIndex);
      buffer.putInt(start);
      buffer.putInt(stop);
      return buffer.array();
    }

    @Override
    public boolean equals(Object o) {
      if (this == o) return true;
      if (o == null || getClass() != o.getClass()) return false;
      StringPoolSpan that = (StringPoolSpan)o;
      return nameIndex == that.nameIndex &&
             start == that.start &&
             stop == that.stop &&
             Objects.equals(parent, that.parent);
    }

    @Override
    public int hashCode() {
      return Objects.hash(nameIndex, start, stop, parent);
    }

    /**
     * Returns a brief description of this span. The representation of this information is subject
     * to change, but below is a typical example:
     *
     * <pre>"StringPoolSpan{foo, start=0, stop=5}"</pre>
     */
    @Override
    public String toString() {
      return String.format("StringPoolSpan{%s, start=%d, stop=%d}",
          parent.getString(nameIndex), start, stop);
    }
  }
}
