#!/usr/bin/env perl
use strict;
use Bio::Seq;
use Bio::SeqIO;
use File::Spec;
use File::Path qw(make_path remove_tree);
use File::Slurp qw(edit_file read_file);
use List::Util qw(first);
use Cwd qw(abs_path);
use Data::Dumper;
use LWP::Simple;
use Text::CSV;
use JSON;
use FindBin;
#use lib "$FindBin::RealBin/../perl5";
#use Abricate::Headers qw(hash_encode hash_decode);

#..............................................................................
# Globals

my %DATABASE = (
  'resfinder'         => \&get_resfinder,
  'plasmidfinder'     => \&get_plasmidfinder,
#  'megares'           => \&get_megares,
  'argannot'          => \&get_argannot,
  'card'              => \&get_card,
#  'ncbibetalactamase' => \&get_ncbibetalactamase,
  'ncbi'              => \&get_ncbi,
  'vfdb'              => \&get_vfdb,
  'ecoh'              => \&get_ecoh,
#  'serotypefinder'    => \&get_serotypefinder,
);
my $DATABASES = join(' ', sort keys %DATABASE);

#..............................................................................
# Command line options

my(@Options, $debug, $outdir, $db, $force);
setOptions();

$db or err("Please choose a --db from: $DATABASES");
exists $DATABASE{$db} or err("Unknown --db '$db', choose from: $DATABASES ");
-d $outdir or err("--outdir '$outdir' does not exist");

my $dir = abs_path( File::Spec->catdir($outdir, $db) );
make_path($dir);
msg("Setting up '$db' in '$dir'");
#my $tmpdir = tempdir("$db-XXXXXXXX", DIR=>$dir, CLEANUP=>0);
#my $tmpdir = "/home/tseemann/git/abricate/db/resfinder/resfinder-6Kuphtvv";
my $tmpdir = "$dir/src";
make_path($tmpdir);

# run the specific function from --db
chdir $tmpdir;
my $seq = $DATABASE{$db}->();
map { is_full_gene($_) } @$seq;
$seq = dedupe_seq($seq);
#print Dumper($seq);
msg("Sorting sequences by ID");
$seq = [ sort { $a->{ID} cmp $b->{ID} } @$seq ];
save_fasta("$dir/sequences", $seq);

msg("Formatting BLASTN database: $dir/sequences");
my $logfile = "$tmpdir/makeblastdb.log";
my $ec = system("makeblastdb -in '$dir/sequences' -title '$db' -dbtype nucl -parse_seqids -hash_index -logfile $logfile");
if ($ec != 0) {
  system("tail '$logfile'");
  err("Error with makign BLAST database. See $logfile");
}
#msg("Run 'abricate --setupdb' to format the database");

msg("Done.");

#..............................................................................

sub download {
  my($url, $dest) = @_;
  if (-r $dest and not $force) {
     msg("Won't re-download existing $dest (use --force)");
     exit(1);
  }
  else {
    msg("Downloading: $url");
    my $ec = mirror($url, $dest);
    msg("Result: $ec");
  }
  msg("Destination: $dest");
  msg("Filesize:", (-s $dest), "bytes");
}

