package main

import (
	"bufio"
	"compress/gzip"
	"context"
	"encoding/gob"
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	_ "net/http/pprof"
	"os"
	"os/exec"
	"path/filepath"
	"regexp"
	"runtime"
	"sort"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"github.com/klauspost/pgzip"
	log "github.com/sirupsen/logrus"
)

type importer struct {
	tagLibraryFile      string
	refFile             string
	outputFile          string
	projectUUID         string
	loglevel            string
	priority            int
	runLocal            bool
	skipOOO             bool
	outputTiles         bool
	saveIncompleteTiles bool
	outputStats         string
	matchChromosome     *regexp.Regexp
	encoder             *gob.Encoder
	batchArgs
}

func (cmd *importer) RunCommand(prog string, args []string, stdin io.Reader, stdout, stderr io.Writer) int {
	var err error
	defer func() {
		if err != nil {
			fmt.Fprintf(stderr, "%s\n", err)
		}
	}()
	flags := flag.NewFlagSet("", flag.ContinueOnError)
	flags.SetOutput(stderr)
	flags.StringVar(&cmd.tagLibraryFile, "tag-library", "", "tag library fasta `file`")
	flags.StringVar(&cmd.refFile, "ref", "", "reference fasta `file`")
	flags.StringVar(&cmd.outputFile, "o", "-", "output `file`")
	flags.StringVar(&cmd.projectUUID, "project", "", "project `UUID` for output data")
	flags.BoolVar(&cmd.runLocal, "local", false, "run on local host (default: run in an arvados container)")
	flags.BoolVar(&cmd.skipOOO, "skip-ooo", false, "skip out-of-order tags")
	flags.BoolVar(&cmd.outputTiles, "output-tiles", false, "include tile variant sequences in output file")
	flags.BoolVar(&cmd.saveIncompleteTiles, "save-incomplete-tiles", false, "treat tiles with no-calls as regular tiles")
	flags.StringVar(&cmd.outputStats, "output-stats", "", "output stats to `file` (json)")
	cmd.batchArgs.Flags(flags)
	matchChromosome := flags.String("match-chromosome", "^(chr)?([0-9]+|X|Y|MT?)$", "import chromosomes that match the given `regexp`")
	flags.IntVar(&cmd.priority, "priority", 500, "container request priority")
	pprof := flags.String("pprof", "", "serve Go profile data at http://`[addr]:port`")
	flags.StringVar(&cmd.loglevel, "loglevel", "info", "logging threshold (trace, debug, info, warn, error, fatal, or panic)")
	err = flags.Parse(args)
	if err == flag.ErrHelp {
		err = nil
		return 0
	} else if err != nil {
		return 2
	} else if cmd.tagLibraryFile == "" {
		fmt.Fprintln(os.Stderr, "cannot import without -tag-library argument")
		return 2
	} else if flags.NArg() == 0 {
		flags.Usage()
		return 2
	}

	if *pprof != "" {
		go func() {
			log.Println(http.ListenAndServe(*pprof, nil))
		}()
	}

	lvl, err := log.ParseLevel(cmd.loglevel)
	if err != nil {
		return 2
	}
	log.SetLevel(lvl)

	cmd.matchChromosome, err = regexp.Compile(*matchChromosome)
	if err != nil {
		return 1
	}

	if !cmd.runLocal {
		err = cmd.runBatches(stdout, flags.Args())
		if err != nil {
			return 1
		}
		return 0
	}

	infiles, err := listInputFiles(flags.Args())
	if err != nil {
		return 1
	}
	infiles = cmd.batchArgs.Slice(infiles)

	taglib, err := cmd.loadTagLibrary()
	if err != nil {
		return 1
	}

	var outw, outf io.WriteCloser
	if cmd.outputFile == "-" {
		outw = nopCloser{stdout}
	} else {
		outf, err = os.OpenFile(cmd.outputFile, os.O_CREATE|os.O_WRONLY, 0777)
		if err != nil {
			return 1
		}
		defer outf.Close()
		if strings.HasSuffix(cmd.outputFile, ".gz") {
			outw = pgzip.NewWriter(outf)
		} else {
			outw = outf
		}
	}
	bufw := bufio.NewWriterSize(outw, 64*1024*1024)
	cmd.encoder = gob.NewEncoder(bufw)

	tilelib := &tileLibrary{taglib: taglib, retainNoCalls: cmd.saveIncompleteTiles, skipOOO: cmd.skipOOO}
	if cmd.outputTiles {
		cmd.encoder.Encode(LibraryEntry{TagSet: taglib.Tags()})
		tilelib.encoder = cmd.encoder
	}
	go func() {
		for range time.Tick(10 * time.Minute) {
			log.Printf("tilelib.Len() == %d", tilelib.Len())
		}
	}()

	err = cmd.tileInputs(tilelib, infiles)
	if err != nil {
		return 1
	}
	err = bufw.Flush()
	if err != nil {
		return 1
	}
	err = outw.Close()
	if err != nil {
		return 1
	}
	if outf != nil && outf != outw {
		err = outf.Close()
		if err != nil {
			return 1
		}
	}
	return 0
}

