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:-}}"

        # Locate the actual CPLEX engine shared library. Two layouts are
        # supported: (a) a full CPLEX_Studio install, where it lives under
        # <studio>/cplex/bin/x86-64_linux/ (the historical assumption,
        # {params.cplex_bin}); (b) a standalone pip-style "cplex" package,
        # where the .so is bundled inside cplex/_internal/ next to the
        # Python bindings themselves. Search both so either layout works.
        CPLEX_SO=""
        if [ -n "{params.cplex_py}" ]; then
            CPLEX_SO=$(find "{params.cplex_py}" -name 'libcplex*.so' 2>/dev/null | head -n 1)
        fi
        if [ -z "$CPLEX_SO" ] && [ -d "{params.cplex_bin}" ]; then
            CPLEX_SO=$(find "{params.cplex_bin}" -name 'libcplex*.so' 2>/dev/null | head -n 1)
        fi
        if [ -n "$CPLEX_SO" ]; then
            export LD_LIBRARY_PATH="$(dirname "$CPLEX_SO"):$CONDA_PREFIX/lib:${{LD_LIBRARY_PATH:-}}"
        else
            export LD_LIBRARY_PATH="{params.cplex_bin}:$CONDA_PREFIX/lib:${{LD_LIBRARY_PATH:-}}"
        fi

        # 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}"
