# Streamlit UI
import streamlit as st
from streamlit_extras.colored_header import colored_header
from streamlit_extras.stylable_container import stylable_container
from streamlit_file_browser import st_file_browser
import datetime
import glob
import gzip
import io
import os
import re
import shutil
import sqlite3
import subprocess
import zipfile
from concurrent.futures import ThreadPoolExecutor, as_completed
from io import StringIO
from pathlib import Path
import pandas as pd
from Bio import Entrez, SeqIO
from ete3 import NCBITaxa
from tqdm import tqdm
import os
import sys
import gzip
import shutil
import subprocess
from pathlib import Path
import tempfile

### Set page config first
st.set_page_config(page_title="APSCALE-blast Database Creator", page_icon="🧬", layout="wide")

Entrez.email = "macher@uni-trier.de"#
accession2taxonomy_parquet = Path('/Volumes/Coruscant/APSCALE_raw_databases/test/APSCALE_databases/accession_taxonomy.parquet.snappy')

## NCBI ACCESSION CONVERTER

def accession_to_taxonomy(accession):
    # Fetch the record summary from NCBI
    handle = Entrez.efetch(db="nucleotide", id=accession, rettype="gb", retmode="text")
    record = SeqIO.read(handle, "genbank")
    handle.close()

    # Try to get TaxID from features
    tax_id = None
    for feature in record.features:
        if feature.type == "source":
            tax_id = feature.qualifiers.get("db_xref", [None])[0]
            if tax_id:
                if tax_id.startswith("taxon:"):
                    tax_id = tax_id.split(":")[1]
                    break

    if not tax_id:
        return {"Accession": accession, "superkingdom": None, "phylum": None,
                "class": None, "order": None, "family": None, "genus": None, "species": None}

    # Fetch taxonomy from Taxonomy database
    handle = Entrez.efetch(db="taxonomy", id=tax_id, retmode="xml")
    tax_record = Entrez.read(handle)[0]
    handle.close()

    # Build a dict of taxonomy
    lineage = {d["Rank"]: d["ScientificName"] for d in tax_record["LineageEx"]}

    result = {
        "Accession": accession,
        "superkingdom": lineage.get("superkingdom"),
        "phylum": lineage.get("phylum"),
        "class": lineage.get("class"),
        "order": lineage.get("order"),
        "family": lineage.get("family"),
        "genus": lineage.get("genus"),
        "species": tax_record.get("ScientificName")
    }

    return result

def make_blast_db(fasta_file, db_folder, title="db", dbtype="nucl"):

    fasta_file = Path(fasta_file)
    db_folder = Path(db_folder)
    db_folder.parent.mkdir(parents=True, exist_ok=True)

    is_gz = fasta_file.suffix == ".gz"
    is_windows = os.name == "nt"

    # -------------------------
    # WINDOWS: unzip first
    # -------------------------
    if is_windows and is_gz:
        with tempfile.NamedTemporaryFile(delete=False, suffix=".fasta") as tmp:
            tmp_fasta = Path(tmp.name)

        with gzip.open(fasta_file, "rb") as f_in, open(tmp_fasta, "wb") as f_out:
            shutil.copyfileobj(f_in, f_out)

        cmd = [
            "makeblastdb",
            "-in", str(tmp_fasta),
            "-title", title,
            "-dbtype", dbtype,
            "-out", str(db_folder)
        ]

        subprocess.run(cmd, check=True)
        tmp_fasta.unlink()

    # -------------------------
    # UNIX: stream via zcat
    # -------------------------
    elif not is_windows and is_gz:
        cmd = (
            f"zcat < {fasta_file} | "
            f"makeblastdb -in - -title {title} -dbtype {dbtype} -out {db_folder}"
        )
        subprocess.run(cmd, shell=True, check=True)

    # -------------------------
    # Plain FASTA (all systems)
    # -------------------------
    else:
        cmd = [
            "makeblastdb",
            "-in", str(fasta_file),
            "-title", title,
            "-dbtype", dbtype,
            "-out", str(db_folder)
        ]
        subprocess.run(cmd, check=True)

def get_fasta_headers(fasta_file):
    fasta_file = Path(fasta_file)
    open_func = gzip.open if fasta_file.suffix == ".gz" else open

    headers = []
    with open_func(fasta_file, "rt") as fh:
        for line in fh:
            if line.startswith(">"):
                headers.append(line[1:].strip())

    return headers

