import os
import gzip
import shutil
import logging
import multiprocessing
import pandas as pd

from typing import Optional
from tqdm import tqdm
from multiprocessing import Pool
from pathlib import Path
from EnteroTyper.bin.accessories import run_subprocess

script = os.path.basename(__file__)


def type_sample(input_assembly: Path, database: Path, outdir: Path, create_db: bool, keep_blast: bool,
                sample_name: str = None):
    """
    Function wrapping all of the functionality of the enterobase_typer script. Can be called directly.
    """

    # Output directory creation/validation
    os.makedirs(str(outdir), exist_ok=False)
    logging.debug(f"Created directory {outdir}")

    logging.debug(f"input_assembly: {input_assembly}")
    logging.debug(f"database: {database}")
    logging.debug(f"outdir: {outdir}")
    logging.debug(f"create_db: {create_db}")

    # Call makeblastdb on each database file if create_db=True
    if create_db:
        database_files = makeblastdb_database(database=database)
    else:
        database_files = get_database_files(database=database)

    # Query input_assembly against each loci in the cgMLST database
    df = multiprocess_blastn_call(database_files, input_assembly, outdir)

    # Grab sample name from input_assembly
    if sample_name is None:
        sample_name = input_assembly.with_suffix("").name.replace(".pilon", "")

    # Prepare detailed report
    detailed_report = generate_detailed_report(df=df, outdir=outdir, sample_name=sample_name)

    # Prepare cgMLST report
    cgmlst_allele_report = generate_cgmlst_report(df=df, outdir=outdir, sample_name=sample_name)

    logging.info(f"cgMLST Allele Report: {cgmlst_allele_report}")
    logging.info(f"Detailed Report: {detailed_report}")

    if keep_blast:
        # Move BLASTn files
        move_blastn_files(indir=outdir, outdir=(outdir / 'blastn_output'))
    else:
        # Delete all blastn files
        delete_blastn_files(indir=outdir)
    return detailed_report


def move_blastn_files(indir: Path, outdir: Path):
    """
    Moves all of the BLASTn output to a new folder 'blastn_output'
    """
    os.makedirs(str(outdir), exist_ok=True)
    blastn_files = list(indir.glob("*.BLASTn"))
    for f in blastn_files:
        shutil.move(str(f), str(outdir / f.name))


def delete_blastn_files(indir: Path):
    blastn_files = list(indir.glob("*.BLASTn"))
    [os.remove(str(f)) for f in blastn_files]


def generate_cgmlst_report(df: pd.DataFrame, outdir: Path, sample_name: str) -> Path:
    """
    Generates the cgMLST report along with a transposed version within the out_dir folder
    """
    cgmlst_allele_report = outdir / f"{sample_name}_cgMLST_Allele_Report.tsv"
    cgmlst_df = get_sequence_type(df)
    cgmlst_df_transposed = cgmlst_df.transpose()
    cgmlst_df_transposed.to_csv(cgmlst_allele_report, sep="\t", header=False)
    return cgmlst_allele_report


def generate_detailed_report(df: pd.DataFrame, outdir: Path, sample_name: str) -> Path:
    """
    Generates the detailed report within the out_dir folder
    """
    output_detailed_report = outdir / f"{sample_name}_BLASTn_Detailed_Report.tsv"
    df.to_csv(output_detailed_report, sep="\t", index=None)
    return output_detailed_report


def get_database_files(database: Path, loci_suffix: str = "*.gz") -> Optional[list]:
    """
    Grabs database files within a provided directory and returns a list of everything present.
    Matches against a given suffix (defaults to *.gz, the expected database file extension)
    """
    database_files = list(database.glob(loci_suffix))
    if len(database_files) > 0:
        return database_files
    else:
        raise EmptyDatabase(f"Could not find database files at {database}. "
                            f"Try re-running the script with the --create_db flag.")


