from pathlib import Path
import shutil
import pandas as pd
import numpy as np
import pysam
import subprocess

from metacarp.app.command import Command
from metacarp.app.utils.loggingutils import logger

def extract_kraken_info(kraken_report: Path, cutoff_prok_fungi: float, cutoff_euk: float, cutoff_allergens: float, usual_suspects: Path, allergen_file: Path)-> pd.DataFrame:
    """
    - Get a list of taxids from the kraken results according to the cutoffs for prokaryotes/fungi and eukaryotes.
    - get info on unclassified reads
    - get info on presence of allergens
    :param kraken_report: kraken report file
    :param cutoff_prok_fungi: cutoff for prokaryotes and fungi
    :param cutoff_euk: cutoff for eukaryotes
    :param cutoff_allergens: cutoff for allergens
    :param usual_suspects: Path to file with usual suspects to include
    :param allergen_file: Path to file with allergens to screen for
    :return: pandas dataframes with unclassified reads, allergens, and selected taxids
    """
    column_names = ['abundance', 'frags1', 'frags2', 'rank', 'taxid', 'name']
    taxa = pd.read_csv(kraken_report, sep='\t', names=column_names)
    taxa['name'] = taxa['name'].str.strip()

    sample = Path(kraken_report).parent.name
    usual = pd.read_csv(usual_suspects, sep='\t')
    allergens_df = pd.read_csv(allergen_file, sep='\t')

    # Screen for unclassified reads
    unclassified_mask = taxa['rank'] == 'U'
    unclassified = taxa.loc[unclassified_mask, ['abundance', 'frags1']].copy()
    unclassified.index = [sample]
    unclassified.columns = ['abundance', 'reads']

    # Screen for allergens
    allergen_mask = (taxa['taxid'].isin(allergens_df['taxid'].tolist())) & (taxa['abundance'] >= cutoff_allergens)
    allergens = taxa.loc[allergen_mask, ['name', 'abundance', 'frags1']].copy()
    allergens.columns = ['species', 'abundance', 'reads']
    allergens = allergens.set_index('species')

    # Assign groups
    conditions = [
        (taxa['rank'] == 'D') & (taxa['name'] == 'Bacteria'),
        (taxa['rank'] == 'K') & (taxa['name'] == 'Metazoa'),
        (taxa['rank'] == 'K') & (taxa['name'] == 'Viridiplantae'),
        (taxa['rank'] == 'K') & (taxa['name'] == 'Fungi'),
        (taxa['rank'] == 'D') & (taxa['name'] == 'Viruses')
    ]
    choices = ['Bacteria', 'Animals', 'Plants', 'Fungi', 'Viruses']
    taxa['group'] = np.select(conditions, choices, default=None)

    # Forward fill group to species level
    taxa['group'] = taxa['group'].ffill()

    # Filter species level
    species_mask = taxa['rank'] == 'S'
    prok_fungi_mask = (taxa['group'].isin(['Bacteria', 'Fungi'])) & \
                      (taxa['abundance'] >= cutoff_prok_fungi) & \
                      (~taxa['name'].str.contains(' sp. ', na=False))
    euk_mask = (taxa['group'].isin(['Animals', 'Plants'])) & \
               (taxa['abundance'] >= cutoff_euk)
    usual_suspects_mask = taxa['taxid'].isin(usual['taxid'].tolist())

    final_mask = species_mask & (prok_fungi_mask | euk_mask | usual_suspects_mask)

    taxids = taxa.loc[final_mask, ['name', 'taxid', 'group', 'abundance']].copy()
    taxids.columns = ['species', 'taxid', 'group', 'abundance']
    taxids = taxids.set_index('species')

    # Add missing usual suspects
    usual['abundance'] = np.nan
    missing_usual = usual[~usual['taxid'].isin(taxids['taxid'])]
    all_taxids = pd.concat([taxids.reset_index(), missing_usual])

    return all_taxids, allergens, unclassified