func (cmd *importer) runBatches(stdout io.Writer, inputs []string) error {
	if cmd.outputFile != "-" {
		// Not yet implemented, but this should write
		// the collection to an existing collection,
		// possibly even an in-place update.
		return errors.New("cannot specify output file in container mode: not implemented")
	}
	runner := arvadosContainerRunner{
		Name:        "lightning import",
		Client:      arvadosClientFromEnv,
		ProjectUUID: cmd.projectUUID,
		APIAccess:   true,
		RAM:         300000000000,
		VCPUs:       64,
		Priority:    cmd.priority,
		KeepCache:   1,
	}
	err := runner.TranslatePaths(&cmd.tagLibraryFile, &cmd.refFile, &cmd.outputFile)
	if err != nil {
		return err
	}
	for i := range inputs {
		err = runner.TranslatePaths(&inputs[i])
		if err != nil {
			return err
		}
	}

	outputs, err := cmd.batchArgs.RunBatches(context.Background(), func(ctx context.Context, batch int) (string, error) {
		runner := runner
		if cmd.batches > 1 {
			runner.Name += fmt.Sprintf(" (batch %d of %d)", batch, cmd.batches)
		}
		runner.Args = []string{"import",
			"-local=true",
			"-loglevel=" + cmd.loglevel,
			"-pprof=:6061",
			fmt.Sprintf("-skip-ooo=%v", cmd.skipOOO),
			fmt.Sprintf("-output-tiles=%v", cmd.outputTiles),
			fmt.Sprintf("-save-incomplete-tiles=%v", cmd.saveIncompleteTiles),
			"-match-chromosome", cmd.matchChromosome.String(),
			"-output-stats", "/mnt/output/stats.json",
			"-tag-library", cmd.tagLibraryFile,
			"-ref", cmd.refFile,
			"-o", "/mnt/output/library.gob.gz",
		}
		runner.Args = append(runner.Args, cmd.batchArgs.Args(batch)...)
		runner.Args = append(runner.Args, inputs...)
		return runner.RunContext(ctx)
	})
	if err != nil {
		return err
	}
	var outfiles []string
	for _, o := range outputs {
		outfiles = append(outfiles, o+"/library.gob.gz")
	}
	fmt.Fprintln(stdout, strings.Join(outfiles, " "))
	return nil
}

func (cmd *importer) tileFasta(tilelib *tileLibrary, infile string) (tileSeq, []importStats, error) {
	var input io.ReadCloser
	input, err := open(infile)
	if err != nil {
		return nil, nil, err
	}
	defer input.Close()
	input = ioutil.NopCloser(bufio.NewReaderSize(input, 8*1024*1024))
	if strings.HasSuffix(infile, ".gz") {
		input, err = pgzip.NewReader(input)
		if err != nil {
			return nil, nil, err
		}
		defer input.Close()
	}
	return tilelib.TileFasta(infile, input, cmd.matchChromosome)
}

func (cmd *importer) loadTagLibrary() (*tagLibrary, error) {
	log.Printf("tag library %s load starting", cmd.tagLibraryFile)
	f, err := open(cmd.tagLibraryFile)
	if err != nil {
		return nil, err
	}
	defer f.Close()
	rdr := ioutil.NopCloser(bufio.NewReaderSize(f, 64*1024*1024))
	if strings.HasSuffix(cmd.tagLibraryFile, ".gz") {
		rdr, err = gzip.NewReader(rdr)
		if err != nil {
			return nil, fmt.Errorf("%s: gzip: %s", cmd.tagLibraryFile, err)
		}
		defer rdr.Close()
	}
	var taglib tagLibrary
	err = taglib.Load(rdr)
	if err != nil {
		return nil, err
	}
	if taglib.Len() < 1 {
		return nil, fmt.Errorf("cannot tile: tag library is empty")
	}
	log.Printf("tag library %s load done", cmd.tagLibraryFile)
	return &taglib, nil
}

