use crate::{
    annotation::{Annotation, Annotations, Caller},
    collection::{vcf::Vcf, Initialize},
    commands::bcftools::{bcftools_concat, bcftools_keep_pass, BcftoolsConfig},
    config::Config,
    helpers::{force_or_not, get_temp_file_path},
    io::vcf::read_vcf,
    runners::{run_wait, DockerRun, Run},
    variant::{
        variant::{RunnerVariants, Variants},
        variant_collection::VariantCollection,
    },
};
use anyhow::{Context, Ok};
use rayon::iter::{IntoParallelRefIterator, ParallelIterator};
use std::{fs, path::Path};
use log::info;

#[derive(Debug, Clone)]
pub struct ClairS {
    pub id: String,
    pub output_dir: String,
    pub output_vcf: String,
    pub output_indels_vcf: String,
    pub vcf_passed: String,
    pub diag_bam: String,
    pub mrd_bam: String,
    pub log_dir: String,
    pub config: Config,
    pub clair3_germline_normal: String,
    pub clair3_germline_tumor: String,
    pub clair3_germline_passed: String,
}

impl Initialize for ClairS {
    fn initialize(id: &str, config: Config) -> anyhow::Result<Self> {
        let id = id.to_string();
        info!("Initialize ClairS for {id}.");
        let log_dir = format!("{}/{}/log/clairs", config.result_dir, &id);

        if !Path::new(&log_dir).exists() {
            fs::create_dir_all(&log_dir)
                .context(format!("Failed  to create {log_dir} directory"))?;
        }

        let output_dir = config.clairs_output_dir(&id);
        fs::create_dir_all(&output_dir).context(format!("Can't create dir: {output_dir}"))?;

        let (output_vcf, output_indels_vcf) = config.clairs_output_vcfs(&id);
        let vcf_passed = format!("{output_dir}/{id}_diag_clairs_PASSED.vcf.gz");

        let diag_bam = config.tumoral_bam(&id);
        let mrd_bam = config.normal_bam(&id);

        let clair3_germline_normal = config.clairs_germline_normal_vcf(&id);
        let clair3_germline_tumor = config.clairs_germline_tumor_vcf(&id);
        let clair3_germline_passed = config.clairs_germline_passed_vcf(&id);

        Ok(Self {
            id,
            output_dir,
            output_vcf,
            output_indels_vcf,
            vcf_passed,
            diag_bam,
            mrd_bam,
            log_dir,
            config,
            clair3_germline_normal,
            clair3_germline_tumor,
            clair3_germline_passed,
        })
    }
}

