/**
    * HBaseWriter
    *
    * $Id$
    *
    * Created on June 23rd, 2007
    *
    * Copyright (C) 2007 stack
    *
   * This file is part of the Heritrix web crawler (crawler.archive.org).
   *
   * Heritrix is free software; you can redistribute it and/or modify
   * it under the terms of the GNU Lesser Public License as published by
   * the Free Software Foundation; either version 2.1 of the License, or
   * any later version.
   *
   * Heritrix 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 Lesser Public License for more details.
   *
   * You should have received a copy of the GNU Lesser Public License
   * along with Heritrix; if not, write to the Free Software
   * Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA
   */
  package com.powerset.heritrix.writer;
  
  import java.io.ByteArrayOutputStream;
  import java.io.IOException;
  
  import org.apache.hadoop.hbase.HBaseConfiguration;
  import org.apache.hadoop.hbase.HColumnDescriptor;
  import org.apache.hadoop.hbase.HConstants;
  import org.apache.hadoop.hbase.HTableDescriptor;
  import org.apache.hadoop.hbase.client.HBaseAdmin;
  import org.apache.hadoop.hbase.client.HTable;
  import org.apache.hadoop.hbase.io.BatchUpdate;
  import org.apache.hadoop.hbase.util.Bytes;
  import org.apache.hadoop.hbase.util.Keying;
  import org.apache.log4j.Logger;
  import org.archive.io.ArchiveFileConstants;
  import org.archive.io.RecordingInputStream;
  import org.archive.io.RecordingOutputStream;
  import org.archive.io.ReplayInputStream;
  import org.archive.io.WriterPoolMember;
  import org.archive.modules.ProcessorURI;
  
  /**
   * Write to HBase.
   * Puts content into the 'content:' column and all else into the 'curi:'
   * column family.  Makes a row key of an url transformation.  Creates table
   * if it does not exist.
   * 
   * <p>Limitations: Hard-coded table schema.
   * @author stack
   */
  public class HBaseWriter extends WriterPoolMember implements ArchiveFileConstants {
    private final Logger LOG = Logger.getLogger(this.getClass().getName());
    private final HTable client;
    // TODO: make this variable configurable in the heritrix sheet: CONTENT_COLUMN_FAMILY
    public static final String CONTENT_COLUMN_FAMILY = "content:";
    // TODO: make this variable configurable in the heritrix sheet: CONTENT_COLUMN
    public static final String CONTENT_COLUMN = CONTENT_COLUMN_FAMILY + "raw_data";
    // TODO: make this variable configurable in the heritrix sheet: CURI_COLUMN_FAMILY
    public static final String CURI_COLUMN_FAMILY = "curi:";
    
    private static final String IP_COLUMN = CURI_COLUMN_FAMILY + "ip";
    private static final String PATH_FROM_SEED_COLUMN =
      CURI_COLUMN_FAMILY + "path-from-seed";
    private static final String IS_SEED_COLUMN = CURI_COLUMN_FAMILY + "is-seed";
    private static final String VIA_COLUMN = CURI_COLUMN_FAMILY + "via";
    private static final String URL_COLUMN = CURI_COLUMN_FAMILY + "url";
    private static final String REQUEST_COLUMN = CURI_COLUMN_FAMILY + "request";
  
    public HBaseWriter(final String master, final String table)
    throws IOException {
      super(null, null, null, false, null);
      if (table == null || table.length() <= 0) {
        throw new IllegalArgumentException("Must specify a table name");
      }
      HBaseConfiguration c = new HBaseConfiguration();
      if (master != null && master.length() > 0) {
        c.set(HConstants.MASTER_ADDRESS, master);
      }
      createCrawlTable(c, table);
      this.client = new HTable(c, table);
    }
  
    protected void createCrawlTable(final HBaseConfiguration c,
        final String table)
    throws IOException {
      HBaseAdmin admin = new HBaseAdmin(c);
      if (admin.tableExists(table)) {
        return;
      }
      HTableDescriptor htd = new HTableDescriptor(table);
      htd.addFamily(new HColumnDescriptor(CONTENT_COLUMN_FAMILY));
      htd.addFamily(new HColumnDescriptor(CURI_COLUMN_FAMILY));
      admin.createTable(htd);
     LOG.info("Created table " + htd.toString());
   }
 
   /**
    * @param curi URI of crawled document
    * @param ip IP of remote machine.
    * @param ros recording input stream that captured the response
    * @param ris recording output stream that captured the GET request
    */
   public void write(final ProcessorURI curi, final String ip,
       final RecordingOutputStream ros, final RecordingInputStream ris)
   throws IOException {
     String url = curi.toString();
     String row = Keying.createKey(url);
     if (LOG.isTraceEnabled()) {
       LOG.trace("Writing " + url + " as " + row.toString());
     }
     BatchUpdate bu = new BatchUpdate(row);
     bu.put(URL_COLUMN, Bytes.toBytes(url));
     bu.put(IP_COLUMN, Bytes.toBytes(ip));
     if (curi.isSeed()) {
       // TODO: Make Bytes.toBytes that takes a boolean.
       bu.put(IS_SEED_COLUMN, Bytes.toBytes(Boolean.TRUE.toString()));
       if (curi.getPathFromSeed() != null
           && curi.getPathFromSeed().trim().length() > 0) {
         bu.put(PATH_FROM_SEED_COLUMN, Bytes.toBytes(curi.getPathFromSeed().trim()));
       }
     }
     String viaStr = (curi.getVia() != null)?
       curi.getVia().toString().trim(): null;
       if (viaStr != null && viaStr.length() > 0) {
         bu.put(VIA_COLUMN, Bytes.toBytes(viaStr));
       }
       // Request
       if (ros.getSize() > 0) {
         add(bu, REQUEST_COLUMN, ros.getReplayInputStream(), (int)ros.getSize());
       }
       // Response
       add(bu, CONTENT_COLUMN, ris.getReplayInputStream(),
         (int)ris.getSize());
       // process the content (optional)
       processContent(bu);
       // Set crawl time.
       bu.setTimestamp(curi.getFetchBeginTime());
       this.client.commit(bu);
   }
   
   /**
    *  This is a stub method and is here to allow extension/overriding for 
    *  custom content parsing, data manipulation and to populate new columns.
    *  
    *  For Example : html parsing, text extraction, analysis and transformation and 
    *  storing the results in new column families/columns using the batch update object.
    *  
    * @param bu
    */
   protected void processContent(BatchUpdate bu) {
 	  //byte[] content = bu.get(CONTENT_COLUMN);
 	  // process content.....
 	  //bu.put("some:new_column", someParsedByteArray);
   }
   
   /*
    * Add ReplayInputStream to the passed BatchUpdate.
    * @param bu
    * @param key
    * @param ris
    * @param baos
    * @throws IOException
    */
   private void add(final BatchUpdate bu, final String key,
     final ReplayInputStream ris, final int size)
   throws IOException {
     ByteArrayOutputStream baos = new ByteArrayOutputStream(size);
     try {
       ris.readFullyTo(baos);
     } finally {
       ris.close();
     }
     baos.close();
     bu.put(key, baos.toByteArray());
   }
 }