package com.netflix.staash.test.core;

import java.io.File;
import java.lang.annotation.Annotation;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;
import org.apache.cassandra.config.CFMetaData;
import org.apache.cassandra.config.KSMetaData;
import org.apache.cassandra.db.ColumnFamilyType;
import org.apache.cassandra.db.marshal.CounterColumnType;
import org.apache.cassandra.db.marshal.TypeParser;
import org.apache.cassandra.exceptions.AlreadyExistsException;
import org.apache.cassandra.service.MigrationManager;
import org.apache.cassandra.service.StorageProxy;
import org.apache.cassandra.service.StorageService;
import org.junit.runner.notification.RunNotifier;
import org.junit.runners.BlockJUnit4ClassRunner;
import org.junit.runners.model.FrameworkMethod;
import org.junit.runners.model.InitializationError;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

public class CassandraRunner extends BlockJUnit4ClassRunner {

  private static final Logger logger = LoggerFactory.getLogger(CassandraRunner.class);
  
  static StaashDeamon staashDaemon;

  static ExecutorService executor = Executors.newSingleThreadExecutor();

  public CassandraRunner(Class<?> klass) throws InitializationError {
    super(klass);
    logger.debug("CassandraRunner constructed with class {}", klass.getName());
  }

  @Override
  protected void runChild(FrameworkMethod method, RunNotifier notifier) {
    logger.debug("runChild invoked on method: " + method.getName());
    RequiresKeyspace rk = method.getAnnotation(RequiresKeyspace.class);
    RequiresColumnFamily rcf = method.getAnnotation(RequiresColumnFamily.class);
    if ( rk != null ) {
      maybeCreateKeyspace(rk, rcf);
    } else if ( rcf != null ) {
      maybeCreateColumnFamily(rcf);
    }

    super.runChild(method, notifier);
  }

  
  @Override
  public void run(RunNotifier notifier) {
    startCassandra();
    RequiresKeyspace rk = null;
    RequiresColumnFamily rcf = null;
    for (Annotation ann : getTestClass().getAnnotations() ) {
      if ( ann instanceof RequiresKeyspace ) {
        rk = (RequiresKeyspace)ann;
      } else if ( ann instanceof RequiresColumnFamily ) {
        rcf = (RequiresColumnFamily)ann;
      }
    }
    if ( rk != null ) {
      maybeCreateKeyspace(rk, rcf);
    } else if ( rcf != null ) {
      maybeCreateColumnFamily(rcf);
    }

    super.run(notifier);
  }

  
  private void maybeCreateKeyspace(RequiresKeyspace rk, RequiresColumnFamily rcf) {
    logger.debug("RequiresKeyspace annotation has keyspace name: {}", rk.ksName());
    List<CFMetaData> cfs = extractColumnFamily(rcf);
    try {
      MigrationManager
              .announceNewKeyspace(KSMetaData.newKeyspace(rk.ksName(),
                      rk.strategy(), KSMetaData.optsWithRF(rk.replication()), false, cfs));
    } catch (AlreadyExistsException aee) {
      logger.info("using existing Keyspace for " + rk.ksName());
      if ( cfs.size() > 0 ) {
        maybeTruncateSafely(rcf);
      }
    } catch (Exception ex) {
      throw new RuntimeException("Failure creating keyspace for " + rk.ksName(),ex);
    }
  }

  private List<CFMetaData> extractColumnFamily(RequiresColumnFamily rcf) {
    logger.debug("RequiresColumnFamily  has name: {} for ks: {}", rcf.cfName(), rcf.ksName());
    List<CFMetaData> cfms = new ArrayList<CFMetaData>();
    if ( rcf != null ) {
      try {
        cfms.add(new CFMetaData(rcf.ksName(), rcf.cfName(),
                ColumnFamilyType.Standard, TypeParser.parse(rcf.comparator()), null));

      } catch (Exception ex) {
        throw new RuntimeException("unable to create column family for: " + rcf.cfName(), ex);
      }
    }
    return cfms;
  }

  private void maybeCreateColumnFamily(RequiresColumnFamily rcf) {
    try {
      CFMetaData cfMetaData;
      if ( rcf.isCounter() ) {
        cfMetaData = new CFMetaData(rcf.ksName(), rcf.cfName(),
                      ColumnFamilyType.Standard, TypeParser.parse(rcf.comparator()), null)
                .replicateOnWrite(false).defaultValidator(CounterColumnType.instance);
      } else {
        cfMetaData = new CFMetaData(rcf.ksName(), rcf.cfName(),
                      ColumnFamilyType.Standard, TypeParser.parse(rcf.comparator()), null);
      }
      MigrationManager.announceNewColumnFamily(cfMetaData);
    } catch(AlreadyExistsException aee) {
      logger.info("CF already exists for " + rcf.cfName());
      maybeTruncateSafely(rcf);
    } catch (Exception ex) {
      throw new RuntimeException("Could not create CF for: " + rcf.cfName(), ex);
    }
  }

  private void maybeTruncateSafely(RequiresColumnFamily rcf) {
    if ( rcf != null && rcf.truncateExisting() ) {
      try {
        StorageProxy.truncateBlocking(rcf.ksName(), rcf.cfName());
      } catch (Exception ex) {
        throw new RuntimeException("Could not truncate column family: " + rcf.cfName(),ex);
      }
    }
  }

  private void startCassandra() {
    if ( staashDaemon != null ) {
      return;
    }
    deleteRecursive(new File("/tmp/staash_cache"));
    deleteRecursive(new File ("/tmp/staash_data"));
    deleteRecursive(new File ("/tmp/staash_log"));
    System.setProperty("cassandra-foreground", "true");
    System.setProperty("log4j.defaultInitOverride","true");
    System.setProperty("log4j.configuration", "log4j.properties");
    System.setProperty("cassandra.ring_delay_ms","1000");
    System.setProperty("cassandra.start_rpc","true"); 
    System.setProperty("cassandra.start_native_transport","true"); 

    executor.execute(new Runnable() {
      public void run() {
          staashDaemon = new StaashDeamon();
          staashDaemon.activate();
      }
    });
    try {
      TimeUnit.SECONDS.sleep(3);
    }
    catch (InterruptedException e) {
      throw new AssertionError(e);
    }
    Runtime.getRuntime().addShutdownHook(new Thread() {
      @Override
      public void run() {
        try {
          logger.error("In shutdownHook");
          stopCassandra();
        } catch (Exception ex) {
          ex.printStackTrace();
        }
      }
    });
  }

  private void stopCassandra() throws Exception {
    if (staashDaemon != null) {
        staashDaemon.deactivate();
      StorageService.instance.stopClient();

    }
    executor.shutdown();
    executor.shutdownNow();
  }

  private static boolean deleteRecursive(File path) {
    if (!path.exists())
      return false;
    boolean ret = true;
    if (path.isDirectory()){
      for (File f : path.listFiles()){
        ret = ret && deleteRecursive(f);
      }
    }
    return ret && path.delete();
  }
}