def collect_all_taxonomy_dfs():
    files = glob.glob('/Volumes/Coruscant/APSCALE_raw_databases/test/tmp/*/db_taxonomy.parquet.snappy')
    df = pd.concat((pd.read_parquet(f) for f in files),ignore_index=True)
    df.to_parquet('/Volumes/Coruscant/APSCALE_raw_databases/test/APSCALE_databases/accession_taxonomy.parquet.snappy', compression='snappy')

def collect_accession_taxonomy(fasta_file, accession2taxonomy_parquet):
    TAX_COLS = [
        'Accession', 'superkingdom', 'phylum',
        'class', 'order', 'family', 'genus', 'species'
    ]

    if accession2taxonomy_parquet.exists():
        accession_df = pd.read_parquet(accession2taxonomy_parquet)
    else:
        accession_df = pd.DataFrame(columns=TAX_COLS)

    existing_accessions = set(accession_df['Accession'])

    fasta_accessions = set(get_fasta_headers(fasta_file))

    missing = fasta_accessions - existing_accessions
    cached = fasta_accessions & existing_accessions

    cached_df = accession_df[
        accession_df['Accession'].isin(cached)
    ]

    new_rows = []
    for acc in tqdm(missing, desc="Fetching taxonomy"):
        row = accession_to_taxonomy(acc)
        if row:
            new_rows.append(row)

    new_df = pd.DataFrame(new_rows, columns=TAX_COLS)

    out_df = pd.concat(
        [accession_df, new_df],
        ignore_index=True
    ).drop_duplicates(subset='Accession')

    return out_df

## DIATBARCODE

def run_diat_barcode(output_path, diat_barcode_xlsx):

    st.write('{} : Starting to collect accession numbers from .xlsx file.'.format(datetime.datetime.now().strftime('%H:%M:%S')))

    diat_barcode_df = pd.read_excel(diat_barcode_xlsx, sheet_name='diatbarcode v12').fillna('')

    st.write('{} : Writing fasta file.'.format(datetime.datetime.now().strftime('%H:%M:%S')))

    fasta_file = diat_barcode_xlsx.replace('.xlsx', '.fasta.gz').replace(' ', '_')
    with gzip.open(fasta_file, 'wt') as f:
        for line in diat_barcode_df[['Sequence ID', 'Sequence']].values.tolist():
            if line[0] != '':
                f.write(f'>{line[0]}\n')
                f.write(f'{line[1]}\n')

    st.write('{} : Finished to convert fasta format.'.format(datetime.datetime.now().strftime('%H:%M:%S')))
    st.write('{} : Starting to generate taxonomy file.'.format(datetime.datetime.now().strftime('%H:%M:%S')))

    records = []
    for line in diat_barcode_df[["Species", "Genus", "Family (following Round, Crawford & Mann 1990)", "Order (following Round, Crawford & Mann 1990)", "Class (following Round, Crawford & Mann 1990)", "Phylum (following Algaebase 2018)", "Subkingdom (following Algaebase 2018)", "Sequence ID"]].values.tolist():
        if line[0] != '':
            records.append(line[::-1])
    records_df = pd.DataFrame(records, columns=['Accession', 'superkingdom', 'phylum', 'class', 'order', 'family', 'genus', 'species'])

    st.write('{} : Finished to convert accession numbers to taxonomy.'.format(datetime.datetime.now().strftime('%H:%M:%S')))
    st.write('{} : Starting to create database.'.format(datetime.datetime.now().strftime('%H:%M:%S')))

    # create database
    db_name = Path(fasta_file).name.replace('.fasta.gz', '').replace('.', '_')
    db_folder = Path(output_path).joinpath(f'db_{db_name}')
    if not os.path.isdir(db_folder):
        os.mkdir(db_folder)
    db_folder = db_folder.joinpath('db')
    command = f'makeblastdb -in {fasta_file} -title db -dbtype nucl -out {db_folder}'
    os.system(command)

    # move taxonomy file
    taxonomy_file = db_folder.parent.joinpath('db_taxonomy.parquet.snappy')
    records_df.to_parquet(taxonomy_file)

    ## zip the folder
    output = Path(output_path).joinpath(f'db_{db_name}')
    shutil.make_archive(output, 'zip', output)

    st.write('{} : Finished to create database.'.format(datetime.datetime.now().strftime('%H:%M:%S')))