var (
	vcfFilenameRe    = regexp.MustCompile(`\.vcf(\.gz)?$`)
	fasta1FilenameRe = regexp.MustCompile(`\.1\.fa(sta)?(\.gz)?$`)
	fasta2FilenameRe = regexp.MustCompile(`\.2\.fa(sta)?(\.gz)?$`)
	fastaFilenameRe  = regexp.MustCompile(`\.fa(sta)?(\.gz)?$`)
)

func listInputFiles(paths []string) (files []string, err error) {
	for _, path := range paths {
		if fi, err := os.Stat(path); err != nil {
			return nil, fmt.Errorf("%s: stat failed: %s", path, err)
		} else if !fi.IsDir() {
			if !fasta2FilenameRe.MatchString(path) {
				files = append(files, path)
			}
			continue
		}
		d, err := os.Open(path)
		if err != nil {
			return nil, fmt.Errorf("%s: open failed: %s", path, err)
		}
		defer d.Close()
		names, err := d.Readdirnames(0)
		if err != nil {
			return nil, fmt.Errorf("%s: readdir failed: %s", path, err)
		}
		sort.Strings(names)
		for _, name := range names {
			if vcfFilenameRe.MatchString(name) {
				files = append(files, filepath.Join(path, name))
			} else if fastaFilenameRe.MatchString(name) && !fasta2FilenameRe.MatchString(name) {
				files = append(files, filepath.Join(path, name))
			}
		}
		d.Close()
	}
	for _, file := range files {
		if fastaFilenameRe.MatchString(file) {
			continue
		} else if vcfFilenameRe.MatchString(file) {
			if _, err := os.Stat(file + ".csi"); err == nil {
				continue
			} else if _, err = os.Stat(file + ".tbi"); err == nil {
				continue
			} else {
				return nil, fmt.Errorf("%s: cannot read without .tbi or .csi index file", file)
			}
		} else {
			return nil, fmt.Errorf("don't know how to handle filename %s", file)
		}
	}
	return
}

func (cmd *importer) tileInputs(tilelib *tileLibrary, infiles []string) error {
	starttime := time.Now()
	errs := make(chan error, 1)
	todo := make(chan func() error, len(infiles)*2)
	allstats := make([][]importStats, len(infiles)*2)
	var encodeJobs sync.WaitGroup
	for idx, infile := range infiles {
		idx, infile := idx, infile
		var phases sync.WaitGroup
		phases.Add(2)
		variants := make([][]tileVariantID, 2)
		if fasta1FilenameRe.MatchString(infile) {
			todo <- func() error {
				defer phases.Done()
				log.Printf("%s starting", infile)
				defer log.Printf("%s done", infile)
				tseqs, stats, err := cmd.tileFasta(tilelib, infile)
				allstats[idx*2] = stats
				var kept, dropped int
				variants[0], kept, dropped = tseqs.Variants()
				log.Printf("%s found %d unique tags plus %d repeats", infile, kept, dropped)
				return err
			}
			infile2 := fasta1FilenameRe.ReplaceAllString(infile, `.2.fa$1$2`)
			todo <- func() error {
				defer phases.Done()
				log.Printf("%s starting", infile2)
				defer log.Printf("%s done", infile2)
				tseqs, stats, err := cmd.tileFasta(tilelib, infile2)
				allstats[idx*2+1] = stats
				var kept, dropped int
				variants[1], kept, dropped = tseqs.Variants()
				log.Printf("%s found %d unique tags plus %d repeats", infile2, kept, dropped)

				return err
			}
		} else if fastaFilenameRe.MatchString(infile) {
			todo <- func() error {
				defer phases.Done()
				defer phases.Done()
				log.Printf("%s starting", infile)
				defer log.Printf("%s done", infile)
				tseqs, stats, err := cmd.tileFasta(tilelib, infile)
				allstats[idx*2] = stats
				if err != nil {
					return err
				}
				totlen := 0
				for _, tseq := range tseqs {
					totlen += len(tseq)
				}
				log.Printf("%s tiled %d seqs, total len %d", infile, len(tseqs), totlen)
				return cmd.encoder.Encode(LibraryEntry{
					CompactSequences: []CompactSequence{{Name: infile, TileSequences: tseqs}},
				})
			}
			// Don't write out a CompactGenomes entry
			continue
		} else if vcfFilenameRe.MatchString(infile) {
			for phase := 0; phase < 2; phase++ {
				phase := phase
				todo <- func() error {
					defer phases.Done()
					log.Printf("%s phase %d starting", infile, phase+1)
					defer log.Printf("%s phase %d done", infile, phase+1)
					tseqs, stats, err := cmd.tileGVCF(tilelib, infile, phase)
					allstats[idx*2] = stats
					var kept, dropped int
					variants[phase], kept, dropped = tseqs.Variants()
					log.Printf("%s phase %d found %d unique tags plus %d repeats", infile, phase+1, kept, dropped)
					return err
				}
			}
		} else {
			panic(fmt.Sprintf("bug: unhandled filename %q", infile))
		}
		encodeJobs.Add(1)
		go func() {
			defer encodeJobs.Done()
			phases.Wait()
			if len(errs) > 0 {
				return
			}
			err := cmd.encoder.Encode(LibraryEntry{
				CompactGenomes: []CompactGenome{{Name: infile, Variants: flatten(variants)}},
			})
			if err != nil {
				select {
				case errs <- err:
				default:
				}
			}
		}()
	}
	go close(todo)
	var tileJobs sync.WaitGroup
	var running int64
	for i := 0; i < runtime.GOMAXPROCS(-1)*2; i++ {
		tileJobs.Add(1)
		atomic.AddInt64(&running, 1)
		go func() {
			defer tileJobs.Done()
			defer atomic.AddInt64(&running, -1)
			for fn := range todo {
				if len(errs) > 0 {
					return
				}
				err := fn()
				if err != nil {
					select {
					case errs <- err:
					default:
					}
				}
				remain := len(todo) + int(atomic.LoadInt64(&running)) - 1
				if remain < cap(todo) {
					ttl := time.Now().Sub(starttime) * time.Duration(remain) / time.Duration(cap(todo)-remain)
					eta := time.Now().Add(ttl)
					log.Printf("progress %d/%d, eta %v (%v)", cap(todo)-remain, cap(todo), eta, ttl)
				}
			}
		}()
	}
	tileJobs.Wait()
	if len(errs) > 0 {
		// Must not wait on encodeJobs in this case. If the
		// tileJobs goroutines exited early, some funcs in
		// todo haven't been called, so the corresponding
		// encodeJobs will wait forever.
		return <-errs
	}
	encodeJobs.Wait()

	go close(errs)
	err := <-errs
	if err != nil {
		return err
	}

	if cmd.outputStats != "" {
		f, err := os.OpenFile(cmd.outputStats, os.O_CREATE|os.O_WRONLY, 0666)
		if err != nil {
			return err
		}
		var flatstats []importStats
		for _, stats := range allstats {
			flatstats = append(flatstats, stats...)
		}
		err = json.NewEncoder(f).Encode(flatstats)
		if err != nil {
			return err
		}
	}

	return nil
}