impl Run for ClairS {
    fn run(&mut self) -> anyhow::Result<()> {
        force_or_not(&self.vcf_passed, self.config.clairs_force)?;

        // Run Docker command if output VCF doesn't exist
        if !Path::new(&self.output_vcf).exists() || !Path::new(&self.output_indels_vcf).exists() {
            let mut docker_run = DockerRun::new(&[
                "run",
                "-d",
                "-v",
                "/data:/data",
                "-v",
                &format!("{}:{}", self.output_dir, self.output_dir),
                "hkubal/clairs:latest",
                "/opt/bin/run_clairs",
                "-T",
                &self.diag_bam,
                "-N",
                &self.mrd_bam,
                "-R",
                &self.config.reference,
                "-t",
                &self.config.clairs_threads.to_string(),
                "-p",
                &self.config.clairs_platform,
                "--enable_indel_calling",
                "--include_all_ctgs",
                "--print_germline_calls",
                "--enable_clair3_germline_output",
                "--output_dir",
                &self.output_dir,
                "-s",
                &format!("{}_diag", self.id),
            ]);
            let report = run_wait(&mut docker_run).context(format!(
                "Error while running ClairS in docker for  {} and {}",
                &self.diag_bam, &self.mrd_bam
            ))?;

            let log_file = format!("{}/clairs_", self.log_dir);
            report
                .save_to_file(&log_file)
                .context(format!("Error while writing logs into {log_file}"))?;
        } else {
            info!("ClairS VCFs exist.");
        }

        // Germline
        if !Path::new(&self.clair3_germline_passed).exists() {
            // let tmp_file = format!("{}.vcf.gz", get_temp_file_path()?.display());
            // let report = bcftools_concat(
            //     vec![
            //         self.clair3_germline_tumor.to_string(),
            //         self.clair3_germline_normal.to_string(),
            //     ],
            //     &tmp_file,
            //     BcftoolsConfig::default(),
            // )
            // .context(format!(
            //     "Error while running bcftools concat for {} and {}",
            //     &self.output_vcf, &self.output_indels_vcf
            // ))?;
            //
            // let log_file = format!("{}/bcftools_concat_", self.log_dir);
            // report
            //     .save_to_file(&log_file)
            //     .context(format!("Error while writing logs into {log_file}"))?;
            //
            let report = bcftools_keep_pass(
                &self.clair3_germline_normal,
                &self.clair3_germline_passed,
                BcftoolsConfig::default(),
            )
            .context(format!(
                "Error while running bcftools keep PASS for {}",
                &self.clair3_germline_passed
            ))?;
            let log_file = format!("{}/bcftools_pass_", self.log_dir);
            report
                .save_to_file(&log_file)
                .context(format!("Error while writing logs into {log_file}"))?;

            // fs::remove_file(&tmp_file).context(format!("Can't remove tmp file {tmp_file}"))?;
        } else {
            info!("ClairS germline VCF exists.");
        }


        if !Path::new(&self.vcf_passed).exists() {
            // Concat output and indels
            let tmp_file = format!("{}.vcf.gz", get_temp_file_path()?.display());
            let report = bcftools_concat(
                vec![
                    self.output_vcf.to_string(),
                    self.output_indels_vcf.to_string(),
                ],
                &tmp_file,
                BcftoolsConfig::default(),
            )
            .context(format!(
                "Error while running bcftools concat for {} and {}",
                &self.output_vcf, &self.output_indels_vcf
            ))?;

            let log_file = format!("{}/bcftools_concat_", self.log_dir);
            report
                .save_to_file(&log_file)
                .context(format!("Error while writing logs into {log_file}"))?;

            let report = bcftools_keep_pass(&tmp_file, &self.vcf_passed, BcftoolsConfig::default())
                .context(format!(
                    "Error while running bcftools keep PASS for {}",
                    &self.output_vcf
                ))?;
            let log_file = format!("{}/bcftools_pass_", self.log_dir);
            report
                .save_to_file(&log_file)
                .context(format!("Error while writing logs into {log_file}"))?;

            fs::remove_file(&tmp_file).context(format!("Can't remove tmp file {tmp_file}"))?;
        } else {
            info!("ClairS PASSED vcf exists.");
        }

        Ok(())
    }
}

impl Variants for ClairS {
    fn variants(&self, annotations: &Annotations) -> anyhow::Result<VariantCollection> {
        let add = vec![Annotation::Callers(Caller::ClairS), Annotation::Somatic];
        info!("Loading variant from ClairS {} with annotations: {:?}", self.id, add);
        let variants = read_vcf(&self.vcf_passed)?;

        variants.par_iter().for_each(|v| {
            annotations.insert_update(v.hash_variant(), &add);
        });
        Ok(VariantCollection {
            variants,
            vcf: Vcf::new(self.vcf_passed.clone().into())?,
        })
    }
}

impl ClairS {
    pub fn germline(&self, annotations: &Annotations) -> anyhow::Result<VariantCollection> {
        let add = vec![Annotation::Callers(Caller::ClairS), Annotation::Germline];
        let variants = read_vcf(&self.clair3_germline_passed)?;
        variants.par_iter().for_each(|v| {
            annotations.insert_update(v.hash_variant(), &add);
        });

        Ok(VariantCollection {
            variants,
            vcf: Vcf::new(self.clair3_germline_passed.clone().into())?,
        })
    }
}

impl RunnerVariants for ClairS {}