## MIDORI2

def midori2_taxonomy(fasta_file):
    all_accession_numbers = []
    with gzip.open(fasta_file, 'rt') as myfile:
        data = myfile.read()
        sequences = SeqIO.parse(StringIO(data), 'fasta')
        for record in sequences:
            accession = record.id
            taxonomy = []
            for i in record.description.split(';')[1:]:
                record_split = i.split('_')
                if len(record_split) == 2:
                    taxonomy.append(i.split('_')[0])
                else:
                    taxonomy.append(' '.join(i.split('_')[:2]))
            all_accession_numbers.append([accession] + taxonomy)

    accession_df = pd.DataFrame(all_accession_numbers, columns=['Accession', 'superkingdom','phylum', 'class', 'order', 'family', 'genus', 'species'])
    return accession_df

def run_midori2(output_path, fasta_file):

    st.write('{} : Starting to collect taxonomy from fasta files.'.format(datetime.datetime.now().strftime('%H:%M:%S')))

    ## collect accession numbers
    accession_df = midori2_taxonomy(fasta_file)

    st.write('{} : Starting to create database.'.format(datetime.datetime.now().strftime('%H:%M:%S')))

    # create database
    db_name = Path(fasta_file).name.replace('.fasta.gz', '')
    db_folder = Path(output_path).joinpath(f'db_{db_name}')
    if not os.path.isdir(db_folder):
        os.mkdir(db_folder)
    db_folder = db_folder.joinpath('db')
    command = f'makeblastdb -in {fasta_file} -title db -dbtype nucl -out {db_folder}'
    os.system(command)

    # write taxonomy file
    taxonomy_file_snappy = db_folder.parent.joinpath('db_taxonomy.parquet.snappy')
    accession_df.to_parquet(taxonomy_file_snappy)

    ## zip the folder
    output = Path(output_path).joinpath(f'db_{db_name}')
    shutil.make_archive(output, 'zip', output)

    st.write('{} : Finished to create database.'.format(datetime.datetime.now().strftime('%H:%M:%S')))

## PR2

def pr2_taxonomy(fasta_file):
    all_accession_numbers = []
    with gzip.open(fasta_file, 'rt') as myfile:
        data = myfile.read()
        sequences = SeqIO.parse(StringIO(data), 'fasta')
        for record in sequences:
            accession = record.id
            taxonomy_dict = {i.split(':')[0]:i.split(':')[1] for i in record.id.split(';')[1].split('=')[1].split(',')}
            taxonomy = []
            for taxon in ['k', 'p', 'c', 'o', 'f', 'g', 's']:
                try:
                    taxonomy.append(taxonomy_dict[taxon].replace('_', ' '))
                except KeyError:
                    taxonomy.append('')
            all_accession_numbers.append([accession] + taxonomy)

    accession_df = pd.DataFrame(all_accession_numbers, columns=['Accession', 'superkingdom','phylum', 'class', 'order', 'family', 'genus', 'species'])
    return accession_df

def run_pr2(output_path, fasta_file):

    st.write('{} : Starting to collect taxonomy from fasta files.'.format(datetime.datetime.now().strftime('%H:%M:%S')))

    ## collect accession numbers
    accession_df = pr2_taxonomy(fasta_file)

    st.write('{} : Starting to create database.'.format(datetime.datetime.now().strftime('%H:%M:%S')))

    # create database
    db_name = Path(fasta_file).name.replace('.fasta.gz', '')
    db_folder = Path(output_path).joinpath(f'db_{db_name}')
    if not os.path.isdir(db_folder):
        os.mkdir(db_folder)
    db_folder = db_folder.joinpath('db')
    command = f'makeblastdb -in {fasta_file} -title db -dbtype nucl -out {db_folder}'
    os.system(command)

    # write taxonomy file
    taxonomy_file_snappy = db_folder.parent.joinpath('db_taxonomy.parquet.snappy')
    accession_df.to_parquet(taxonomy_file_snappy)

    ## zip the folder
    output = Path(output_path).joinpath(f'db_{db_name}')
    shutil.make_archive(output, 'zip', output)

    st.write('{} : Finished to create database.'.format(datetime.datetime.now().strftime('%H:%M:%S')))

