import argparse
import os
import subprocess
import sys
from os.path import getsize
import pandas as pd
from functools import reduce
from Bio import SeqIO
from tqdm import tqdm
import datetime
import time
import mvip

def fetch_arguments(parser):
    parser.set_defaults(func=main)
    parser.set_defaults(program="MVP_00_set_up_MVP")
    parser.add_argument(
        "--input", "-i",
        type=str,
        help="Path to the working directory where MVP will be run.",
    )
    parser.add_argument(
        "--metadata", "-m",
        type=str,
        help="Path to the metadata that will be use to run MVP.",
    )
    parser.add_argument(
        "--skip_install_databases",
        action="store_true",
        default=False,
        help="Install geNomad and CheckV databases in the respective repositories. (default = False).",
    )
    parser.add_argument(
        '--genomad_db_path', 
        type=str,
        help='Path to the directory where geNomad database will be installed.'
        )
    parser.add_argument(
        '--checkv_db_path', 
        type=str,
        help='Path to the directory where CheckV database will be installed.'
        )
    parser.add_argument(
        "--skip_check_errors",
        action="store_true",
        default=False,
        help="Skip to run sequence data error checking. (default = False).",
    )
    
# Global variable to track the current step number
current_step = 0
# Dictionary to track sub-step counts for each function
sub_step_counts = {}

def print_step(func_name):
    global current_step
    # If it's the first call to this function, increment the step number
    if func_name not in sub_step_counts:
        current_step += 1
        sub_step_counts[func_name] = 1
        return f"Step {current_step}"  # Return without a letter for the first call
    else:
        # Increment the sub-step count for subsequent calls within the same function
        sub_step_counts[func_name] += 1
        sub_step_letter = chr(64 + sub_step_counts[func_name])  # 'A' is 65 in ASCII
        return f"Step {current_step}{sub_step_letter}"


def convert_metadata_to_unix(args):
    print(f"\n\033[1m{print_step(convert_metadata_to_unix)}: Converting the metadata file...\033[0m")

    subprocess.run(['dos2unix', args["metadata"]])

def check_metadata(args):
    print(f"\n\033[1m{print_step(check_metadata)}: Checking for missing, incorrectly named mandatory columns or extra columns in the metadata file...\033[0m")

    # Read the metadata file
    metadata = pd.read_csv(args["metadata"], sep='\t')

    # Define the mandatory and optional columns
    mandatory_columns = ['Sample_number', 'Sample', 'Assembly_Path', 'Read_Path']
    optional_columns = ['Variable']

    missing_columns = [col for col in mandatory_columns if col not in metadata.columns]
    if missing_columns:
        print(f"Error: Missing mandatory column(s) in metadata: {', '.join(missing_columns)}")
        sys.exit(1)

    # Check for extra columns
    extra_columns = [col for col in metadata.columns if col not in mandatory_columns + optional_columns]
    if extra_columns:
        print(f"Error: Extra column(s) found in metadata: {', '.join(extra_columns)}")
        sys.exit(1)

    print("The metadata looks fine. Proceeding to the next steps.")

def check_paths(args):
    print(f"\n\033[1m{print_step(check_paths)}: Checking if the assembly or read file paths provided in the metadata are correct...\033[0m")

    metadata = pd.read_csv(args["metadata"], sep='\t')
    nonexistent_assembly_paths = metadata.loc[~metadata['Assembly_Path'].apply(lambda path: path == '' or (isinstance(path, str) and os.path.exists(path))), 'Assembly_Path']
    nonexistent_read_paths = metadata.loc[~metadata['Read_Path'].apply(lambda path: path == '' or (isinstance(path, str) and os.path.exists(path))), 'Read_Path']

    if not nonexistent_assembly_paths.empty or not nonexistent_read_paths.empty:
        if not nonexistent_assembly_paths.empty:
            print(f"Error: The following Assembly_Path values correspond to non-existent files: {', '.join(nonexistent_assembly_paths)}")
        if not nonexistent_read_paths.empty:
            nonexistent_read_paths_str = nonexistent_read_paths.astype(str)
            print(f"Error: The following Read_Path values correspond to non-existent files: {', '.join(nonexistent_read_paths_str)}")
        sys.exit(1)
    
    print("The assembly and read file paths are correct. Proceeding to the next steps.")