#..............................................................................
sub get_resfinder {
  my $name = "resfinder_db";
  my $url = "https://bitbucket.org/genomicepidemiology/$name.git";

  if (-r $name and not $force) {
    msg("Won't overwrite existing $name (use --force)");
#    exit(1);
  }
  else {
    msg("Nuking existing folder: $name");
    remove_tree("./$name");
    msg("Cloning $url to $name");
    system("git clone --quiet $url $name");
  }

  my %anno;
  open NOTES, '<', "$name/notes.txt";
  while (<NOTES>) {
    chomp;
    chop; # remove the DOS CR ending
    next if m/^#/;
    my @x = split m/:/;
    $anno{$x[0]} = join(':', @x[2..$#x]);
  }
  msg("get_resfinder: notes.txt has", scalar(keys %anno), "genes");
  #print Dumper(\%anno);

  my @seq;
  for my $fasta (<$name/*.fsa>) {
    push @seq, @{ load_fasta($fasta) };
  }

  for my $seq (@seq) {
    my($id,$copy,$acc) = $seq->{ID} =~ m/^(.*)_(\d+)_(\w+)$/;
    #msg("resfinder: $1 $2 $3", $anno{$1});
    $seq->{ID}   = "${id}_${copy}";
    $seq->{ACC}  = $acc;
    $seq->{DESC} = $anno{$id} || '';
  }
  
  return \@seq;
}

#..............................................................................
sub get_serotypefinder {
  my $name = "serotypefinder_db";
  my $url = "https://bitbucket.org/genomicepidemiology/$name.git";

  if (-r $name and not $force) {
    msg("Won't overwrite existing $name (use --force)");
#    exit(1);
  }
  else {
    msg("Nuking existing folder: $name");
    remove_tree("./$name");
    msg("Cloning $url to $name");
    system("git clone --quiet $url $name");
  }

  my @seq;
  for my $fasta (<$name/*.fsa>) {
    push @seq, @{ load_fasta($fasta) };
  }

  # >fliC_44444_AY250028_H52
  # FIXME - this is already in EcOH database!
  for my $seq (@seq) {
    my($id,$copy,$acc) = $seq->{ID} =~ m/^(.*)_(\d+)_(\w+)$/;
    #msg("serotypefinder: $1 $2 $3", $anno{$1});
    $seq->{ID}   = "${id}_${copy}";
    $seq->{ACC}  = $acc;
    #$seq->{DESC} = $anno{$id} || '';
  }

  return \@seq;
}

#..............................................................................
sub get_tag {
  my($f, $tag) = @_;
  if ($f->has_tag($tag)) {
    my($val) = $f->get_tag_values($tag);
    return $val;
  }
  return '';
}

#..............................................................................
sub get_ncbi {
  my $name = "ncbi.gbk";
  my $cmd = "esearch -db nucleotide -query '313047[BioProject]' | efetch -db nuccore -format gbwithparts";

  if (-r $name and not $force) {
    msg("Won't overwrite existing $name (use --force)");
#    exit(1);
  }
  else {
    msg("Running: $cmd");
    system("$cmd | tee $name | grep LOCUS");
  }

  # LOCUS       NG_054996               1152 bp    DNA     linear   CON 13-JUN-2017
  #     gene            1..1194
  #                     /gene="blaPDC"
  #                     /locus_tag="A7J11_05084"
  #                     /allele="blaPDC-210"
  #     CDS             1..1194
  #                     /gene="blaPDC"
  #                     /locus_tag="A7J11_05084"
  #                     /allele="blaPDC-210"
  #                     /codon_start=1
  #                     /transl_table=11
  #                     /product="class C beta-lactamase PDC-210"
  #                     /protein_id="WP_087587974.1"

  my %seen;
  my @seq;
  my $in = Bio::SeqIO->new(-file=>$name, -format=>"genbank");
  while (my $rec = $in->next_seq) {
    my $cds = first { $_->primary_tag eq 'CDS' } $rec->get_SeqFeatures;
    $cds or die("Could not find a CDS in", $rec->id);
    my $s = {
      ID   => get_tag($cds, 'allele') || get_tag($cds, 'gene') || get_tag($cds,'locus_tag'),
      ACC  => get_tag($cds, 'locus_tag'),
      DESC => get_tag($cds, 'product'),
      SEQ  => $cds->spliced_seq->seq,
    };

#    if ($seen{ $s->{ID} }++) {
#      msg("Duplicate", $s->{ID}, "x", $seen{ $s->{ID} } );
#      my $end = (split ' ', $s->{DESC})[-1];
#      msg("\tend=$end");
#      if ( lc(substr($s->{ID},0,3)) eq lc(substr($end,0,3)) ) {
#        msg("\t",substr($end,0,3));
#      }
#    }

    push @seq, $s;
    msg("[$name]", 0+@seq, "|", $s->{ID}, "|", $s->{DESC});
  }
  return \@seq;
}

#..............................................................................
sub get_plasmidfinder {
  my $name = "plasmidfinder";
  my $zip = "$name.zip";
  download("https://cge.cbs.dtu.dk/cge/download_data.php?folder=$name&filename=$zip&submit=$zip", $zip);
  system("unzip -j -u $zip");

  my @seq;
  for my $fasta (<*.fsa>) {
    push @seq, @{ load_fasta($fasta) };
  }

  for my $seq (@seq) {
    $seq->{DESC} = $seq->{ID};  # no desc, so use ORIGINAL ID as desc
    my($id,$acc) = ($seq->{ID} =~ m/^(.*)_(([A-Z]+|NC_)\d+(\.\d+)?)$/);
    $id =~ s/_+$//g;
    $seq->{ID}   = $id || $seq->{ID};
    $seq->{ACC}  = $acc || '';
    whine("Parsed empty ID:", $seq->{DESC}, "=> id='$id' acc='$acc' seq=".substr($seq->{SEQ},0,10)) if not $id;
  }
  
  return \@seq;
}

#..............................................................................
sub get_megares {
  err("MEGARES is currently not supported due issues with its confused metadata.");
  my $zip = "megares.zip";
  download('https://megares.meglab.org/download/megares_v1.01.zip', $zip);
  system("unzip -j -u $zip");
  my $seqs = load_fasta( glob("megares_database_*.fasta") );
  
  # >Bla|OXA-223|JN248564|1-825|825|betalactams|Class_D_betalactamases|OXA
  # >gi|558607393|gb|KF513179.1|betalactams|Class_A_betalactamases|TEM
  # >CARD|phgb|AM408573|10027-11077|ARO:3003441|cfrA|Phenicol|Florfenicol_methyltransferases|CFRA
  for my $s (@$seqs) {
    # they use | for field sep AND within some fields, arrg
    $s->{ID} =~ s/\|?(gi|gb)\|(\S+)\|/$2/;
    my($source,$gene,$acc,$loc,undef,$class,$desc,$family) = split m/\|/, $s->{ID};
    $s->{ID} = $source eq 'CARD' ? $class : $gene;
    $s->{ACC} = "$acc:$loc";
    $s->{DESC} = join(' ', $desc, $class, $family);
  }
  return $seqs;
}

#..............................................................................
sub get_argannot {
  my $fasta = 'arg-annot.fa';
  download(
    'http://www.mediterranee-infection.com/arkotheque/client/ihumed/_depot_arko/articles/1425/argannot-aa-v3-march2017_doc.fasta',
#    'http://www.mediterranee-infection.com/arkotheque/client/ihumed/_depot_arko/articles/691/argannot-nt_doc.fasta',
    $fasta
  );
  
  # fix syntax errors in the FASTA file...
  edit_file { s/\\//g } $fasta;

  my $seqs = load_fasta($fasta);

  #  0             1         2               3 
  # >(AGly)Aac2-Ie:NC_011896:3039059-3039607:549
  for my $s (@$seqs) {
    my @x = split m/:/, $s->{ID};
    $s->{ID} = $x[0];
    $s->{ACC} = $x[1].':'.$x[2];
    $s->{DESC} = '';
  }

  return $seqs;
}

#..............................................................................
sub get_card {
  download(
    'https://card.mcmaster.ca/download/0/broadstreet-v2.0.0.tar.gz',
    "card.tar.bz2"   # yes, it's really BZ2 not GZ ...
  );
#  my $fasta = "./nucleotide_fasta_protein_homolog_model.fasta";
  my $jsonfile = "./card.json";
  system("tar xf card.tar.bz2 $jsonfile");

  # JSON
#  my $json = read_file($json, binmode => ':utf8');
  my $json = read_file($jsonfile);
  my $card = from_json( $json, { latin1=>1 });
#  my $card = decode_json( read_file($json) );
#  my $card = from_json( read_file($json, binmode => ':utf8'), { utf8=>1 } );
#  print Dumper($meta);
  my @seq;
  for my $g (values %$card) {
    next unless ref($g) eq 'HASH';
#    msg(Dumper($g));
    next unless $g->{model_type} eq "protein homolog model";  # only 'acquired' genes
    my $id =  $g->{model_name};
    err("$id has {model_param}{snp}") if exists $g->{model_param}{snp};
#    msg("CARD: $id");
#    print STDERR Dumper($g);
    my $dna =  $g->{model_sequences}{sequence} or err("$id: no {model_sequences}{sequence} found");
    my($key) = sort keys %$dna;  # first key
    $dna = $dna->{$key} or err("$id: invalid key '$key'"); 
    $dna = $dna->{dna_sequence} or err("$id: no dna_sequence");
#    msg(Dumper($dna)) if $id eq 'OXA-25';
    $id =~ s/\s+/_/g;
    push @seq, {
      ID   => $id,
      ACC  => $dna->{accession}.':'.($dna->{fmin}+1).'-'.($dna->{fmax}+1),
      DESC => ($g->{ARO_description} || $g->{ARO_accession}),
      SEQ  => $dna->{sequence},
    };
#    msg(Dumper($seq[-1]));
  }

  return \@seq;
}

#..............................................................................
sub get_vfdb {
  download('http://www.mgc.ac.cn/VFs/Down/VFDB_setA_nt.fas.gz', 'vfdb.fa.gz');
  system("gzip -f -d -c vfdb.fa.gz > vfdb.fa");
  my $seqs = load_fasta("vfdb.fa");

  # >VFG000676(gb|AAD32411) (lef) anthrax toxin lethal factor precursor [Anthrax toxin (VF0142)] [Bacillus anthracis str. Sterne]
  for my $s (@$seqs) {
#    print STDERR Dumper($s);
    $s->{ID} =~ m/^(\w+)\(\w+\|(\w+)\)$/;
    $s->{ACC} = $2 if $2;
    $s->{DESC} =~ m/^\((.*?)\)/;
    $s->{ID}  = $1 if $1;
#    print STDERR Dumper($s); exit;
  }

  return $seqs;
}

#..............................................................................
sub get_ncbibetalactamase {
  my $fasta = "ncbi.fa";
  download('ftp://ftp.ncbi.nlm.nih.gov/pathogen/betalactamases/Allele-dna.fa', $fasta);
  my $tab = "ncbi.tab";
  download('ftp://ftp.ncbi.nlm.nih.gov/pathogen/betalactamases/Allele.tab', $tab);

  # >ACD12694.1 EU650653.1:1-1173
  my $seqs = load_fasta($fasta);
  
  # ACC-1 ACD12694.1 EU650653.1 blaACC-1 1 1173 + cephalosporin-hydrolyzing class C beta-lactamase ACC-1
  my %anno;
  my @anno = read_file($tab);
  msg("Read", 0+@anno, "annotations");
  foreach (@anno) {
    next if m/^#/;
    chomp;
    my($name,$id,$acc,$gene,$begin,$end,undef,$product) = split m/\t/;
    $anno{$id} = {
      ID => $gene,
      DESC => $product,
      ACC => "$acc:$begin-$end",
    };
  }
#  print Dumper(\%anno);

  for my $s (@$seqs) {
    my $id = $s->{ID};
    next unless exists $anno{$id};
    $s->{ID} = $anno{$id}{ID};
    $s->{ACC} = $anno{$id}{ACC};
    $s->{DESC} = $anno{$id}{DESC};
  }
#  print Dumper($seqs);

  return $seqs;
}

#..............................................................................
sub get_ecoh {
  my $fasta = "EcOH.fa";
  download('https://raw.githubusercontent.com/katholt/srst2/master/data/EcOH.fasta', $fasta);

  # https://github.com/katholt/srst2#generating-srst2-compatible-clustered-database-from-raw-sequences
  # [clusterID]__[gene]__[allele]__[seqID] [other stuff]
  # >1__fliC__fliC-H1__1 AB028471.1;flagellin;H1
  # >8__wzx__wzx-O41__246 AB811617.1;O antigen flippase;O41
  # >9__wzy__wzy-OgN31__597 LC125932.1;O antigen polyermase;OgN31
  my $seqs = load_fasta($fasta);

  for my $s (@$seqs) {
    my @id = split m/__/, $s->{ID};
    my @desc = split m';', $s->{DESC};
    $s->{ID} = $id[2];
    $s->{ACC} = shift(@desc);
    $s->{DESC} = join(' ', @desc);
  }
#  print Dumper($seqs);
  return $seqs;
}

#..............................................................................
sub is_full_gene {
  my($s) = @_;

  my $id = $s->{ID};
  my $L = length($s->{SEQ});
  if ($L % 3 != 0) {
    whine("$id - length $L bp is not multiple of 3");
    return;
  }
  if ($s->{SEQ} !~  m/^[AGCT]+$/) {
    whine("$id - has non-AGTC bases");
    # don't return though
  }

  my $seq = Bio::Seq->new( -id=>$s->{ID}, -seq=>$s->{SEQ} );
  my $aa = $seq->translate->seq;
  
  if ($aa =~ m/\*./) {
    whine("$id - has internal stop codons, trying revcom");
    $aa = $seq->revcom->translate->seq;
    if ($aa =~ m/\*./) {
      whine("$id - revcom has internal stop codons too");
      return 0;
    }
    else {
      msg("$id - revcom resolves problem, hooray!");
      $s->{SEQ} = $seq->revcom->seq;
    }
  }                                                                                                                                                                                                            
}

#..............................................................................
sub dedupe_seq {
  my($seq) = @_;
  my %seen;
  my $good = [];
  for my $s (@$seq) {
    if ($seen{ $s->{SEQ} }) {
      whine("duplicate", length($s->{SEQ}), "bp sequence:", 
        $s->{ID}, '~', $seen{$s->{SEQ}} );
    }
    else {
      push @$good, $s;
    }
    $seen{ $s->{SEQ} } .= ' '.$s->{ID};
  }
  msg( "dedupe_seq: read", scalar(@$seq), "/ kept", scalar(@$good) );
  return $good;
}

#..............................................................................
sub load_fasta {
  my($fasta) = @_;
  my %seen;
  my $list;
  msg("load_fasta: $fasta");
  my $in = Bio::SeqIO->new(-file=>$fasta, -format=>'fasta');
  while (my $seq = $in->next_seq) {
    my $id = $seq->id or err("Empty ID in $fasta");
    if ($seen{$id}) {
      whine("Duplicate ID '$id' in $fasta");
      $id = $id.'_dupe';
    }
    $seen{$id}++;
    my $dna = uc($seq->seq);
    $dna =~ s/[^AGTC]/N/g;
    push @$list, {
      ID   => $id,
      ACC  => '',
      DESC => $seq->desc,
      SEQ  => $dna,
    }
  }
  msg("load_fasta: read", scalar(@$list), "sequences");
  return $list;
}

#..............................................................................
sub save_fasta {
  my($fasta, $seq) = @_;
  msg("save_fasta: $fasta");
  my %seen;
  my $out = Bio::SeqIO->new(-file=>">$fasta", -format=>'fasta');
  for my $s (@$seq) {
    $seen{$s->{ID}}++;
    my $freq = $seen{$s->{ID}};
    whine("seen $s->{ID} now $freq times") if $freq > 1;
#    print Dumper($s);
    my $obj = Bio::Seq->new(
      -id   => join('~~~', $db, $s->{ID}, $s->{ACC}),
      -desc => ($s->{DESC} || $s->{ID}),
      -seq  => $s->{SEQ},
    );
#    $obj->desc( hash_encode($s) );
    $out->write_seq($obj);
#    $seen{ $s->{ID} }++;
  }
  msg("save_fasta: wrote", scalar(@$seq), "sequences");
}

#----------------------------------------------------------------------
sub msg {
  print STDERR "@_\n";
}

#----------------------------------------------------------------------
sub whine {
  print STDERR "WARNING: @_\n";
}

#----------------------------------------------------------------------
sub err {
  msg("ERROR:", @_);
  exit(1);
}

#----------------------------------------------------------------------
# Option setting routines

sub setOptions {
  use Getopt::Long;

  @Options = (
    {OPT=>"help",    VAR=>\&usage,             DESC=>"This help"},
    {OPT=>"debug!",  VAR=>\$debug, DEFAULT=>0, DESC=>"Verbose debug output"},
    {OPT=>"dbdir=s", VAR=>\$outdir, DEFAULT=>"$FindBin::RealBin/../db", DESC=>"Parent folder"},
    {OPT=>"db=s",    VAR=>\$db, DEFAULT=>"",  DESC=>"Options: $DATABASES" },
    {OPT=>"force!",  VAR=>\$force, DEFAULT=>0, DESC=>"Force download even if exists"},
  );

#  @ARGV or usage();

  &GetOptions(map {$_->{OPT}, $_->{VAR}} @Options) || usage(1);

  # Now setup default values.
  foreach (@Options) {
    if (defined($_->{DEFAULT}) && !defined(${$_->{VAR}})) {
      ${$_->{VAR}} = $_->{DEFAULT};
    }
  }
}

sub usage {
  my($exitcode) = @_;
  $exitcode = 0 if $exitcode eq 'help'; # what gets passed by getopt func ref
  $exitcode ||= 0;
  select STDERR if $exitcode; # write to STDERR if exitcode is error
  
  print "Usage: $0 [options]\n";
  foreach (@Options) {
    printf "  --%-13s %s%s.\n",$_->{OPT},$_->{DESC},
           defined($_->{DEFAULT}) ? " (default '$_->{DEFAULT}')" : "";
  }
  exit($exitcode);
}
 
#----------------------------------------------------------------------

