from pathlib import Path

# This function generates an ollama batch script for running llmflux

def create_ollama_batch_script(
    account: str,
    partition: str,
    nodes: str,
    gpus_per_node: str,
    time: str,
    memory: str,
    cpus_per_task: str,
    logs_dir: Path,
    input_file: Path,
    output_file: Path,
    slurm_config,
) -> list[str]:
    # Create SLURM job script
    job_script = [
        "#!/bin/bash",
        f"#SBATCH --job-name=llm_processor",
        f"#SBATCH --account={account}",
        f"#SBATCH --partition={partition}",
        f"#SBATCH --nodes={nodes}",
        f"#SBATCH --gpus-per-node={gpus_per_node}",
        f"#SBATCH --time={time}",
        f"#SBATCH --mem={memory}",
        f"#SBATCH --cpus-per-task={cpus_per_task}",
        f"#SBATCH --output={logs_dir}/%j.out",
        f"#SBATCH --error={logs_dir}/%j.err",
    ]

    # Add extra SBATCH directives if provided
    if slurm_config.extra_sbatch_args:
        for key, value in slurm_config.extra_sbatch_args.items():
            job_script.append(f"#SBATCH --{key}={value}")

    job_script.extend([
        "",
        "# Load required modules",
        "module purge",
        "",
        "# Try loading GCC",
        "for gcc_version in '11.4.0' '11.3.0'; do",
        "    if module load gcc/$gcc_version &>/dev/null; then",
        "        echo \"Loaded gcc/$gcc_version\"",
        "        break",
        "    fi",
        "done",
        "",
        "# Try loading CUDA",
        "for cuda_version in '12.2.1' '11.7.0'; do",
        "    if module load cuda/$cuda_version &>/dev/null; then",
        "        echo \"Loaded cuda/$cuda_version\"",
        "        break",
        "    fi",
        "done",
        "",
        "# Create all necessary directories",
        "mkdir -p $DATA_INPUT_DIR $DATA_OUTPUT_DIR $MODELS_DIR $LOGS_DIR $CONTAINERS_DIR $APPTAINER_TMPDIR $APPTAINER_CACHEDIR",
        "",
        "# Start Ollama server",
        "mkdir -p $OLLAMA_HOME $OLLAMA_MODELS",
        "",
        "# Build container if needed (or if forced)",
        "if [ \"$LLMFLUX_FORCE_REBUILD\" = \"1\" ] || [ ! -f \"$CONTAINERS_DIR/llm_processor.sif\" ]; then",
        "    echo \"Building Container in ${CONTAINERS_DIR}\"",
        "    export APPTAINER_DEBUG=1",
        "    apptainer build --force $CONTAINERS_DIR/llm_processor.sif $CONTAINER_DEF",
        "    echo Container built successfully: ${?}",
        "fi",
        "",
        "# Start server with clean environment",
        "# All APPTAINERENV_* variables are automatically passed in (prefix removed)",
        "OLLAMA_DEBUG=1 apptainer exec --nv --cleanenv \\",
        "    --bind $DATA_INPUT_DIR:/app/data/input,$DATA_OUTPUT_DIR:/app/data/output,$MODELS_DIR:/app/models,$LOGS_DIR:/app/logs,$OLLAMA_HOME:$OLLAMA_HOME \\",
        "    $CONTAINERS_DIR/llm_processor.sif \\",
        "    ollama serve &",
        "",
        "OLLAMA_PID=$!",
        "",
        "# Wait for server",
        "for i in {1..60}; do",
        "    if curl -s \"http://localhost:$OLLAMA_PORT/api/version\" &>/dev/null; then",
        "        echo \"Ollama server started\"",
        "        break",
        "    fi",
        "    if ! ps -p $OLLAMA_PID > /dev/null; then",
        "        echo \"Ollama server died\"",
        "        exit 1",
        "    fi",
        "    echo \"Waiting... ($i/60)\"",
        "    sleep 1",
        "done",
        "",
        "# Pull model if needed",
        "MODEL_NAME=\"${OLLAMA_MODEL_NAME:-llama3.2:3b}\"",
        "echo \"Checking if model ${MODEL_NAME} exists...\"",
        "",
        # "# Extract base model name for Ollama (e.g. llama3.2:3b -> llama3.2)",
        # "if [[ \"$MODEL_NAME\" == *\":\"* ]]; then",
        # "    BASE_MODEL=$(echo \"$MODEL_NAME\" | cut -d':' -f1)",
        # "    echo \"Using base model name for Ollama: $BASE_MODEL\"",
        # "else",
        # "    BASE_MODEL=\"$MODEL_NAME\"",
        # "fi",
        "",
        "# Check if model exists, try to pull if it doesn't",
        "if ! curl -s \"http://localhost:$OLLAMA_PORT/api/tags\" | grep -q \"\\\"name\\\":\\\"$MODEL_NAME\\\"\"; then",
        "    echo \"Model not found, pulling base model ${MODEL_NAME}...\"",
        "    curl -X POST \"http://localhost:$OLLAMA_PORT/api/pull\" -d '{\"name\": \"'\"$MODEL_NAME\"'\"}' -H \"Content-Type: application/json\"",
        "    if [ $? -ne 0 ]; then",
        "        echo \"Failed to pull model ${MODEL_NAME}\"",
        "        echo \"Available models:\"",
        "        curl -s \"http://localhost:$OLLAMA_PORT/api/tags\" | grep -o '\"name\":\"[^\"]*\"' || echo \"None found\"",
        "        exit 1",
        "    else",
        "        echo \"Successfully pulled model ${MODEL_NAME}\"",
        "    fi",
        "else",
        "    echo \"Model ${MODEL_NAME} already exists\"",
        "fi",
        "",
        "# Run processor",
        f"python3 -c \"",
        "import sys",
        "import os",
        "sys.path.append('$PROJECT_ROOT')",
        "from llmflux.core.config import Config",
        "from llmflux.processors import BatchProcessor",
        "",
        "# Ensure OLLAMA environment variables are available in Python",
        "ollama_port = os.environ.get('OLLAMA_PORT')",
        "if ollama_port:",
        "    os.environ['OLLAMA_HOST'] = f'http://localhost:{ollama_port}'",
        "",
        "# Load model configuration",
        "config = Config()",
        "model_identifier = os.environ.get('APPTAINERENV_MODEL_IDENTIFIER', 'Llama-3.2-3B-Instruct')",
        "engine = os.environ.get('APPTAINERENV_ENGINE', 'ollama')",
        "",
        "try:",
        "    model_config = config.load_model_config(model_identifier)",
        "    model_config.engine = engine",
        "except Exception as e:",
        "    print(f'Error loading model config for {model_identifier}: {e}')",
        "    model_config = config.load_model_config('Llama-3.2-3B-Instruct')",
        "    model_config.engine = engine",
        "",
        "# Create batch processor with JSONL input",
        "batch_processor = BatchProcessor(",
        "    model_config=model_config,",
        "    batch_size=int(os.environ.get('BATCH_SIZE', '4')),",
        "    save_frequency=int(os.environ.get('SAVE_FREQUENCY', '50')),",
        "    max_retries=int(os.environ.get('MAX_RETRIES', '3')),",
        "    retry_delay=float(os.environ.get('RETRY_DELAY', '1.0'))",
        ")",
        "",
        "# Prepare run kwargs",
        "run_kwargs = {}",
        "",
        "# Add any other kwargs from environment variables",
        "for key in ['max_tokens', 'temperature', 'top_p', 'top_k']:",
        "    if key.upper() in os.environ:",
        "        run_kwargs[key] = os.environ[key.upper()]",
        "",
        f"batch_processor.run('{input_file}', '{output_file}', 'ollama', **run_kwargs)",
        "\"",
        "",
        "# Cleanup",
        "pkill -f \"ollama serve\" || true",
        "# Only remove temporary directories that we created",
        "if [ -d \"$APPTAINER_TMPDIR\" ] && [ -w \"$APPTAINER_TMPDIR\" ]; then",
        "    rm -rf \"$APPTAINER_TMPDIR\"",
        "fi",
        "if [ -d \"$APPTAINER_CACHEDIR\" ] && [ -w \"$APPTAINER_CACHEDIR\" ]; then",
        "    rm -rf \"$APPTAINER_CACHEDIR\"",
        "fi"
    ])
    return job_script