def check_sequence_data(args):
    if args["skip_check_errors"]:
        print(f"\n\033[1m{print_step(check_sequence_data)}: Skipping verifying that the input assembly files contain only valid nucleotide characters and do not have duplicate headers...\033[0m")
        return
    
    print(f"\n\033[1m{print_step(check_sequence_data)}: Verifying that the input assembly files contain only valid nucleotide characters and does not have duplicate headers...\033[0m")
    metadata = pd.read_csv(args["metadata"], sep='\t')
    for row_index, (_, row) in enumerate(metadata.iterrows(), start=1):
        input_assembly_file = row['Assembly_Path']
        
        with open(input_assembly_file, "r") as input_fasta:
            headers_seen = set()
            invalid_chars_found = False
            duplicate_headers_found = False
            
            for record in SeqIO.parse(input_fasta, "fasta"):
                header = record.id
                if any(base not in "ACTGNactgn" for base in record.seq):
                    print(f"Error: Invalid nucleotide characters found in {input_assembly_file} for sequence {record.id}\n")
                    invalid_chars_found = True
                if header in headers_seen:
                    print(f"Error: Duplicate header found in {input_assembly_file}: {header}\n")
                    duplicate_headers_found = True
                headers_seen.add(header)
            
            if invalid_chars_found or duplicate_headers_found:
                sys.exit(1)

    print("The input assembly files look fine. Proceeding to the next steps.")

def create_output_folders(args):
    print(f"\n\033[1m{print_step(create_output_folders)}: Creating MVP directories where outputs will be stored...\033[0m")

    output_folders = ['00_DATABASES', '01_GENOMAD', '02_CHECK_V', '03_CLUSTERING', '04_READ_MAPPING', '05_VOTU_TABLES', '06_FUNCTIONAL_ANNOTATION']
    for folder_name in output_folders:
        output_folder = os.path.join(args["input"], folder_name)
        if not os.path.exists(output_folder):
            os.makedirs(output_folder)

def install_databases(args):
    if args["skip_install_databases"]:
        genomad_db_path = None
        checkv_db_path = None
        print(f"\n\033[1m{print_step(install_databases)}: Skipping to install databases because --install_databases is False...\033[0m")
    else:
        database_path = os.path.join(args["input"], '00_DATABASES')
        os.makedirs(database_path, exist_ok=True)

        # Install geNomad database
        if args['genomad_db_path']:
            genomad_db_path = os.path.join(args['genomad_db_path'], 'genomad_db')
        else:
            genomad_db_path = os.path.join(database_path, 'genomad_db')
        
        if os.path.exists(genomad_db_path):
            print(f"\n\033[1m{print_step(install_databases)}: Skipping to install geNomad database because {genomad_db_path} already exists...\033[0m")
        else:
            print(f"\n\033[1m{print_step(install_databases)}: Creating geNomad database...\033[0m")
            subprocess.run(['genomad', 'download-database', database_path])

        # Install CheckV database
        if args['checkv_db_path']:
            checkv_db_path = os.path.join(args['checkv_db_path'], 'checkv-db-v1.5')
        else:
            checkv_db_path = os.path.join(database_path, 'checkv-db-v1.5')

        if os.path.exists(checkv_db_path):
            print(f"\n\033[1m{print_step(install_databases)}: Skipping to install Check_V database because {checkv_db_path} already exists...\033[0m")
        else:
            print(f"\n\033[1m{print_step(install_databases)}: Creating CheckV database...\033[0m")
            subprocess.run(['checkv', 'download_database', database_path])

        additional_dbs = ['PhrogDB_v14', 'RdRP_DB', 'Pfam_A_DB', 'dbAPIS']
        for db in additional_dbs:
            db_path = os.path.join(database_path, db)
            if os.path.exists(db_path) and os.listdir(db_path):
                print(f"\n\033[1m{print_step(install_databases)}: Skipping install of {db} because it already exists...\033[0m")
            else:
                print(f"\n\033[1m{print_step(install_databases)}: Downloading {db}...\033[0m")
                subprocess.run(['wget', '-r', '-np', '-nH', '--cut-dirs=3', '--reject="index.html*"', f'https://portal.nersc.gov/cfs/m342/MViP/{db}/', '-P', database_path])

    return genomad_db_path, checkv_db_path