func (cmd *importer) tileGVCF(tilelib *tileLibrary, infile string, phase int) (tileseq tileSeq, stats []importStats, err error) {
	if cmd.refFile == "" {
		err = errors.New("cannot import vcf: reference data (-ref) not specified")
		return
	}
	args := []string{"bcftools", "consensus", "--fasta-ref", cmd.refFile, "-H", fmt.Sprint(phase + 1), infile}
	indexsuffix := ".tbi"
	if _, err := os.Stat(infile + ".csi"); err == nil {
		indexsuffix = ".csi"
	}
	if out, err := exec.Command("docker", "image", "ls", "-q", "lightning-runtime").Output(); err == nil && len(out) > 0 {
		args = append([]string{
			"docker", "run", "--rm",
			"--log-driver=none",
			"--volume=" + infile + ":" + infile + ":ro",
			"--volume=" + infile + indexsuffix + ":" + infile + indexsuffix + ":ro",
			"--volume=" + cmd.refFile + ":" + cmd.refFile + ":ro",
			"lightning-runtime",
		}, args...)
	}
	consensus := exec.Command(args[0], args[1:]...)
	consensus.Stderr = os.Stderr
	stdout, err := consensus.StdoutPipe()
	defer stdout.Close()
	if err != nil {
		return
	}
	err = consensus.Start()
	if err != nil {
		return
	}
	defer consensus.Wait()
	tileseq, stats, err = tilelib.TileFasta(fmt.Sprintf("%s phase %d", infile, phase+1), stdout, cmd.matchChromosome)
	if err != nil {
		return
	}
	err = stdout.Close()
	if err != nil {
		return
	}
	err = consensus.Wait()
	if err != nil {
		err = fmt.Errorf("%s phase %d: bcftools: %s", infile, phase, err)
		return
	}
	return
}

func flatten(variants [][]tileVariantID) []tileVariantID {
	ntags := 0
	for _, v := range variants {
		if ntags < len(v) {
			ntags = len(v)
		}
	}
	flat := make([]tileVariantID, ntags*2)
	for i := 0; i < ntags; i++ {
		for hap := 0; hap < 2; hap++ {
			if i < len(variants[hap]) {
				flat[i*2+hap] = variants[hap][i]
			}
		}
	}
	return flat
}