import os
import glob
import sys
import yaml

# Load the config file
with open("config.yaml", "r") as file:
    config = yaml.safe_load(file)
    
# -----------------------------------------------------------------------------
# INPUT PARSING (PYTHON)
# -----------------------------------------------------------------------------

# 1. DATABASE LOGIC
MAG_DIR = config["mags_dir"]
MAG_EXTS = config["mag_extensions"]
mag_files = []

# Recursive scan to find genomes even in subfolders
for ext in MAG_EXTS:
    mag_files.extend(glob.glob(os.path.join(MAG_DIR, f"**/*{ext}"), recursive=True))

if not mag_files:
    raise ValueError(f"No MAG files found in {MAG_DIR} with extensions {MAG_EXTS}")

DB_PATH = f"{config['output_dir']}/database/genomes.syldb"


# 2. SAMPLE LOGIC (PAIRED-END DETECTION)
SAMPLE_DIR = config["samples_dir"]
SAMPLE_EXTS = config["sample_extensions"]

raw_r1_files = []
for ext in SAMPLE_EXTS:
    raw_r1_files.extend(glob.glob(os.path.join(SAMPLE_DIR, f"*_1{ext}")))
    raw_r1_files.extend(glob.glob(os.path.join(SAMPLE_DIR, f"*_R1{ext}")))

if not raw_r1_files:
    raise ValueError(f"No Forward read files (_1 or _R1) found in {SAMPLE_DIR}")

SAMPLES = {}
for r1_file in raw_r1_files:
    fname = os.path.basename(r1_file)
    matched_ext = None
    for ext in SAMPLE_EXTS:
        if fname.endswith(ext):
            matched_ext = ext
            break
            
    if not matched_ext:
        continue

    # Detect sample ID and corresponding R2 file
    if f"_1{matched_ext}" in fname:
        sample_id = fname.replace(f"_1{matched_ext}", "")
        r2_file = r1_file.replace(f"_1{matched_ext}", f"_2{matched_ext}")
    elif f"_R1{matched_ext}" in fname:
        sample_id = fname.replace(f"_R1{matched_ext}", "")
        r2_file = r1_file.replace(f"_R1{matched_ext}", f"_R2{matched_ext}")
    else:
        continue

    if os.path.exists(r2_file):
        SAMPLES[sample_id] = {"r1": r1_file, "r2": r2_file}
    else:
        print(f"WARNING: Paired R2 file not found for {sample_id}. Skipping.")

SAMPLE_IDS = list(SAMPLES.keys())
print(f"DEBUG: Found {len(SAMPLE_IDS)} paired samples.")

# -----------------------------------------------------------------------------
# RULES
# -----------------------------------------------------------------------------

rule all:
    input:
        f"{config['output_dir']}/final_report.csv",
        f"{config['output_dir']}/final_report.html"

# Rule 1: Scalable Database Build (Double Extension Fix)
rule sylph_sketch_db:
    input:
        mags = mag_files
    output:
        db = DB_PATH,
        manifest = f"{config['output_dir']}/database/mags_list.txt"
    conda:
        "./envs/sylph.yaml"
    threads:
        config["sylph"]["threads"]
    run:
        import os
        
        # Use variables directly
        input_mags = input.mags
        output_manifest = output.manifest
        output_db = output.db
        t_threads = threads
        
        # Get absolute paths
        abs_mags = [os.path.abspath(m) for m in input_mags]
        
        # Write manifest
        with open(output_manifest, 'w', newline='\n') as f:
            for mag in abs_mags:
                f.write(f"{mag}\n")
        
        # Run sylph AND fix the double extension
        # Paths are quoted -- a directory with a space in its name would
        # otherwise break unquoted shell word-splitting.
        shell(
            f"""
            sylph sketch -t {t_threads} -l "{output_manifest}" -o "{output_db}"

            # Check if Sylph added a double extension and fix it
            if [ -f "{output_db}.syldb" ]; then
                mv "{output_db}.syldb" "{output_db}"
            fi
            """
        )

# Rule 2: Sketch Paired Reads
rule sylph_sketch_sample:
    input:
        r1 = lambda wildcards: SAMPLES[wildcards.sample]["r1"],
        r2 = lambda wildcards: SAMPLES[wildcards.sample]["r2"]
    output:
        sketch = f"{config['output_dir']}/sketches/{{sample}}.sylsp"
    conda:
        "./envs/sylph.yaml"
    threads:
        config["sylph"]["threads"]
    log:
        "logs/sketch_samples/{sample}.log"
    shell:
        """
        out_dir=$(dirname "{output.sketch}")
        sylph sketch -t {threads} -1 "{input.r1}" -2 "{input.r2}" -d "$out_dir" 2> "{log}"

        # Rename logic to handle Sylph's automatic naming
        generated_file=$(find "$out_dir" -name "*{wildcards.sample}*.sylsp" | head -n 1)
        if [ -f "$generated_file" ]; then
            mv "$generated_file" "{output.sketch}"
        else
            echo "Error: Could not find generated sketch file for {wildcards.sample}"
            exit 1
        fi
        """

# Rule 3: Profile Samples against Database
rule sylph_profile:
    input:
        db = DB_PATH,
        # FIX: Explicit path ensures Snakemake connects this to Rule 2
        sample_sketch = f"{config['output_dir']}/sketches/{{sample}}.sylsp"
    output:
        tsv = f"{config['output_dir']}/profiles/{{sample}}.tsv"
    conda:
        "./envs/sylph.yaml"
    threads:
        config["sylph"]["threads"]
    params:
        c = config["sylph"]["c"]
    log:
        "logs/profile/{sample}.log"
    shell:
        """
        sylph profile -t {threads} --min-number-kmers 5 -c {params.c} "{input.db}" "{input.sample_sketch}" > "{output.tsv}" 2> "{log}"
        """

# Rule 4: Merge Results
rule merge_results:
    input:
        tsvs = expand(f"{config['output_dir']}/profiles/{{sample}}.tsv", sample=SAMPLE_IDS)
    output:
        csv = f"{config['output_dir']}/final_report.csv",
        html = f"{config['output_dir']}/final_report.html"
    conda:
        "./envs/sylph.yaml"
    log:
        "logs/merge_results.log"
    shell:
        """
        python merge_results.py --inputs {input.tsvs} --output_csv "{output.csv}" --output_html "{output.html}" 2> "{log}"
        """