def multiprocess_blastn_call(database_files: list, input_assembly: Path, outdir: Path) -> [pd.DataFrame]:
    """
    Calls blastn on every database file against the input_assembly, reads top hit of results into a dataframe per loci,
    then returns a single dataframe of the combined results
    """
    df_params = [(database_file, input_assembly, outdir) for database_file in database_files]

    with Pool(multiprocessing.cpu_count() - 1) as p:
        df_list = p.starmap(closest_allele_df, df_params)
    df = combine_dataframes(df_list)

    # Sort values
    df = df.sort_values(by=['sseqid'])

    # Get reverse complement of minus strands
    df['qseq_strand_aware'] = df.apply(get_reverse_complement_row, axis=1)

    # Drop extraneous columns
    df = df.drop(['qseq', 'sstrand', 'score'], axis=1)

    # Sort again
    df = df.sort_values(by=['locus'])

    return df


def closest_allele_df(database_file: Path, input_assembly: Path, outdir: Path) -> pd.DataFrame:
    """
    Queries the input_assembly against a given database file, returns the top hit as a dataframe
    """
    database_file = database_file.with_suffix(".blastDB")
    blastn_file, locus_name = call_blastn(database_file=database_file, query_fasta=input_assembly, out_dir=outdir)
    df = parse_blastn(blastn_file)
    df = process_blastn_df(df, locus_name)
    return df


def get_sequence_type(df: pd.DataFrame) -> pd.DataFrame:
    loci = df['locus'].unique()
    sequence_type_dict = {}
    for locus in loci:
        sequence_type_df = df[df['locus'] == locus]
        sequence_type = sequence_type_df['sseqid'].iloc[0]
        if str(sequence_type) != 'nan':
            sequence_type = sequence_type.rsplit('_', 1)[-1]
        else:
            sequence_type = "NA"
        sequence_type_dict[locus] = sequence_type
    df = pd.DataFrame.from_dict(sequence_type_dict, orient='index').transpose()
    return df


def process_blastn_df(df: pd.DataFrame, locus_name: str) -> pd.DataFrame:
    """
`   Sorts and filters the DataFrame according to pident and a newly calculated lratio, evaluates top hits, then
    returns a DataFrame with only the top hit along with a 'hit_type' classification.
    :param df: DataFrame containing parsed *.BLASTn results
    :param locus_name: Name of the locus represented in the *.BLASTn file
    :return: df containing only the top hit from the parsed *.BLASTn
    """
    logging.debug(f"Processing locus: {locus_name}...")
    # Filter by length to slen ratio
    df['lratio'] = df['length'] / df['slen']
    df['locus'] = locus_name

    # Filter junk from DataFrame
    df = df.query("lratio >= 0.90 & pident >= 90.0 & lratio <=1.10")

    # Sort values so the best hits are at the top
    df = df.sort_values(["pident", "lratio"], ascending=False)
    df = df.reset_index(drop=True)

    # TODO: Revisit this and rework it. The hit classification is not great.
    hit_type = None
    if len(df) == 1:
        # CASE 1: PERFECT SINGLE HIT
        if df.pident[0] == 100.0 and df.lratio[0] == 1.0:
            hit_type = "PERFECT_HIT"
    elif len(df) > 1 and df.pident[0] == 100.0 and df.lratio[0] == 1.0:
        # CASE 2: PERFECT DUPLICATE - top two hits both have 100% pident and lratio
        if df.pident[1] == 100.0 and df.lratio[1] == 1.0:
            hit_type = "PERFECT_DUPLICATE"
        # CASE 3: PERFECT HIT w/very close second hit
        elif (0.99 <= df.lratio[1] <= 1.01) or df.pident[1] >= 0.99:
            hit_type = "PERFECT_HIT_WITH_POTENTIAL_PARALOG"
        # CASE 4: PERFECT HIT w/slightly less close second hit
        else:
            hit_type = "PERFECT_HIT"
    elif len(df) == 0:  # no close hits
        hit_type = "NO_MATCH"
        data = []
        data.insert(0, {"qseqid": "NA", "sseqid": "NA", "slen": "NA",
                        "length": "NA", "qstart": "NA", "qend": "NA",
                        "pident": "NA", "score": "NA", "locus": locus_name,
                        "sstrand": "NA", "qseq": "NA", "lratio": "NA"})
        df = pd.concat([pd.DataFrame(data), df], ignore_index=True, sort=False)
    else:
        hit_type = "CLOSEST_HIT"

    df["hit_type"] = hit_type
    return df.head(1)