def get_taxids(all_taxids: Path) -> set:
    """
    Get list of taxids
    :param all_taxids: species detected above cutoff with kraken2 + usual suspects
    :return: list of unique taxids
    """
    taxids = pd.read_csv(all_taxids, sep='\t')
    return taxids.taxid.astype(str).tolist()

def download_genomes(all_taxids: Path, dir_intermediary: Path, dir_genomes: Path) -> None:
    """
    Download reference genomes for a list of taxids
    :param all_taxids: species detected above cutoff with kraken2 + usual suspects
    :param dir_intermediary: directory for intermediary download files
    :param dir_genomes: directory to collect downloaded genomes
    :return: None
    """
    taxids = get_taxids(all_taxids)
    dir_intermediary.mkdir(parents=True, exist_ok=True)

    for taxid in taxids:
        logger.info(f'Getting genome for {taxid}')
        taxid_dir = dir_intermediary / taxid
        taxid_dir.mkdir(exist_ok=True)

        outfile = taxid_dir / 'dataset.zip'
        try:
            subprocess.run(
                ['datasets', 'download', 'genome', 'taxon', taxid,
                    '--reference',
                    '--filename', str(outfile)],
                capture_output=True, check=True)

            command = Command(' '.join([
                'unzip',
                '-qq',
                f'-o {dir_intermediary.absolute()}/{taxid}/dataset.zip',
                f'-d {dir_intermediary.absolute()}/{taxid}'
                ]))
            command.run(dir_intermediary.parent)

            # sort in reverse order so refseq (GCF..) will be used if available, otherwise genbank (GCA..)
            command = Command(' '.join([
                f'ref=$(find {dir_intermediary.absolute()}/{taxid}/ncbi_dataset/data/GC*/*.fna | sort -r | head -n 1)',
                '&&',
                f'cp $ref {dir_genomes.absolute()}/{taxid}.fna'
            ]))
            command.run(dir_genomes.parent)

        except subprocess.CalledProcessError as e:
            logger.info(f'Error occurred: {e}.')
            #logger.info(f'Error output: {e.stderr}.')
            logger.info(f'Warning: genome download failed for {taxid}.')

        # Clean up only the specific taxid subdirectory, not the entire intermediary directory
        if taxid_dir.exists():
            shutil.rmtree(taxid_dir)

def make_bed_files(all_taxids: Path, dir_genome: Path) -> None:
    """
    Make bed files for reference genomes
    :param all_taxids: species detected above cutoff with kraken2 + usual suspects
    :param dir_genome: directory with downloaded genomes
    :return: None
    """
    taxids = get_taxids(all_taxids)
    for taxid in taxids:
        try:
            ref_genome = pysam.FastaFile(str(list(dir_genome.glob(taxid + '*.fna'))[0]))
        except (IndexError, FileNotFoundError, OSError):
            print('Warning: skipping taxid ' + str(taxid), flush=True)
            continue
        region_names = []
        region_lens = []
        for seq_name in ref_genome.references:
            region_names.append(seq_name)
            region_lens.append(ref_genome.get_reference_length(seq_name))

        bed_file = dir_genome / f'{taxid}.bed'
        with open(str(bed_file), 'w') as file:
            for region_name, region_len in zip(region_names, region_lens):
                line = f"{region_name}\t0\t{region_len}\n"
                file.write(line)

