from snakemake.logging import logger
from snakemake.utils import min_version, update_config
import re
import os
from datetime import datetime
import shutil
import sys
import secrets
import glob as _glob
min_version("8.16.0") # declare the the lowest version of snakemake that can run this script

# configfile: "configs/params.yml"


# validate the parameter settings
# validate(config, "../resources/schemas/config_schema.yaml")


# docker [image]
container: "docker://condaforge/mambaforge:latest"

# ---------------------------------------------------------------------------
# Centralised path constants — used by all rules
# ---------------------------------------------------------------------------
WORK_DIR  = config["work_dir"]
RUN_NAME  = os.path.basename(WORK_DIR)
RESULTS   = os.path.join(WORK_DIR, "Results")
LOGS      = os.path.join(WORK_DIR, "logs")
SETTINGS  = os.path.join(WORK_DIR, "edentity_pipeline_settings")

# ---------------------------------------------------------------------------
# Symlink input FASTQs into a stable directory with a normalised naming scheme
# ---------------------------------------------------------------------------
def _setup_fastq_symlinks(raw_data_dir, dest_dir):
    """
    Organise fastq file names to fit the expected naming format by symlinking
    raw files into dest_dir.  Only recreates a symlink when it is missing or
    points to a different source.
    """
    os.makedirs(dest_dir, exist_ok=True)
    gz = True
    for file in os.listdir(raw_data_dir):
        if ".fastq" not in file:
            logger.error(f"{file} is not a valid fastq file: skipping")
            continue
        base_name, extension = re.split(r'_R[12][_.]', file)
        base_name = base_name.split("fastq")[0]
        gz = file.endswith(".gz")
        ext = ".fastq.gz" if gz else ".fastq"
        if "R1" in file:
            read = "R1"
        elif "R2" in file:
            read = "R2"
        else:
            logger.error(f"Missing file(s) for sample {base_name}")
            continue
        suffix = extension.split('fastq')[0]
        new_file_name = f"{base_name}{suffix}_{read}{ext}"
        src  = os.path.abspath(os.path.join(raw_data_dir, file))
        dest = os.path.abspath(os.path.join(dest_dir, new_file_name))
        if os.path.islink(dest) and os.readlink(dest) == src:
            continue
        if os.path.islink(dest):
            os.remove(dest)
        os.symlink(src, dest)
    return gz



def _collect_sample_names(fastq_dir, ext):
    """
    Loop through each sample in the SAMPLES list.
    For each sample, determine the base name by removing the "_R1" or "_R2" suffix if present.
    Construct the file paths for the R1 and R2 fastq.gz files using the base name.
    Check if both R1 and R2 files exist in the specified raw data directory.
    If both files exist, append the base name to the SAMPLE_NAMES list.
    If either file is missing, log an error message indicating the missing file(s) for the sample.
    """
    names = []
    for r1 in _glob.glob(os.path.join(fastq_dir, f"*_R1{ext}")):
        base_name = os.path.basename(r1).replace(f"_R1{ext}", "")
        r2 = r1.replace("_R1", "_R2")
        if os.path.exists(r2):
            names.append(base_name)
        else:
            logger.error(f"Missing R2 for sample {base_name}: {r2}")
    return names


fastq_files = os.path.join(WORK_DIR, "input_fastq_files")  # stable symlink dir with normalised filenames
GZ = _setup_fastq_symlinks(config['raw_data_dir'], fastq_files)  # returns True if inputs are gzipped

ext = ".fastq.gz" if GZ else ".fastq"
SAMPLE_NAMES = _collect_sample_names(fastq_files, ext)  # list of sample base names with paired R1/R2

onstart:
    print("Starting the pipeline")
    update_config(config,
        { "runID": f"MBR_{secrets.token_hex(8)}",
        "start_time": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
        "snakemake_version": snakemake.__version__,
        "commandline_settings": sys.argv})  
# flow of execution
# ruleorder: fastpQC > merge > trimming > filter > dereplication > denoise > removeChimera > searchExact > esv_table > multiqc

rule all:           
    input:
        # MultiQC report
        os.path.join(RESULTS, "report", f"{RUN_NAME}_multiqc_reports", f"{RUN_NAME}_multiqc_report.html"),

        # custom multiqc data
        os.path.join(RESULTS, "report", f"{RUN_NAME}_custom_multiqc_data_mqc.txt"),

        # ESV table
        os.path.join(RESULTS, "report", f"{RUN_NAME}_ESV_table.tsv"),

        # summary report
        os.path.join(RESULTS, "report", f"{RUN_NAME}_summary_report.tsv")

onsuccess:
    print("pipeline completed successfully")
    

include: "rules/merge.smk"
include: "rules/trimming.smk"
include: "rules/filter.smk"
include: "rules/dereplication.smk"
include: "rules/denoise.smk"
include: "rules/chimera.smk"
include: "rules/search_exact.smk"