def generate_summary_report(args, genomad_db_path, checkv_db_path, formatted_start_time, formatted_end_time, duration):
    print(f"\n\033[1m{print_step(generate_summary_report)}: Generating summary report for Module 00...\033[0m")

    # Create a dictionary to hold argument descriptions and their default values
    argument_defaults = {
        '--input': args["input"],
        '--metadata': args["metadata"],
        '--skip_check_errors': args["skip_check_errors"],
        '--skip_install_databases': args["skip_install_databases"],
        '--genomad_db_path': genomad_db_path,
        '--checkv_db_path': checkv_db_path
    }

    # Create the module header
    module_header = """
****************************************************************************
*********                          MODULE 00                       *********
****************************************************************************
"""

    # Write a summary line with script arguments and their default values
    summary_line = "00_set_up_mvp.py"
    for arg, default in argument_defaults.items():
        if default is not None:
            summary_line += f" {arg} {default}"

    # Create a summary report file
    summary_report_path = os.path.join(args["input"], 'MVP_00_Summary_Report.txt')

    with open(summary_report_path, 'w') as summary_report:
        # Write the module header
        summary_report.write(module_header.strip() + '\n')
        summary_report.write(summary_line + '\n\n')
        summary_report.write(f"Start Time: {formatted_start_time}\n")
        summary_report.write(f"End Time: {formatted_end_time}\n")
        summary_report.write(f"Running Time: {duration.total_seconds():.2f} seconds\n\n")
        summary_report.write("Tools and versions:\n")

        # List of tools
        tools = [
            "tqdm", "genomad", "pandas", "checkv", "bowtie2", "minimap2", 
            "samtools", "seqtk", "scikit-learn", "numpy", "parallel", 
            "coverm", "python", "dos2unix", "mmseqs2", "mafft", "trimal", 
            "fasttree", "vRhyme", "taxopy", "table2asn"
        ]

        # For each tool, find it in conda list and get its version
        for tool in tools:
            try:
                result = subprocess.run(
                    ["conda", "list", tool], capture_output=True, text=True, check=True
                )
                lines = result.stdout.splitlines()
                if len(lines) > 3:
                    version_line = lines[3]  # Assuming the version info is in the 4th line
                    version = version_line.split()[1]
                    summary_report.write(f"{tool} {version}\n")
                else:
                    summary_report.write(f"{tool} version not found\n")
            except subprocess.CalledProcessError:
                summary_report.write(f"{tool} version not found\n")

        # Add section for databases and versions
        summary_report.write("\nDatabases and versions:\n")
        if genomad_db_path:
            genomad_db_name = os.path.basename(genomad_db_path)
            # Read version from version.txt
            version_file_path = os.path.join(genomad_db_path, 'version.txt')
            if os.path.exists(version_file_path):
                with open(version_file_path, 'r') as version_file:
                    genomad_version = version_file.read().strip()
                summary_report.write(f"{genomad_db_name} {genomad_version}\n")
            else:
                summary_report.write(f"{genomad_db_name}: version file not found\n")
        else:
            summary_report.write("geNomad database: not installed\n")
        
        if checkv_db_path:
            checkv_db_name = os.path.basename(checkv_db_path)
            # Extract version number from checkv_db_path
            checkv_version = checkv_db_name.split('-')[-1].replace('v', '')
            summary_report.write(f"{checkv_db_name} {checkv_version}\n")
        else:
            summary_report.write("CheckV database: not installed\n")

    print("Summary report for Module 00 has been created:", summary_report_path)

def main(args):
    # Capture start time
    start_time = datetime.datetime.now()
    formatted_start_time = start_time.strftime("%Y-%m-%d %H:%M:%S") 
    print(f"\nStart Time: {formatted_start_time}")

    # Usage of run_step for the dos2unix step
    convert_metadata_to_unix(args)

    check_metadata(args)

    # Usage of run_step for the check_paths step
    check_paths(args)

    # Usage of run_step for the check_sequence_data step
    check_sequence_data(args)

    create_output_folders(args)

    genomad_db_path, checkv_db_path = install_databases(args)

    # Capture end time
    end_time = datetime.datetime.now()
    formatted_end_time = end_time.strftime("%Y-%m-%d %H:%M:%S")

    # Calculate duration
    duration = end_time - start_time

    # Add this function to your script
    generate_summary_report(args, genomad_db_path, checkv_db_path, formatted_start_time, formatted_end_time, duration)

    message1 = "\033[1mModule 00 finished: you are all set to use MVP!\n\033[0m"
    message2 = "Your metadata looks fine, as well as your read and assembly files."
    message3 = "All the directories needed for MVP have been generated in your working directory."
    if not args["skip_install_databases"]:
        message4 = "geNomad and CheckV databases have been installed in their respective directories."
    message5 = "\n\033[1mYou can now proceed to the next step of the MVP script: Module 01!\033[0m"
    line_of_stars = '*' * len(message3)

    print()
    print(line_of_stars)
    print(message1)
    print(f"Start Time: {formatted_start_time}")
    print(f"End Time: {formatted_end_time}")
    print(f"Running Time: {duration.total_seconds():.2f} seconds\n")
    print(message2)
    print(message3)
    if not args["skip_install_databases"]:
        print(message4)
    print(message5)
    print(line_of_stars)
    print()

    print("Please don't forget to cite MVP and the following softwares used in this module:")
    print("- Camargo, A. P., Roux, S., Schulz, F., Babinski, M., Xu, Y., Hu, B., Chain, P. S. G., Nayfach, S., & Kyrpides, N. C. You can move, but you can’t hide: identification of mobile genetic elements with geNomad. bioRxiv (2023), DOI: 10.1101/2020.11.01.361691")
    print("\n- Nayfach, S., Camargo, A.P., Schulz, F. et al. CheckV assesses the quality and completeness of metagenome-assembled viral genomes. Nat Biotechnol 39, 578–585 (2021). https://doi.org/10.1038/s41587-020-00774-7\n")
