# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/00_core.ipynb.

# %% auto 0
__all__ = ['PACKAGE_NAME', 'DEV_MODE', 'PACKAGE_DIR', 'PROJECT_DIR', 'config', 'set_env_variables', 'get_config',
           'show_project_env_vars', 'tool_is_present', 'tools_are_present', 'SamplesheetConfig', 'get_samplesheet',
           'hello_world', 'cli', 'create_bash_script', 'execute_job']

# %% ../nbs/00_core.ipynb 4
# Need the pingme for a few functions, this can be considered a static var

import importlib
import importlib.util
import os

PACKAGE_NAME: str = "pingme"  # Make sure to adjust this to your package name
DEV_MODE: bool = (
    False  # set below to override, as this is in an export block it'll be exported while the dev mode section is not
)

PACKAGE_DIR = None
try:
    spec = importlib.util.find_spec(PACKAGE_NAME)
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    PACKAGE_DIR = os.path.dirname(module.__file__)
except ImportError:
    DEV_MODE = True
except AttributeError:
    DEV_MODE = True
PROJECT_DIR = os.getcwd()  # override value in dev mode
if PROJECT_DIR.endswith("nbs"):
    DEV_MODE = True
    PROJECT_DIR = os.path.split(PROJECT_DIR)[0]

# %% ../nbs/00_core.ipynb 10
# standard libs
import os
import re

# Common to template
# add into settings.ini, requirements, package name is python-dotenv, for conda build ensure `conda config --add channels conda-forge`
import dotenv  # for loading config from .env files, https://pypi.org/project/python-dotenv/
import envyaml  # Allows to loads env vars into a yaml file, https://github.com/thesimj/envyaml
import fastcore  # To add functionality related to nbdev development, https://github.com/fastai/fastcore/
import pandas  # For sample sheet manipulation
from fastcore import (
    test,
)
from fastcore.script import (
    call_parse,
)  # for @call_parse, https://fastcore.fast.ai/script

# Project specific libraries

# %% ../nbs/00_core.ipynb 13
import importlib
import importlib.util


def set_env_variables(config_path: str, overide_env_vars: bool = True) -> bool:
    # Load dot env sets environmental values from a file, if the value already exists it will not be overwritten unless override is set to True.
    # If we have multiple .env files then we need to apply the one which we want to take precedence last with overide.

    # Order of precedence: .env file > environment variables > default values
    # When developing, making a change to the config will not be reflected until the environment is restarted

    # Set the env vars first, this is needed for the card.yaml to replace ENV variables
    # NOTE: You need to adjust PROJECT_NAME to your package name for this to work, the exception is only for dev purposes
    # This here checks if your package is installed, such as through pypi or through pip install -e  [.dev] for development. If it is then it'll go there and use the config files there as your default values.
    try:
        dotenv.load_dotenv(f"{PACKAGE_DIR}/config/config.default.env", override=False)
    except Exception as e:
        print(f"Error: {PACKAGE_DIR}/config/config.default.env does not exist")
        return False

    # 2. set values from file:
    if os.path.isfile(config_path):
        dotenv.load_dotenv(config_path, override=overide_env_vars)

    return True

# %% ../nbs/00_core.ipynb 15
import importlib
import importlib.util


def get_config(config_path: str = None, overide_env_vars: bool = True) -> dict:
    if config_path is None:
        config_path = ""
    # First sets environment with variables from config_path, then uses those variables to fill in appropriate values in the config.yaml file, the yaml file is then returned as a dict
    # If you want user env variables to take precedence over the config.yaml file then set overide_env_vars to False
    set_env_variables(config_path, overide_env_vars)

    config: dict = envyaml.EnvYAML(
        os.environ.get(
            "CORE_YAML_CONFIG_FILE", f"{PACKAGE_DIR}/config/config.default.yaml"
        ),
        strict=False,
    ).export()

    return config

# %% ../nbs/00_core.ipynb 17
# create a os.PathLike object
config = get_config(os.environ.get("CORE_CONFIG_FILE", ""))

# %% ../nbs/00_core.ipynb 19
def show_project_env_vars(config: dict) -> None:
    # Prints out all the project environment variables
    # This is useful for debugging and seeing what is being set
    for k, v in config.items():
        # If ENV var starts with PROJECTNAME_ then print
        if k.startswith(config["CORE_PROJECT_VARIABLE_PREFIX"]):
            print(f"{k}={v}")

# %% ../nbs/00_core.ipynb 21
import shutil  # using shell utilities


def tool_is_present(tool_name: str) -> bool:
    # Checks if a tool is present in current env
    return shutil.which(tool_name) is not None

# %% ../nbs/00_core.ipynb 22
import sys


def tools_are_present(tool_names: list) -> bool:
    # Checks if a list of tools are present in current env
    tools_present: bool = True
    for tool in tool_names:
        if not tool_is_present(tool):
            print(f"Tool {tool} is not present in current environment", file=sys.stderr)
            tools_present = False
    return tools_present