def split_bam(all_taxids: Path, bam: Path, reads_basename: Path, dir_genomes: Path) -> None:
    """
    Split alignment to concatenated reference into separate bam files per taxid
    :param all_taxids: species detected above cutoff with kraken2 + usual suspects
    :param bam: bam file with alignments to all reference genomes
    :param reads_basename: directory to save split bam files
    :param dir_genomes: directory with downloaded genomes
    :return: None
    """
    taxids = get_taxids(all_taxids)
    bamfile = pysam.AlignmentFile(bam, 'rb')
    genome_region_dict = {}
    genome_bamfiles = {}
    for taxid in taxids:
        region_names = []
        region_file = dir_genomes / f'{taxid}.bed'
        # create dict with reference genomes and their respective regions
        try:
            with open(str(region_file), 'r') as file:
                for line in file:
                    columns = line.strip().split()
                    region_names.append(columns[0])
        except (FileNotFoundError, IndexError, OSError):
            print("Warning: Skipped {0} because file was missing".format(taxid), flush=True)
            continue
        genome_region_dict[taxid] = region_names

    region_to_taxid = {}
    for taxid, regions in genome_region_dict.items():
        for region in regions:
            region_to_taxid[region] = taxid

        # create a bam file for this genome, the header has to be identical to the original bam file for the splitting to work
        bamfile_name = reads_basename / f'{taxid}.bam'
        genome_bamfiles[taxid] = pysam.AlignmentFile(str(bamfile_name), 'wb', template=bamfile)

    for read in bamfile:
        if not read.is_unmapped:
            reference = read.reference_name
            taxid = region_to_taxid.get(reference)
            if taxid:
                genome_bamfiles[taxid].write(read)

def sort_index_stats(taxid: str, reads_basename: Path, tmp_dir: Path, dir_genomes: Path) -> set:
    """
    Calculate alignment metrics (depth and breadth of coverage and number of mapped reads) for a given taxid
    :param taxid: taxid to calculate alignment metrics for
    :param reads_basename: directory with split bam files
    :param dir_genomes: directory with downloaded genomes
    :return: tuple of (depth of coverage, breadth of coverage, number of mapped reads)
    """
    refgen_bam = str(list(reads_basename.glob(taxid + '.bam'))[0])
    # sort bam file
    command = Command(' '.join([
        'samtools sort',
        refgen_bam,
        '-T', f'{tmp_dir}',
        '>',
        f'{reads_basename}/{taxid}.sorted.bam',
    ]))
    command.run(tmp_dir)

    # calculate the amount of mapped reads
    command = Command(' '.join([
        'samtools view',
        '-F 0x4',
        f'{reads_basename}/{taxid}.sorted.bam',
        '| cut -f1',
        '| sort',
        '| uniq',
        '| wc -l',
    ]))
    command.run(tmp_dir)
    reads = int(command.stdout)

    # calculate depth and breadth of coverage if there are at least 5 mapped reads:
    if reads > 4:
        depthfile = taxid + '_depth.tsv'
        command = Command(' '.join([
            'samtools depth',
            '-aa',
            '-b', f'{dir_genomes}/{taxid}.bed',
            f'{reads_basename}/{taxid}.sorted.bam',
            '>',
            f'{reads_basename}/{depthfile}'
        ]))
        command.run(tmp_dir)

# Calculate depth using awk (constant memory usage)
        command = Command(
            f"awk '{{sum+=$3; count++; if($3>0) cov++}} END {{print sum/count, (cov/count)*100}}' {reads_basename}/{depthfile}")
        command.run(tmp_dir)
        depth_str, breadth_str = command.stdout.strip().split()
        depth = float(depth_str)
        breadth = float(breadth_str)
    else:
        depth = 0.0
        breadth = 0.0

    return depth, breadth, reads

def sanitize_input_name(name: str, extension: str) -> str:
    """
    Sanitizes the input file name.
    :param name: Name
    :param extension: Expected file extension (e.g., 'bam' or 'fasta')
    :return: None
    """
    invalid_chars = '/!@#$\\'

    # Replace spaces by dashes
    name = name.replace(' ','_')

    # Avoid double dot before the extension
    if name.endswith('.'):
        name = name[:-1]

    # Add extension
    if not name.endswith(f'.{extension}'):
        return f'{name}.{extension}'

    # Return sample name without invalid characters
    return ''.join(c for c in name if c not in invalid_chars)