#!/usr/bin/perl -w

#
# tasuke_refseq_redis.pl
# http://tasuke.dna.affrc.go.jp/
#
# Author: Ryutaro Itoh
#
# Copyright (c) 2013 National Institute of Agrobiological Sciences   
#
# This software is released under the MIT License.
#
# http://opensource.org/licenses/mit-license.php
#

use strict;
use DBI;
use Getopt::Long;
use Redis;

my $DELETE       = 0;
my $DB           = '';
my $HOST         = 'localhost';
my $USER         = '';
my $PASS         = '';
my $REDISDB      = '';
my $REDISHOST    = 'localhost:6379';
my $FILE    = '';
my $RDBMS = '';
my $TABLESPACE = '';


################################################################
# Checking parameters and database connection
################################################################
GetOptions( 
            'r|restore'     => \$DELETE,
            'db|database=s' => \$DB,
            'h|host=s'      => \$HOST,
            'u|user=s'      => \$USER,
            'b|backend=s'   => \$RDBMS,
            's|space=s'   => \$TABLESPACE,
            'p|password=s'  => \$PASS,
            'z=s'           => \$REDISDB,
            'x=s'           => \$REDISHOST,
            'f|file=s'           => \$FILE
);

if ( !$DB || !$USER ||(!$DELETE && !-e$FILE)) {
    _error_method();
    exit;
}
if(!$RDBMS){
    $RDBMS = 'mysql';
}

my $connect;
if($RDBMS eq "oracle"){
    if(!$TABLESPACE){
        _error_method();
        exit;
    }
    $connect = DBI->connect('DBI:Oracle:host='.$HOST.';SERVICE_NAME='.$DB.';',$USER,$PASS)
    or die "Can't connect to database: $DBI::errstr\n";
    if(!_check_exists_tablespace($TABLESPACE)){
        print "Tablespace: '$TABLESPACE' is not found.";
        _exit(1);
    }
}else{
    $connect= DBI->connect('DBI:mysql:database='.$DB.';host='.$HOST,$USER,$PASS)
    or die "Can't connect to database: $DBI::errstr\n";
}

my $redis = Redis->new(server => $REDISHOST) or die "Can't connect to redis server.\n";
if($REDISDB ne ""){
   $redis->select($REDISDB) or die "Can't connect to database:$DB.\n"; 
}

$SIG{'INT'} = '_SIGINThandler';

my @config = $redis->config_get("maxmemory");
my $info = $redis->info();
my %infohash = %$info ;
if(_check_maxmem() != 0){
    print "Exit.\n";
    _exit(1);
}

###########################################################
# Restore
###########################################################
my @chrs = _chrs();
if($DELETE){
    _print_time("Start");
    for(my $i=0; $i<$#chrs+1; $i++){
        my $len = _get_chr_len($chrs[$i]);
        my $key = int($len/100)+1;
        for(my $x=0;$x<$key;$x++){
            if($redis->exists($chrs[$i].":".$x)){
                $redis->del($chrs[$i].":".$x);
            }
        }
        print "$chrs[$i] was deleted.\n";
    }
    $redis->del("TASUKE_REDIS");
    $redis->save();
    _print_time("Finished!");
    _exit();
}else{
    if($redis->exists("TASUKE_REDIS")){
        print "Reference sequence is already installed.\n";
        print "If you want to install again, please run this tool again with \"-r\" option for restore the reference sequence.\n";
        _exit(1);
    }
}