# %% ../nbs/00_core.ipynb 24
import pandas as pd
from pydantic import BaseModel, field_validator
from pathlib import Path
from typing import Any


class SamplesheetConfig(BaseModel):
    path: str
    delimiter: str = "\t"
    header: int = 0
    columns: list | None = None

    # Custom validator to check if the file exists
    @field_validator("path")
    def check_file_exists(cls, value):
        if not Path(value).is_file():
            raise ValueError(f"The file at path '{value}' does not exist.")
        return value

    # Override __init__ to accept a dictionary directly, for backwards compatibility probably should just use parse_obj
    def __init__(self, config: dict[str, Any]):
        super().__init__(**config)  # Unpack dictionary internally


def get_samplesheet(samplesheet_config: SamplesheetConfig) -> pd.DataFrame:
    # Load the sample sheet into a pandas dataframe
    # If columns is not None then it will only load those columns
    # If the sample sheet is a csv then it will load it as a csv, otherwise it will assume it's a tsv

    # Expected samplesheet_config:
    # sample_sheet:
    #   path: path/to/sample_sheet.tsv
    #   delimiter: '\t' # Optional, will assume , for csv and \t otherwises
    #   header: 0 # Optional, 0 indicates first row is header, None indicates no header
    #   columns: ['column1', 'column2', 'column3'] # Optional, if not provided all columns will be used

    # Example sample sheet:
    # #sample_id	file_path	metadata1	metadata2
    # Sample1	/path/to/sample1.fasta	value1	option1
    # Sample2	/path/to/sample2.fasta	value2	option2
    # Sample3	/path/to/sample3.fasta	value3	option1
    # Sample4	/path/to/sample4.fasta	value1	option2
    # Sample5	/path/to/sample5.fasta	value2	option1

    try:
        # note when we have a header the first column may begin with a #, so we need to remove it
        df = pd.read_csv(
            samplesheet_config.path,
            delimiter=samplesheet_config.delimiter,
            header=samplesheet_config.header,
            comment=None,
        )
    except Exception as e:
        print(
            "Error: Could not load sample sheet into dataframe, you have a problem with your sample sheet or the configuration."
        )
        raise e

    # Check the first header has a # in it, if so remove it for only that item
    if df.columns[0].startswith("#"):
        df.columns = [col.lstrip("#") for col in df.columns]
    # Ensure the sample sheet has the correct columns
    if samplesheet_config.columns is not None and not all(
        [col in df.columns for col in samplesheet_config.columns]
    ):
        raise ValueError("Error: Sample sheet does not have the correct columns")
    # also drop columns which are not needed
    if samplesheet_config.columns is not None:
        df = df[samplesheet_config.columns]

    # Clean the df of any extra rows that can be caused by empty lines in the sample sheet
    df = df.dropna(how="all")
    return df

# %% ../nbs/00_core.ipynb 27
def hello_world(name: str = "Not given") -> str:
    return f"Hello World! My name is {name}"

# %% ../nbs/00_core.ipynb 31
from fastcore.script import call_parse


@call_parse
def cli(
    name: str,  # Your name
    config_file: str = None,  # config file to set env vars from
):
    """
    This will print Hello World! with your name
    """
    config = get_config(config_file)  # Set env vars and get config variables
    if name is not None:
        config["example"]["input"]["name"] = name

    print(hello_world(config["example"]["input"]["name"]))

# %% ../nbs/00_core.ipynb 35
import subprocess


def create_bash_script(script_path, script_content, slurm_params=None):
    """
    Create a bash script with the given content.
    """

    # Add slurm headers to the script
    if slurm_params:
        with open(script_path, "r") as script_file:
            script_content = script_file.read()
        script_content = (
            "#!/bin/bash\n"
            + "\n".join(
                [f"#SBATCH --{key}={value}" for key, value in slurm_params.items()]
            )
            + "\n"
            + script_content
        )
        with open(script_path, "w") as script_file:
            script_file.write(script_content)
    else:
        with open(script_path, "w") as script_file:
            script_file.write(script_content)


def execute_job(script_path, use_slurm=False, slurm_params=None):
    """
    Executes a bash script either locally (bash) or via Slurm.
    """
    if use_slurm:
        # Build sbatch command
        sbatch_command = ["sbatch"]
        if slurm_params:
            for key, value in slurm_params.items():
                sbatch_command.append(f"--{key}={value}")
        sbatch_command.append(script_path)

        # Submit to Slurm
        result = subprocess.run(sbatch_command, capture_output=True, text=True)
        if result.returncode == 0:
            print(f"Job submitted to Slurm: {result.stdout.strip()}")
        else:
            print(f"Error submitting to Slurm: {result.stderr.strip()}")
    else:
        # Run locally
        result = subprocess.run(["bash", script_path], capture_output=True, text=True)
        if result.returncode == 0:
            print(f"Script executed locally: {result.stdout.strip()}")
        else:
            print(f"Error executing script locally: {result.stderr.strip()}")