## TRNL

def run_trnl(output_path, fasta_file):

    st.write('{} : Starting to create database.'.format(datetime.datetime.now().strftime('%H:%M:%S')))

    # create database
    db_name = Path(fasta_file).name.replace('.fasta.gz', '')
    db_folder = Path(output_path).joinpath(f'db_{db_name}')
    os.makedirs(db_folder, exist_ok=True)
    db_folder = db_folder.joinpath('db')
    make_blast_db(fasta_file, db_folder, title="db", dbtype="nucl")

    st.write('{} : Starting to collect taxonomy from fasta files.'.format(datetime.datetime.now().strftime('%H:%M:%S')))

    accession_df = collect_accession_taxonomy(fasta_file, accession2taxonomy_parquet)

    # write taxonomy file
    taxonomy_file_snappy = db_folder.parent.joinpath('db_taxonomy.parquet.snappy')
    accession_df.to_parquet(taxonomy_file_snappy)

    ## zip the folder
    output = Path(output_path).joinpath(f'db_{db_name}')
    shutil.make_archive(output, 'zip', output)

    st.write('{} : Finished to create database.'.format(datetime.datetime.now().strftime('%H:%M:%S')))

    ####################################################################################################################

## UNITE

def unite_taxonomy(fasta_file):
    all_accession_numbers = []
    with gzip.open(fasta_file, 'rt') as myfile:
        data = myfile.read()
        sequences = SeqIO.parse(StringIO(data), 'fasta')
        for record in sequences:
            accession = record.id
            taxonomy_dict = {i.split('__')[0]:i.split('__')[1] for i in record.id.split('|')[4].split(';')}
            taxonomy = []
            for taxon in ['k', 'p', 'c', 'o', 'f', 'g', 's']:
                try:
                    taxonomy.append(taxonomy_dict[taxon].replace('_', ' '))
                except KeyError:
                    taxonomy.append('')
            all_accession_numbers.append([accession] + taxonomy)

    accession_df = pd.DataFrame(all_accession_numbers, columns=['Accession', 'superkingdom','phylum', 'class', 'order', 'family', 'genus', 'species'])
    return accession_df

def run_unite(output_path, fasta_file):

    st.write('{} : Starting to collect taxonomy from fasta files.'.format(datetime.datetime.now().strftime('%H:%M:%S')))

    ## collect accession numbers
    accession_df = unite_taxonomy(fasta_file)

    st.write('{} : Starting to create database.'.format(datetime.datetime.now().strftime('%H:%M:%S')))

    # create database
    db_name = Path(fasta_file).name.replace('.fasta.gz', '')
    db_folder = Path(output_path).joinpath(f'db_{db_name}')
    if not os.path.isdir(db_folder):
        os.mkdir(db_folder)
    db_folder = db_folder.joinpath('db')
    command = f'makeblastdb -in {fasta_file} -title db -dbtype nucl -out {db_folder}'
    os.system(command)

    # write taxonomy file
    taxonomy_file_snappy = db_folder.parent.joinpath('db_taxonomy.parquet.snappy')
    accession_df.to_parquet(taxonomy_file_snappy)

    ## zip the folder
    output = Path(output_path).joinpath(f'db_{db_name}')
    shutil.make_archive(output, 'zip', output)

    st.write('{} : Finished to create database.'.format(datetime.datetime.now().strftime('%H:%M:%S')))

## SILVA

# Initialize ete3 NCBI database
ncbi = NCBITaxa()

def create_acc2taxid_db(gz_path, db_path):
    conn = sqlite3.connect(db_path)
    c = conn.cursor()
    c.execute("DROP TABLE IF EXISTS acc2taxid")
    c.execute("CREATE TABLE acc2taxid (accession TEXT PRIMARY KEY, taxid INTEGER)")

    with gzip.open(gz_path, "rt") as f:
        next(f)  # skip header
        batch = []
        for i, line in enumerate(f):
            fields = line.strip().split("\t")
            if len(fields) >= 3:
                accession = fields[0]
                taxid = int(fields[2])
                batch.append((accession, taxid))
            if i % 100000 == 0:
                c.executemany("INSERT OR IGNORE INTO acc2taxid VALUES (?, ?)", batch)
                conn.commit()
                batch = []
        if batch:
            c.executemany("INSERT OR IGNORE INTO acc2taxid VALUES (?, ?)", batch)
            conn.commit()

    conn.close()