###########################################################
# Set sequences to the redis
###########################################################
my $CHR = "";
my $FIELD = 100;
my $CRTKEY = "";
my $CRTFIELD = 0;
my $pointer = 1;
_print_time("Start");
open IN, '<', $FILE or die "file open error: $!";
while(<IN>){
    chomp($_);
    if ($_ =~ /^>/i){
        $CHR = $_;
        $CHR =~ s/>//g;
        $CHR =~ s/^\s*(.*?)\s*$/$1/;
        my @chrline = split(/ /, $CHR);
        $CHR = $chrline[0];
        $pointer = 1;
        print $CHR."\n";
    }else{
        my $line_len = length($_);
        my @line = split(//,$_);
        foreach my $x (@line){
            $CRTKEY =  $CHR.":".int($pointer/100);
            $CRTFIELD = $pointer%100;
            eval {
                my $ret = $redis->hset($CRTKEY,$CRTFIELD,$x);
            };
            if($@) {
                print "Setting sequences to the database was failed.\n $CHR was deleted.";
                _exit(1);
            }
             $pointer++;
        }
    }
}
eval {
    my $ret = $redis->hset("TASUKE_REDIS","FLAG",1);
};
if($@) {
    print "Setting sequences to the database was failed.\n $CHR was deleted.";
    _exit(1);
}
$redis->save();
close(IN);
_print_time("Finished");
_exit();


###########################################################
# Functions
###########################################################
### check memory ####################################
sub _check_maxmem{
    my $ret = 0;
    my $used_mem = 0;
    my $file_size = 0;
    my $free_mem = 0;
    
    #file size
    $file_size = -s $FILE;
    #set maxmemory
    my @config_maxmem = $redis->config_get("maxmemory");
    if($config_maxmem[1] == 0){
        print "Set any value to maxmemory in redis config file.\n";
        print "Exit.\n";
        $ret = 1;
    }
    #used_memory
    my $info = $redis->info();
    my %infohash = %$info;
    if($infohash{"used_memory"} > ($config_maxmem[1]*0.75)){
        print "[Caution] Used memory: $used_mem%\n";
    }else{
        $used_mem = int($infohash{"used_memory"}/$config_maxmem[1]*100);
        print "[Notice] Used memory: $used_mem%\n"; 
    }
    $free_mem = $config_maxmem[1]-$infohash{"used_memory"};
    if($free_mem < $file_size){
        print "[Warning] It is possible that memory space will exceed the limit.\n";
        print "Free space: ".$free_mem." B\n";
        print "File size: ".$file_size." B\n";
        print "> Are you sure setting data to the database? [y|n] ";
        while (defined(my $line = <STDIN>)){
            chomp($line);
            if ( $line eq "y" || $line eq "Y" ) {
                last;
            } elsif ( $line eq "n" || $line eq "N" ) {
                $ret = 1;
                last;
            } else {
                print "> Are you sure setting data to the database? [y|n] ";
            }
        }   
    }
    return $ret;
}

### Execute DB ####################################
sub _dbexe {
    my ($query) = @_;
    my $query_handle = $connect->prepare($query);
    $query_handle->execute() or warn "Excute error: $DBI::errstr\n";
}

#### Usage #######################################
sub _error_method {
    print "\nUpgrading Database\n\n";
    print "tasuke_refseq_redis.pl -db [database] -u [user] -p [password] -f [reference(FASTA)]\n";
    print "Optional: (-x [redis database] -z [redis host])\n\n";
}

#### Time #######################################
sub _print_time {
    my ($str) = @_;
    my ($sec,$min,$hour,$mday,$mon,$year,$wday,$yday,$isdst) = localtime;
    $year += 1900;
    $mon += 1;
    my $datetime = sprintf "%04d-%02d-%02d %02d:%02d:%02d", $year, $mon, $mday, $hour, $min, $sec;
    print "$str : $datetime\n";
}

### chromosome ################################################
sub _chrs {
    my $val="";
    my @c;
    my $query = "SELECT id FROM sys_chromosome";
    my $query_handle = $connect->prepare($query);
    $query_handle->execute();
    $query_handle->bind_columns(undef, \$val);
    while($query_handle->fetch()) {
        push @c, $val;
    }
    return @c;
}

sub _get_chr_len {
    my ($seq) = @_;
    my $val="";
    my $query = "SELECT size FROM sys_chromosome where id='$seq'";
    my $query_handle = $connect->prepare($query);
    $query_handle->execute();
    $query_handle->bind_columns(undef, \$val);
    $query_handle->fetch();
    return $val;
}

#### Disconnection and exit #######################################
sub _exit {
    my ($arg) = @_;
    my $code = 0;
    if(defined $arg){$code = $arg;}
    $connect->disconnect or warn "Disconnection error: $DBI::errstr\n";
    exit($code);
}

### SIGINT ###
sub _SIGINThandler{
    print "\n\nFinished.\n";
    _exit(1);
}
