#!/usr/bin/env perl

use strict;
use warnings;
use Carp;
use Vcf;

my $opts = parse_params();
do_consensus($opts);

exit;

#--------------------------------

sub error
{
    my (@msg) = @_;
    if ( scalar @msg )
    {
        croak @msg;
    }
    die
        "Usage: cat ref.fa | vcf-consensus [OPTIONS] in.vcf.gz > out.txt\n",
        "Options:\n",
        "   -h, -?, --help                   This help message.\n",
        "   -s, --sample <name>              \n",
        "\n";
}


sub parse_params
{
    my $opts = { };
    while (my $arg=shift(@ARGV))
    {
        if ( $arg eq '-?' || $arg eq '-h' || $arg eq '--help' ) { error(); }
        if ( $arg eq '-s' || $arg eq '--sample' ) { $$opts{sample}=shift(@ARGV); next; }
        if ( -e $arg && !exists($$opts{vcf_file}) ) { $$opts{vcf_file}=$arg; next; }
        error("Unknown parameter \"$arg\". Run -h for help.\n");
    }
    return $opts;
}

sub do_consensus
{
    my ($opts) = @_;
    
    my $vcf = Vcf->new(file=>$$opts{vcf_file});
    $vcf->parse_header;
    if ( exists($$opts{sample}) )
    {
        if ( !exists($$vcf{has_column}{$$opts{sample}}) ) { error("No such sample: $$opts{sample}"); }
        $$opts{vcf} = $vcf; 
        $$opts{sample_col} = $$vcf{has_column}{$$opts{sample}};
    }
    my $chrs = $vcf->get_chromosomes();
    my %chrs = map { $_=>0 } @$chrs;

    my ($chr,$vcf_pos,$warned,$vcf_line);
    while (my $line=<STDIN>)
    {
        if ( $line=~/^>([^:\s]+)/ ) 
        {
            flush_fa_buffer($opts,0);
            $chr = $1;
            my $rest = $';
            $$opts{fa_pos} = ($rest=~/^:(\d+)-\d+$/) ? $1 : 1;
            $$opts{fa_idx} = 0;
            if ( exists($chrs{$chr}) ) { $chrs{$chr}=1; }
            $vcf->open(region=>$chr);
            print $line;
            next;
        }

        chomp($line);
        $$opts{fa_buf} .= $line;
        $$opts{fa_len} += length($line);

        while ( defined($vcf_line = $vcf->next_data_array()) )
        {
            if ( $$opts{fa_pos}+$$opts{fa_len}-$$opts{fa_idx}<=$$vcf_line[1] )
            {
                $vcf->_unread_line($vcf_line);
                flush_fa_buffer($opts,60);
                last;
            }
            apply_variant($opts,$vcf_line);
        }

        if ( !defined $vcf_line ) { flush_fa_buffer($opts,60); }
    }
    flush_fa_buffer($opts,0);

    for my $chr (keys %chrs)
    {
        if ( !$chrs{$chr} ) { warn("The sequence \"$chr\" not found in the fasta file.\n"); }
    }
}

sub flush_fa_buffer
{
    my ($opts,$len) = @_;
    while ( $$opts{fa_len} && $$opts{fa_len}>=60 )
    {
        print substr($$opts{fa_buf},0,60,''), "\n";
        $$opts{fa_len} -= 60;
        $$opts{fa_pos} += 60 - $$opts{fa_idx};
        $$opts{fa_idx}  = 0;
    }
    if ( $len or !$$opts{fa_len} ) { return; }
    print $$opts{fa_buf},"\n";
    $$opts{fa_pos} += $$opts{fa_len}-$$opts{fa_idx};
    $$opts{fa_len} = 0;
    $$opts{fa_buf} = '';
    $$opts{fa_idx} = 0;
}

sub apply_variant
{
    my ($opts,$vline) = @_;

    if ( $$vline[4] eq '.' ) { return; }

    my $pos = $$vline[1] - $$opts{fa_pos} + $$opts{fa_idx};
    if ( $pos<0 or $pos>=$$opts{fa_len} ) { error("FIXME: $$vline[0]:$$vline[1] .. $$opts{fa_pos},$$opts{fa_len}\n"); }

    # Sanity check
    my $ref_len = length($$vline[3]);
    if ( substr($$vline[3],0,$ref_len) ne substr($$opts{fa_buf},$pos,$ref_len) ) 
    { 
        error(sprintf "The fasta sequence does not match the REF at $$vline[0]:$$vline[1]. %s(%s) in .fa, %s in .vcf\n", 
            substr($$opts{fa_buf},$pos,$ref_len),substr($$opts{fa_buf},$pos+1,$ref_len+5),substr($$vline[3],0,$ref_len));
    }

    my $alt_len = length($$vline[4]);
    my $alt;
    if ( !exists($$opts{sample_col}) )
    {
        my $idx;
        $alt = ($idx=index($$vline[4],','))==-1 ? $$vline[4] : substr($$vline[4],0,$idx);
    }
    else
    {
        my @als = $$opts{vcf}->split_gt($$vline[$$opts{sample_col}-1]);
        for my $al (@als)
        {
            if ( $al eq '0' ) { next; }
            $alt = $$opts{vcf}->get_field($$vline[4],$al-1,',');
            last;
        }
        if ( !defined $alt ) { return; }
    }

    substr($$opts{fa_buf},$pos,$ref_len,$alt);
    $$opts{fa_len} += $alt_len - $ref_len;
    $$opts{fa_pos} += $ref_len;
    $$opts{fa_idx} += $alt_len;
}