def get_taxid_from_sqlite(accession, db_path):
    conn = sqlite3.connect(db_path)
    c = conn.cursor()
    c.execute("SELECT taxid FROM acc2taxid WHERE accession=?", (accession,))
    result = c.fetchone()
    conn.close()
    return result[0] if result else None

def get_taxonomy_from_accession_ete3(accession):
    try:
        # Step 1: Try local lookup
        taxid = get_taxid_from_sqlite(accession, db_path)

        # Step 2: If not found, fallback to Entrez
        if not taxid:
            handle = Entrez.esummary(db="nuccore", id=accession)
            summary = Entrez.read(handle)
            handle.close()
            taxid = int(summary[0]["TaxId"])

        # Step 3: Use ete3 to get lineage
        lineage = ncbi.get_lineage(taxid)
        names = ncbi.get_taxid_translator(lineage)
        ranks = ncbi.get_rank(lineage)

        # Map desired ranks
        wanted_ranks = {
            'superkingdom': 'unclassified',
            'phylum': 'unclassified',
            'class': 'unclassified',
            'order': 'unclassified',
            'family': 'unclassified',
            'genus': 'unclassified',
            'species': 'unclassified'
        }

        for tid in lineage:
            rank = ranks.get(tid)
            name = names.get(tid)
            if rank in wanted_ranks:
                wanted_ranks[rank] = name

        wanted_ranks['Accession'] = accession
        return wanted_ranks

    except Exception as e:
        return {'Accession': accession, 'Error': str(e)}

def zip_to_gz(fasta_file):
    # Determine the output .gz file path
    gz_path = os.path.splitext(fasta_file)[0] + '.gz'

    temp_dir = 'temp_unzip'
    os.makedirs(temp_dir, exist_ok=True)

    # Unzip the contents
    with zipfile.ZipFile(fasta_file, 'r') as zip_ref:
        zip_ref.extractall(temp_dir)

    # Gzip the contents
    with open(gz_path, 'wb') as gz_file:
        for root, _, files in os.walk(temp_dir):
            for file in files:
                file_path = os.path.join(root, file)
                with open(file_path, 'rb') as f_in:
                    with gzip.open(gz_file, 'wb') as f_out:
                        shutil.copyfileobj(f_in, f_out)

    # Clean up the temporary directory
    shutil.rmtree(temp_dir)
    os.remove(fasta_file)

    return gz_path

def silva_taxonomy(fasta_file, taxonomy_xlsx):
    # Read existing taxonomy
    try:
        all_accession_numbers_df = pd.read_excel(taxonomy_xlsx)
    except FileNotFoundError:
        all_accession_numbers_df = pd.DataFrame(columns=['Accession', 'superkingdom', 'phylum', 'class', 'order', 'family', 'genus', 'species'])

    known_accessions = set(all_accession_numbers_df['Accession'].values)
    all_accession_numbers_df = all_accession_numbers_df.drop_duplicates(subset='Accession', keep='last')
    accession_dict = all_accession_numbers_df.set_index('Accession').T.to_dict('list')

    # Read sequences
    with gzip.open(fasta_file, 'rt') as handle:
        sequences = list(SeqIO.parse(handle, 'fasta'))

    n_seqs = len(sequences)

    results = []

    def process_record(record):
        id = record.id
        accession = id.split('.')[0]

        if id in known_accessions:
            return [id] + accession_dict[id]
        else:
            tax = get_taxonomy_from_accession_ete3(accession)
            if 'Error' in tax:
                assignment = record.description.split(';')[-1]
                return [id] + ['unclassified']*6 + [assignment]
            else:
                return [id] + [str(i) for i in list(tax.values())[:-1]]

    # Use threads or processes depending on I/O vs CPU time of get_taxonomy
    with ThreadPoolExecutor(max_workers=8) as executor:
        futures = {executor.submit(process_record, record): record for record in sequences}
        for i, future in enumerate(as_completed(futures), 1):
            print(f"Processed {i}/{n_seqs}")
            results.append(future.result())

    # Build output DataFrame
    accession_df = pd.DataFrame(results, columns=['Accession', 'superkingdom', 'phylum', 'class', 'order', 'family', 'genus', 'species'])

    # Combine with old and deduplicate
    out_df = pd.concat([all_accession_numbers_df, accession_df], ignore_index=True).drop_duplicates(subset=['Accession'])

    # Save
    out_df.to_excel(taxonomy_xlsx, index=False)

    return accession_df

