import os
from glob import glob

configfile: "config.yaml"

# Global Variables from Config
GENOME_DIR = config["genome_dir"]
OUT = config["output_dir"]
EXT = config.get("assembly_extension", ".fasta")
THREADS = config.get("threads", 1)

# Define internal paths
MODELS = os.path.join(OUT, "models")
STATS = os.path.join(OUT, "model_stats")
MEMOTE = os.path.join(OUT, "memote_reports")
SUMMARY = config.get("summary_csv", os.path.join(OUT, "model_summary.csv"))

# Discover genomes
GENOMES = glob(os.path.join(GENOME_DIR, f"*{EXT}"))
SAMPLES = [os.path.splitext(os.path.basename(g))[0] for g in GENOMES]

rule all:
    input:
        SUMMARY

rule carveme_draft:
    input:
        genome = os.path.join(GENOME_DIR, "{sample}" + EXT)
    output:
        # CHANGE 1: Switched .xml to .sbml
        sbml = os.path.join(MODELS, "{sample}.sbml") 
    conda:
        "envs/py39.yaml"
    threads: THREADS
    params:
        solver = config["carveme"].get("solver", "glpk"),
        cplex_py = config["carveme"].get("cplex_lib_path", ""),
        media = config["carveme"].get("media", ""),
        cplex_bin = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(config["carveme"].get("cplex_lib_path", "")))), "bin/x86-64_linux")
    shell:
        """
        # 1. Setup environment variables
        export PYTHONPATH="{params.cplex_py}:${{PYTHONPATH:-}}"
        export LD_LIBRARY_PATH="{params.cplex_bin}:$CONDA_PREFIX/lib:${{LD_LIBRARY_PATH:-}}"
        
        # 2. Build the media argument string
        MEDIA_ARG=""
        if [ ! -z "{params.media}" ]; then
            MEDIA_ARG="-g {params.media}"
        fi

        # 3. Run CarveMe
        # CHANGE 2: Added --prodigal to fix your low gene counts
        carve --dna "{input.genome}" \
            -o "{output.sbml}" \
            $MEDIA_ARG \
            --solver {params.solver} \
            --diamond-args "-p {threads}"
        """

rule model_stats:
    input:
        # CHANGE 3: Updated to .sbml to match rule carveme_draft
        sbml = os.path.join(MODELS, "{sample}.sbml")
    output:
        stats = os.path.join(STATS, "{sample}.tsv")
    conda:
        "envs/py39.yaml"
    shell:
        "python generate_model_stats.py --model {input.sbml} --out {output.stats}"

rule memote_report:
    input:
        # CHANGE 4: Updated to .sbml to match rule carveme_draft
        sbml = os.path.join(MODELS, "{sample}.sbml")
    output:
        html = os.path.join(MEMOTE, "{sample}.html"),
        json = os.path.join(MEMOTE, "{sample}.json")
    conda:
        "envs/memote.yaml"
    shell:
        """
        memote report snapshot --filename {output.html} {input.sbml}
        memote run {input.sbml} --save-json {output.json}
        """

rule compile_summary:
    input:
        stats = expand(os.path.join(STATS, "{sample}.tsv"), sample=SAMPLES)
    output:
        SUMMARY
    conda:
        "envs/py39.yaml"
    shell:
        "python compile_model_summary.py --stats-dir {STATS} --out {output}"
