content: /**
   * 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.solr.response;
  
  import org.apache.commons.csv.CSVPrinter;
  import org.apache.commons.csv.CSVStrategy;
  import org.apache.lucene.document.Document;
  import org.apache.lucene.document.Fieldable;
  import org.apache.solr.common.SolrDocument;
  import org.apache.solr.common.SolrDocumentList;
  import org.apache.solr.common.SolrException;
  import org.apache.solr.common.params.SolrParams;
  import org.apache.solr.common.util.DateUtil;
  import org.apache.solr.common.util.FastWriter;
  import org.apache.solr.common.util.NamedList;
  import org.apache.solr.request.SolrQueryRequest;
  import org.apache.solr.request.uninverted.UnInvertedField;
  import org.apache.solr.schema.FieldType;
  import org.apache.solr.schema.SchemaField;
  import org.apache.solr.schema.StrField;
  import org.apache.solr.search.DocIterator;
  import org.apache.solr.search.DocList;
  import org.apache.solr.search.SolrIndexSearcher;
  
  import java.io.CharArrayWriter;
  import java.io.IOException;
  import java.io.Writer;
  import java.util.*;
  
  /**
   * @version $Id: CSVResponseWriter.java 979838 2010-07-27 20:13:15Z yonik $
   */
  
  public class CSVResponseWriter implements QueryResponseWriter {
  
    public void init(NamedList n) {
    }
  
    public void write(Writer writer, SolrQueryRequest req, SolrQueryResponse rsp) throws IOException {
      CSVWriter w = new CSVWriter(writer, req, rsp);
      try {
        w.writeResponse();
      } finally {
        w.close();
      }
    }
  
    public String getContentType(SolrQueryRequest request, SolrQueryResponse response) {
      // using the text/plain allows this to be viewed in the browser easily
      return CONTENT_TYPE_TEXT_UTF8;
    }
  }
  
  
  class CSVWriter extends TextResponseWriter {
    static String SEPARATOR = "separator";
    static String ENCAPSULATOR = "encapsulator";
    static String ESCAPE = "escape";
  
    static String CSV = "csv.";
    static String CSV_SEPARATOR = CSV + SEPARATOR;
    static String CSV_ENCAPSULATOR = CSV + ENCAPSULATOR;
    static String CSV_ESCAPE = CSV + ESCAPE;
  
    static String MV = CSV+"mv.";
    static String MV_SEPARATOR = MV + SEPARATOR;
    static String MV_ENCAPSULATOR = MV + ENCAPSULATOR;
    static String MV_ESCAPE = MV + ESCAPE;
  
    static String CSV_NULL = CSV + "null";
    static String CSV_HEADER = CSV + "header";
    static String CSV_NEWLINE = CSV + "newline";
  
    char[] sharedCSVBuf = new char[8192];
  
    // prevent each instance from creating it's own buffer
    class CSVSharedBufPrinter extends CSVPrinter {
      public CSVSharedBufPrinter(Writer out, CSVStrategy strategy) {
        super(out, strategy);
        super.buf = sharedCSVBuf;
      }
  
      public void reset() {
        super.newLine = true;
        // update our shared buf in case a new bigger one was allocated
        sharedCSVBuf = super.buf;
      }
    }
  
    // allows access to internal buf w/o copying it
    static class OpenCharArrayWriter extends CharArrayWriter {
      public char[]  getInternalBuf() { return buf; }
    }
  
    // Writes all data to a char array,
    // allows access to internal buffer, and allows fast resetting.
    static class ResettableFastWriter extends FastWriter {
      OpenCharArrayWriter cw = new OpenCharArrayWriter();
      char[] result;
      int resultLen;
  
      public ResettableFastWriter() {
        super(new OpenCharArrayWriter());
        cw = (OpenCharArrayWriter)sink;
      }
  
      public void reset() {
        cw.reset();
        pos=0;
      }
  
      public void freeze() throws IOException {
        if (cw.size() > 0) {
          flush();
          result = cw.getInternalBuf();
          resultLen = cw.size();
        } else {
          result = buf;
          resultLen = pos;
        }
      }
  
      public int getFrozenSize() { return resultLen; }
      public char[] getFrozenBuf() { return result; }
    }
  
  
    static class CSVField {
      String name;
      SchemaField sf;
      CSVSharedBufPrinter mvPrinter;  // printer used to encode multiple values in a single CSV value
  
      // used to collect values
      List<Fieldable> values = new ArrayList<Fieldable>(1);  // low starting amount in case there are many fields
      int tmp;
    }
  
    int pass;
    Map<String,CSVField> csvFields = new LinkedHashMap<String,CSVField>();
  
    Calendar cal;  // for formatting date objects
  
    CSVStrategy strategy;  // strategy for encoding the fields of documents
    CSVPrinter printer;
    ResettableFastWriter mvWriter = new ResettableFastWriter();  // writer used for multi-valued fields
  
    String NullValue;
    boolean returnScore = false;
  
  
    public CSVWriter(Writer writer, SolrQueryRequest req, SolrQueryResponse rsp) {
      super(writer, req, rsp);
    }
  
    public void writeResponse() throws IOException {
      SolrParams params = req.getParams();
  
      strategy = new CSVStrategy(',', '"', CSVStrategy.COMMENTS_DISABLED, CSVStrategy.ESCAPE_DISABLED, false, false, false, true);
      CSVStrategy strat = strategy;
  
      String sep = params.get(CSV_SEPARATOR);
      if (sep!=null) {
        if (sep.length()!=1) throw new SolrException( SolrException.ErrorCode.BAD_REQUEST,"Invalid separator:'"+sep+"'");
        strat.setDelimiter(sep.charAt(0));
      }
  
      String nl = params.get(CSV_NEWLINE);
      if (nl!=null) {
        if (nl.length()==0) throw new SolrException( SolrException.ErrorCode.BAD_REQUEST,"Invalid newline:'"+nl+"'");
        strat.setPrinterNewline(nl);
      }
  
      String encapsulator = params.get(CSV_ENCAPSULATOR);
      String escape = params.get(CSV_ESCAPE);
      if (encapsulator!=null) {
        if (encapsulator.length()!=1) throw new SolrException( SolrException.ErrorCode.BAD_REQUEST,"Invalid encapsulator:'"+encapsulator+"'");
        strat.setEncapsulator(encapsulator.charAt(0));
      }
  
      if (escape!=null) {
        if (escape.length()!=1) throw new SolrException( SolrException.ErrorCode.BAD_REQUEST,"Invalid escape:'"+escape+"'");
        strat.setEscape(escape.charAt(0));
        if (encapsulator == null) {
          strat.setEncapsulator( CSVStrategy.ENCAPSULATOR_DISABLED);
        }
      }
  
      if (strat.getEscape() == '\\') {
        // If the escape is the standard backslash, then also enable
        // unicode escapes (it's harmless since 'u' would not otherwise
        // be escaped.
        strat.setUnicodeEscapeInterpretation(true);
      }
  
      printer = new CSVPrinter(writer, strategy);
      
  
      CSVStrategy mvStrategy = new CSVStrategy(strategy.getDelimiter(), CSVStrategy.ENCAPSULATOR_DISABLED, CSVStrategy.COMMENTS_DISABLED, '\\', false, false, false, false);
      strat = mvStrategy;
  
      sep = params.get(MV_SEPARATOR);
      if (sep!=null) {
        if (sep.length()!=1) throw new SolrException( SolrException.ErrorCode.BAD_REQUEST,"Invalid mv separator:'"+sep+"'");
        strat.setDelimiter(sep.charAt(0));
      }
  
      encapsulator = params.get(MV_ENCAPSULATOR);
      escape = params.get(MV_ESCAPE);
  
      if (encapsulator!=null) {
        if (encapsulator.length()!=1) throw new SolrException( SolrException.ErrorCode.BAD_REQUEST,"Invalid mv encapsulator:'"+encapsulator+"'");
        strat.setEncapsulator(encapsulator.charAt(0));
        if (escape == null) {
          strat.setEscape(CSVStrategy.ESCAPE_DISABLED);
        }
      }
  
      escape = params.get(MV_ESCAPE);
      if (escape!=null) {
        if (escape.length()!=1) throw new SolrException( SolrException.ErrorCode.BAD_REQUEST,"Invalid mv escape:'"+escape+"'");
        strat.setEscape(escape.charAt(0));
        // encapsulator will already be disabled if it wasn't specified
      }
  
      returnScore = returnFields != null && returnFields.contains("score");
      boolean needListOfFields = returnFields==null || returnFields.size()==0 || (returnFields.size()==1 && returnScore) || returnFields.contains("*");
      Collection<String> fields = returnFields;
  
      Object responseObj = rsp.getValues().get("response");
      if (needListOfFields) {
        if (responseObj instanceof SolrDocumentList) {
          // get the list of fields from the SolrDocumentList
          fields = new LinkedHashSet<String>();
          for (SolrDocument sdoc: (SolrDocumentList)responseObj) {
            fields.addAll(sdoc.getFieldNames());
          }
        } else {
          // get the list of fields from the index
          fields = req.getSearcher().getFieldNames();
        }
        if (returnScore) {
          fields.add("score");
        } else {
          fields.remove("score");
        }
      }
  
      CSVSharedBufPrinter csvPrinterMV = new CSVSharedBufPrinter(mvWriter, mvStrategy);
  
      for (String field : fields) {
        if (field.equals("score")) {
          CSVField csvField = new CSVField();
          csvField.name = "score";
          csvFields.put("score", csvField);
          continue;
        }
  
        SchemaField sf = schema.getFieldOrNull(field);
        if (sf == null) {
          FieldType ft = new StrField();
          sf = new SchemaField(field, ft);
        }
  
        // if we got the list of fields from the index, only list stored fields
        if (returnFields==null && sf != null && !sf.stored()) {
          continue;
        }
  
        // check for per-field overrides
        sep = params.get("f." + field + '.' + CSV_SEPARATOR);
        encapsulator = params.get("f." + field + '.' + CSV_ENCAPSULATOR);
        escape = params.get("f." + field + '.' + CSV_ESCAPE);
  
        CSVSharedBufPrinter csvPrinter = csvPrinterMV;
        if (sep != null || encapsulator != null || escape != null) {
          // create a new strategy + printer if there were any per-field overrides
          strat = (CSVStrategy)mvStrategy.clone();
          if (sep!=null) {
            if (sep.length()!=1) throw new SolrException( SolrException.ErrorCode.BAD_REQUEST,"Invalid mv separator:'"+sep+"'");
            strat.setDelimiter(sep.charAt(0));
          }
          if (encapsulator!=null) {
            if (encapsulator.length()!=1) throw new SolrException( SolrException.ErrorCode.BAD_REQUEST,"Invalid mv encapsulator:'"+encapsulator+"'");
            strat.setEncapsulator(encapsulator.charAt(0));
            if (escape == null) {
              strat.setEscape(CSVStrategy.ESCAPE_DISABLED);
            }
          }
          if (escape!=null) {
            if (escape.length()!=1) throw new SolrException( SolrException.ErrorCode.BAD_REQUEST,"Invalid mv escape:'"+escape+"'");
            strat.setEscape(escape.charAt(0));
            if (encapsulator == null) {
              strat.setEncapsulator(CSVStrategy.ENCAPSULATOR_DISABLED);
            }
          }        
          csvPrinter = new CSVSharedBufPrinter(mvWriter, strat);
        }
  
  
        CSVField csvField = new CSVField();
        csvField.name = field;
        csvField.sf = sf;
        csvField.mvPrinter = csvPrinter;
        csvFields.put(field, csvField);
      }
  
      NullValue = params.get(CSV_NULL, "");
  
      if (params.getBool(CSV_HEADER, true)) {
        for (CSVField csvField : csvFields.values()) {
          printer.print(csvField.name);
        }
        printer.println();
      }
  
  
      if (responseObj instanceof DocList) {
        writeDocList(null, (DocList)responseObj, null, null);
      } else if (responseObj instanceof SolrDocumentList) {
        writeSolrDocumentList(null, (SolrDocumentList)responseObj, null, null);
      }
  
    }
  
    @Override
    public void close() throws IOException {
      if (printer != null) printer.flush();
      super.close();
    }
  
    @Override
    public void writeNamedList(String name, NamedList val) throws IOException {
    }
  
    @Override
    public void writeDoc(String name, Document doc, Set<String> returnFields, float score, boolean includeScore) throws IOException {
      pass++;
  
      for (Fieldable field: doc.getFields()) {
        CSVField csvField = csvFields.get(field.name());
        if (csvField == null) continue;
        if (csvField.tmp != pass) {
          csvField.tmp = pass;
          csvField.values.clear();
        }
        csvField.values.add(field);
      }
  
      for (CSVField csvField : csvFields.values()) {
        if (csvField.name.equals("score")) {
          writeFloat("score", score);
          continue;
        }
        if (csvField.tmp != pass) {
          writeNull(csvField.name);
          continue;
        }
  
        if (csvField.sf.multiValued() || csvField.values.size() > 1) {
          mvWriter.reset();
          csvField.mvPrinter.reset();
          // switch the printer to use the multi-valued one
          CSVPrinter tmp = printer;
          printer = csvField.mvPrinter;
          for (Fieldable fval : csvField.values) {
            csvField.sf.getType().write(this, csvField.name, fval);
          }
          printer = tmp;  // restore the original printer
  
          mvWriter.freeze();
          printer.print(mvWriter.getFrozenBuf(), 0, mvWriter.getFrozenSize(), true);
        } else {
          assert csvField.values.size() == 1;
          csvField.sf.getType().write(this,csvField.name,csvField.values.get(0));
        }
      }
  
      printer.println();
    }
  
    //NOTE: a document cannot currently contain another document
    List tmpList;
    @Override
    public void writeSolrDocument(String name, SolrDocument doc, Set<String> returnFields, Map pseudoFields) throws IOException {
      if (tmpList == null) {
        tmpList = new ArrayList(1);
        tmpList.add(null);
      }
  
      for (CSVField csvField : csvFields.values()) {
        Object val = doc.getFieldValue(csvField.name);
        int nVals = val instanceof Collection ? ((Collection)val).size() : (val==null ? 0 : 1);
        if (nVals == 0) {
          writeNull(csvField.name);
          continue;
        }
  
        if ((csvField.sf != null && csvField.sf.multiValued()) || nVals > 1) {
          Collection values;
          // normalize to a collection
          if (val instanceof Collection) {
            values = (Collection)val;
          } else {
            tmpList.set(0, val);
            values = tmpList;
          }
  
          mvWriter.reset();
          csvField.mvPrinter.reset();
          // switch the printer to use the multi-valued one
          CSVPrinter tmp = printer;
          printer = csvField.mvPrinter;
          for (Object fval : values) {
            writeVal(csvField.name, fval);
          }
          printer = tmp;  // restore the original printer
  
          mvWriter.freeze();
          printer.print(mvWriter.getFrozenBuf(), 0, mvWriter.getFrozenSize(), true);
  
        } else {
          // normalize to first value
          if (val instanceof Collection) {
            Collection values = (Collection)val;
            val = values.iterator().next();
          }
          writeVal(csvField.name, val);
        }
      }
  
      printer.println();
    }
  
    @Override
    public void writeDocList(String name, DocList ids, Set<String> fields, Map otherFields) throws IOException {
      int sz=ids.size();
      SolrIndexSearcher searcher = req.getSearcher();
      DocIterator iterator = ids.iterator();
      for (int i=0; i<sz; i++) {
        int id = iterator.nextDoc();
        Document doc = searcher.doc(id, fields);
        writeDoc(null, doc, fields, (returnScore ? iterator.score() : 0.0f), returnScore);
      }
    }
  
    Map scoreMap = new HashMap(1);
    @Override
    public void writeSolrDocumentList(String name, SolrDocumentList docs, Set<String> fields, Map otherFields) throws IOException {
      for (SolrDocument doc : docs) {
        writeSolrDocument(name, doc, fields, otherFields);
      }
    }
  
    @Override
    public void writeStr(String name, String val, boolean needsEscaping) throws IOException {
      printer.print(val, needsEscaping);
    }
  
    @Override
    public void writeMap(String name, Map val, boolean excludeOuter, boolean isFirstVal) throws IOException {
    }
  
    @Override
    public void writeArray(String name, Object[] val) throws IOException {
    }
  
    @Override
    public void writeArray(String name, Iterator val) throws IOException {
    }
  
    @Override
    public void writeNull(String name) throws IOException {
      printer.print(NullValue);
    }
  
    @Override
    public void writeInt(String name, String val) throws IOException {
      printer.print(val, false);
    }
  
    @Override
    public void writeLong(String name, String val) throws IOException {
      printer.print(val, false);
    }
  
    @Override
    public void writeBool(String name, String val) throws IOException {
      printer.print(val, false);
    }
  
    @Override
    public void writeFloat(String name, String val) throws IOException {
      printer.print(val, false);
    }
  
    @Override
    public void writeDouble(String name, String val) throws IOException {
      printer.print(val, false);
    }
  
    @Override
    public void writeDate(String name, Date val) throws IOException {
      StringBuilder sb = new StringBuilder(25);
      cal = DateUtil.formatDate(val, cal, sb);
      writeDate(name, sb.toString());
    }
  
    @Override
    public void writeDate(String name, String val) throws IOException {
      printer.print(val, false);
    }
  
    @Override
    public void writeShort(String name, String val) throws IOException {
      printer.print(val, false);
    }
  
    @Override
    public void writeByte(String name, String val) throws IOException {
      printer.print(val, false);
    }
  }
  