def run_silva(output_path, fasta_file, taxonomy_xlsx):

    print('{} : Starting to collect taxonomy from fasta files.'.format(datetime.datetime.now().strftime('%H:%M:%S')))

    ## collect accession numbers
    accession_df = silva_taxonomy(fasta_file, taxonomy_xlsx)

    print('{} : Starting to create database.'.format(datetime.datetime.now().strftime('%H:%M:%S')))

    # create database
    db_name = Path(fasta_file).name.replace('.fasta.gz', '')
    db_folder = Path(output_path).joinpath(f'db_{db_name}')
    if not os.path.isdir(db_folder):
        os.mkdir(db_folder)
    db_folder = db_folder.joinpath('db')
    command = f'makeblastdb -in {fasta_file} -title db -dbtype nucl -out {db_folder}'
    os.system(command)

    # write taxonomy file
    taxonomy_file_snappy = db_folder.parent.joinpath('db_taxonomy.parquet.snappy')
    accession_df.to_parquet(taxonomy_file_snappy)

    ## zip the folder
    output = Path(output_path).joinpath(f'db_{db_name}')
    shutil.make_archive(output, 'zip', output)

    print('{} : Finished to create database.'.format(datetime.datetime.now().strftime('%H:%M:%S')))

    ####################################################################################################################

def run_silva_filtered(output_path, fasta_file, taxonomy_xlsx):

    print('{} : Starting to collect taxonomy from fasta files.'.format(datetime.datetime.now().strftime('%H:%M:%S')))

    ## collect accession numbers
    accession_df = silva_taxonomy(fasta_file, taxonomy_xlsx)

    print('{} : Starting to filter fasta files.'.format(datetime.datetime.now().strftime('%H:%M:%S')))

    accession_df_filtered = accession_df[accession_df['phylum'] != 'unclassified']
    accession_keep = set(accession_df_filtered['Accession'])
    fasta_file_filtered = Path(str(Path(fasta_file)).replace('.fasta.gz', '_sp.fasta.gz'))

    with gzip.open(fasta_file, "rt") as in_handle, gzip.open(fasta_file_filtered, "wt") as out_handle:
        for record in SeqIO.parse(in_handle, "fasta"):
            if record.id in accession_keep:
                SeqIO.write(record, out_handle, "fasta")

    print('{} : Starting to create database.'.format(datetime.datetime.now().strftime('%H:%M:%S')))

    # create database
    db_name = Path(fasta_file).name.replace('.fasta.gz', '_sp')
    db_folder = Path(output_path).joinpath(f'db_{db_name}')
    if not os.path.isdir(db_folder):
        os.mkdir(db_folder)
    db_folder = db_folder.joinpath('db')
    command = f'makeblastdb -in {fasta_file_filtered} -title db -dbtype nucl -out {db_folder}'
    os.system(command)

    # write taxonomy file
    taxonomy_file_snappy = db_folder.parent.joinpath('db_taxonomy.parquet.snappy')
    accession_df.to_parquet(taxonomy_file_snappy)

    ## zip the folder
    output = Path(output_path).joinpath(f'db_{db_name}')
    shutil.make_archive(output, 'zip', output)

    print('{} : Finished to create database.'.format(datetime.datetime.now().strftime('%H:%M:%S')))

    ####################################################################################################################

### GENERAL FUNCTIONS

def zip_to_gz(fasta_file):
    # Determine the output .gz file path
    gz_path = os.path.splitext(fasta_file)[0] + '.gz'

    temp_dir = 'temp_unzip'
    os.makedirs(temp_dir, exist_ok=True)

    # Unzip the contents
    with zipfile.ZipFile(fasta_file, 'r') as zip_ref:
        zip_ref.extractall(temp_dir)

    # Gzip the contents
    with open(gz_path, 'wb') as gz_file:
        for root, _, files in os.walk(temp_dir):
            for file in files:
                file_path = os.path.join(root, file)
                with open(file_path, 'rb') as f_in:
                    with gzip.open(gz_file, 'wb') as f_out:
                        shutil.copyfileobj(f_in, f_out)

    # Clean up the temporary directory
    shutil.rmtree(temp_dir)
    os.remove(fasta_file)

    return gz_path

