from modules.software_base import WorkflowRunnerBase
import os
from workflow.log import logcolor

class AdapterRemovalRunner(WorkflowRunnerBase):
    def __init__(self, logger, manifest_list, input_path="", output_path="", additional_args="", threads=2, stepname="", mode=None):
        super().__init__(
            logger=logger,
            manifest_list=manifest_list,
            input_path=input_path,
            output_path=output_path,
            additional_args=additional_args,
            threads=threads,
            stepname=stepname,
            software_name="AdapterRemoval",
            unneeded_args=["--file1", "--file2", "--basename", "--output1", "--output2", "--threads"],
            output_type="fastq",
            mode=mode
        )

    def find_input_files(self, prev_step_files):
        if "--collapse" in self.additional_args:
            for line in prev_step_files:
                if line["output2"]:
                    self.input_file_list.append({
                        "id": line["id"],
                        "input1": line["output1"],
                        "input2": line["output2"],
                        "output1": f"{os.path.join(self.output_path, line['id'])}.{self.output_type}",
                        "output2": None
                    })
                else:
                    self.input_file_list.append({
                        "id": line["id"],
                        "input1": line["output1"],
                        "input2": None,
                        "output1": f"{os.path.join(self.output_path, line['id'])}.{self.output_type}",
                        "output2": None
                    })
        else:
            for line in prev_step_files:
                if line["output2"]:
                    self.input_file_list.append({
                        "id": line["id"],
                        "input1": line["output1"],
                        "input2": line["output2"],
                        "output1": f"{os.path.join(self.output_path, line['id'])}_R1.{self.output_type}",
                        "output2": f"{os.path.join(self.output_path, line['id'])}_R2.{self.output_type}",
                    })
                else:
                    self.input_file_list.append({
                        "id": line["id"],
                        "input1": line["output1"],
                        "input2": None,
                        "output1": f"{os.path.join(self.output_path, line['id'])}.{self.output_type}",
                        "output2": None
                    })
        if self.leaf:
            return prev_step_files
        else:
            return self.input_file_list

    def build_command_queue(self):
        for item in self.input_file_list:
            base_cmd = [
                self.software_name,
                f"--basename {os.path.join(self.output_path, item['id'])}",
                f"--threads {self.threads}",
                self.additional_args
            ]

            # Some specific rules for adapterremoval
            # depending on the contents of the Manifest file

            # If collapsing to one output file
            if "--collapse" in self.additional_args:
                if item["input2"]:   # FWD and REV
                    base_cmd += [
                        f"--file1 {item['input1']}",
                        f"--file2 {item['input2']}",
                        f"--outputcollapsedtruncated {os.path.join(self.output_path, item['id'])}.truncated",
                        f"--outputcollapsed {os.path.join(self.output_path, item['id'])}.{self.output_type}"
                    ]
                else:               # FWD only
                    base_cmd += [
                        f"--file1 {item['input1']}",
                        f"--outputcollapsedtruncated {os.path.join(self.output_path, item['id'])}.truncated",
                        f"--outputcollapsed {os.path.join(self.output_path, item['id'])}.{self.output_type}"
                    ]
            
            # If not collapsing (two output files)
            else:
                if item["input2"]:   # FWD and REV
                    base_cmd += [
                        f"--output1 {item['output1']}",
                        f"--output2 {item['output2']}",
                        f"--file1 {item['input1']}",
                        f"--file2 {item['input2']}"
                    ]
                else:               # FWD only
                    base_cmd += [
                        f"--output1 {item['output1']}",
                        f"--file1 {item['input1']}"
                    ]

            self.command_queue.append(" ".join(base_cmd))