def call_makeblastdb(db_file: Path):
    """
    Makes a system call to makeblastdb on a given database file. Can handle *.gz, *.fasta, or no suffix.
    :param db_file: Path to database file
    """
    db_name = db_file.with_suffix(".blastDB")
    if db_file.suffix == ".gz":
        cmd = f"gunzip -c {db_file} | makeblastdb -in - -parse_seqids -dbtype nucl -out {db_name} -title {db_name}"
        run_subprocess(cmd, get_stdout=True)
    elif db_file.suffix == ".fasta":
        cmd = f"makeblastdb -in {db_file} -parse_seqids -dbtype nucl -out {db_name} -title {db_name}"
        run_subprocess(cmd, get_stdout=True)
    elif db_file.suffix == "":
        os.rename(str(db_file), str(db_file.with_suffix(".fasta")))
        cmd = f"makeblastdb -in {db_file} -parse_seqids -dbtype nucl -out {db_name} -title {db_name}"
        run_subprocess(cmd, get_stdout=True)
    else:
        logging.debug("Invalid file format provided to call_makeblastdb()")


def makeblastdb_database(database: Path, loci_suffix: str = "*.gz") -> list:
    """
    Calls makeblastdb on every *.gz file in a given directory. Intended to function with the files retrieved via the
    EnterobasePull script (https://github.com/bfssi-forest-dussault/EnterobasePull).
    :param database: Path to database retrieved with EnterobasePull
    :param loci_suffix: Suffix of files to run makeblastdb on
    """
    logging.debug(f"Calling makeblastdb on database at {database}")
    db_files = list(database.glob(loci_suffix))
    for f in tqdm(db_files):
        call_makeblastdb(f)
    return db_files


def get_allele_length(database_file: Path) -> int:
    with gzip.open(database_file, "rt") as handle:
        lines = [line.strip() for line in handle]
        for line in lines:
            if line[0] != ">":
                return len(line)


def call_blastn(database_file: Path, query_fasta: Path, out_dir: Path) -> (Path, str):
    """
    Calls blastn against a query *.fasta file using the provided Enterobase DB target, dumps output into out_dir
    :param database_file: Formatted (makeblastdb) BLAST database file
    :param query_fasta: Sequence to query against database file
    :param out_dir: Path to directory to store output
    :return: BLASTn output (Path) and the base name of the target being searched (str)
    """
    # Dynamically set word_size
    allele_length = get_allele_length(database_file=database_file.with_suffix(".gz"))
    word_size = round(allele_length / 4)

    locus_name = database_file.with_suffix('').name
    reference_name = query_fasta.with_suffix('').name
    out_file = out_dir / Path(reference_name + "." + locus_name + ".BLASTn")
    cmd = f"blastn -query {query_fasta} -db {database_file} -out {out_file} " \
        f"-outfmt '6 qseqid sseqid slen length qstart qend pident score sstrand qseq' -word_size {word_size}"
    run_subprocess(cmd)
    return out_file, locus_name


def combine_dataframes(dfs: [pd.DataFrame]) -> pd.DataFrame:
    """
    Receives a list of DataFrames and concatenates them. They must all have the same header.
    :param dfs: List of DataFrames
    :return: Single concatenated DataFrame
    """
    df = pd.concat(dfs, sort=False)
    return df


def parse_blastn(blastn_file: Path) -> pd.DataFrame:
    """
    Parses *.BLASTn file generated by call_blastn(), then returns the df
    :param blastn_file: file path to *.BLASTn file
    :return: DataFrame contents of *.BLASTn file
    """
    headers = ["qseqid", "sseqid", "slen", "length", "qstart", "qend", "pident", "score", "sstrand", "qseq"]
    df = pd.read_csv(blastn_file, delimiter="\t", names=headers)
    return df


def get_reverse_complement_row(row) -> str:
    """
    Takes DataFrame row and returns the reverse complement of the sequence if the strand is 'minus'
    :param row:
    :return:
    """
    complement_dict = {'A': 'T',
                       'C': 'G',
                       'G': 'C',
                       'T': 'A'}
    sequence = row['qseq'].upper()
    if row['sstrand'] == 'minus':
        reverse_complement = "".join(complement_dict.get(base, base) for base in reversed(sequence))
        return reverse_complement
    else:
        return sequence


class EmptyDatabase(Exception):
    pass