def open_folder(folder_path):
    # Get the current operating system
    current_os = platform.system()

    # Open the folder based on the OS
    try:
        if current_os == "Windows":
            subprocess.Popen(f'explorer "{folder_path}"')
        elif current_os == "Darwin":  # macOS
            subprocess.Popen(['open', folder_path])
        else:  # Linux
            subprocess.Popen(['xdg-open', folder_path])
    except Exception as e:
        st.write(f"Failed to open folder: {e}")

def open_file(path: Path):
    path = str(path)
    if platform.system() == "Darwin":
        subprocess.run(["open", path])
    elif platform.system() == "Windows":
        os.startfile(path)
    else:
        subprocess.run(["xdg-open", path])

### STREAMLIT FUNCTIONS

def create_dropdown(title, description, icon, filetype, run_function, working_dir, input_folder, weblink):

    with st.expander(f"{icon} {title}", expanded=False):
        st.markdown(f"**{description}**")
        st.write("")

        if st.button(f"🔍 Learn more about {title}", key=f"learn_{title}"):
            st.info("External link will be added later.")

        # load real files WITHOUT modifying filenames
        file_names = sorted(glob.glob(str(input_folder / f'*{filetype}*')))
        file_names_dict = {Path(i).name: Path(i) for i in file_names}
        selected_file_key = st.selectbox("Select a database file:", list(file_names_dict.keys()), key=f"{title}_db_file")

        # create database
        if selected_file_key != None:
            selected_file = file_names_dict[selected_file_key]
            if st.button(f"⚙️ Create {title} database", key=f"create_{title}"):
                if selected_file.suffix == '.zip':
                    selected_file = zip_to_gz(fasta_file)
                run_function(working_dir, str(selected_file))
                st.balloons()
        st.write("___")

def create_dropdown_taxonomy_required(title, description, icon, filetype, run_function, working_dir, input_folder, weblink):

    with st.expander(f"{icon} {title}", expanded=False):
        st.markdown(f"**{description}**")
        st.write("")

        if st.button(f"🔍 Learn more about {title}", key=f"learn_{title}"):
            st.info("External link will be added later.")

        # load real files WITHOUT modifying filenames
        file_names = sorted(glob.glob(str(input_folder / f'*{filetype}*')))
        file_names_dict = {Path(i).name: Path(i) for i in file_names}
        selected_file_key = st.selectbox("Select a database file:", list(file_names_dict.keys()),
                                         key=f"{title}_db_file")

        # load real files WITHOUT modifying filenames
        taxonomy_file_names = (
                sorted(glob.glob(str(input_folder / '*.txt'))) +
                sorted(glob.glob(str(input_folder / '*.xlsx')))
        )
        taxonomy_file_names_dict = {Path(i).name: Path(i) for i in taxonomy_file_names}
        selected_taxonomy_file_key = st.selectbox("Select a taxonomy file:", list(taxonomy_file_names_dict.keys()),
                                         key=f"{title}_taxonomy_file")

        # create database
        if selected_file_key != None and selected_taxonomy_file_key != None:
            selected_file = file_names_dict[selected_file_key]
            selected_taxonomy_file = taxonomy_file_names_dict[selected_taxonomy_file_key]
            if st.button(f"⚙️ Create {title} database", key=f"create_{title}"):
                if selected_file.suffix == '.zip':
                    selected_file = zip_to_gz(fasta_file)
                run_function(working_dir, str(selected_file), selected_taxonomy_file)
                st.balloons()
        st.write("___")

def create_dropdown_silva(title, description, icon, filetype, run_function, working_dir, input_folder, weblink):

    with st.expander(f"{icon} {title}", expanded=False):
        st.markdown(f"**{description}**")
        st.write("")

        if st.button(f"🔍 Learn more about {title}", key=f"learn_{title}"):
            st.info("External link will be added later.")

        # load real files WITHOUT modifying filenames
        file_names = sorted(glob.glob(str(input_folder / f'*{filetype}*')))
        file_names_dict = {Path(i).name: Path(i) for i in file_names}
        selected_file_key = st.selectbox("Select a database file:", list(file_names_dict.keys()),
                                         key=f"{title}_db_file")

        # load real files WITHOUT modifying filenames
        taxonomy_file_names = (
                sorted(glob.glob(str(input_folder / '*.txt'))) +
                sorted(glob.glob(str(input_folder / '*.xlsx')))
        )
        taxonomy_file_names_dict = {Path(i).name: Path(i) for i in taxonomy_file_names}
        selected_taxonomy_file_key = st.selectbox("Select a taxonomy file:", list(taxonomy_file_names_dict.keys()),
                                         key=f"{title}_taxonomy_file")

        # create database
        if selected_file_key != None and selected_taxonomy_file_key != None:
            selected_file = file_names_dict[selected_file_key]
            selected_taxonomy_file = taxonomy_file_names_dict[selected_taxonomy_file_key]
            if st.button(f"⚙️ Create {title} database", key=f"create_{title}"):
                if selected_file.suffix == '.zip':
                    selected_file = zip_to_gz(fasta_file)
                run_function(working_dir, str(selected_file), selected_taxonomy_file)
                st.balloons()
        st.write("___")

### GUI ###
st.title("🗃️ APSCALE-blast Database Creator")
st.markdown("A simple tool to build reference databases for APSCALE-blast.")
st.write("___")

## SIDE BAR
if st.sidebar.button("🔄 Refresh files and folders", use_container_width=True):
    pass

### collect input path

colored_header(
    label="Select a working directory",
    description="Database will be written to the selected folder.",
    color_name="blue-70"
)

col1, col2 = st.columns(2)
working_dir = False
with col1:
    st.text_input('Please provide a PATH to your output folder:', key='output_path')
with col2:
    output_path = Path(st.session_state['output_path'])
    st.write('Chosen working directory:')
    if output_path != Path('.'):
        if output_path.exists():
            st.success('The PATH exists and you are ready to create databases!')
            working_dir = output_path / "APSCALE_databases"
            input_folder = working_dir / 'input_files'
            if not input_folder.exists():
                if st.button('Create required folders.'):
                    os.makedirs(input_folder, exist_ok=True)
        else:
            st.warning('The provided working directory does not exists!')
    else:
        st.warning('Please provide a PATH.')

colored_header(
    label="Create your own reference database",
    description="Choose a curated reference source and build your custom BLAST database.",
    color_name="blue-70"
)

if working_dir == False:
    st.warning('Please select a working directory first.')
else:
    if working_dir.exists() and input_folder.exists():
        create_dropdown("DiatBarcode",
                        "A curated diatom reference database for freshwater biomonitoring.",
                        "🟩",
                        ".xlsx",
                        run_diat_barcode,
                        working_dir,
                        input_folder,
                        "weblink",
                        )

        create_dropdown("MIDORI2",
                        "Extensive mitochondrial reference database for metazoan biodiversity.",
                        "🟦",
                        ".fasta",
                        run_midori2,
                        working_dir,
                        input_folder,
                        "weblink",
                        )

        create_dropdown("PR2",
                        "Protist Ribosomal Reference database for 18S metabarcoding.",
                        "🟪",
                        ".fasta",
                        run_pr2,
                        working_dir,
                        input_folder,
                        "weblink",
                        )

        create_dropdown_silva("SILVA",
                        "Comprehensive ribosomal RNA database (16S/18S).",
                        "🟫",
                        ".fasta",
                        None,
                        working_dir,
                        input_folder,
                        "weblink",
                        )

        create_dropdown_taxonomy_required("trnL",
                        "Plant chloroplast trnL reference database.",
                        "🌿",
                        ".fasta",
                        run_trnl,
                        working_dir,
                        input_folder,
                        "weblink",
                        )

        create_dropdown("UNITE",
                        "Fungal ITS reference database.",
                        "🍄",
                        ".fasta",
                        run_unite,
                        working_dir,
                        input_folder,
                        "weblink",
                        )

        create_dropdown("CRABS",
                        "Reference database supporting crustacean metabarcoding.",
                        "🦀",
                        ".fasta",
                        None,
                        working_dir,
                        input_folder,
                        "weblink",
                        )


