# src/exp_kit/schedulers/base.py
from abc import ABC, abstractmethod
from typing import List, Optional, Dict, Tuple
from pathlib import Path
import subprocess

class Scheduler(ABC):
  """Abstract base class for all job schedulers."""

  def generate_script(self, name: str, **kwargs) -> str:
    """
    Generates the full submission script using a template method pattern.
    """
    header = [
      "#!/bin/bash",
      "# ======================================================================",
      "# This file was automatically generated by SbatchMan.",
      "# Manual edits are possible but should be done with care.",
      "# ======================================================================\n",
    ]

    # Get scheduler-specific lines from the subclass implementation.
    scheduler_directives = self._generate_scheduler_directives(name, **kwargs)

    # Add a command to CD into the submission directory
    working_dir_setup = [
      "\n# Change to the submission directory",
      'cd "{CWD}"',
      "\n# Update status to RUNNING",
      'if [ -f "{EXP_DIR}/metadata.json" ]; then',
      '  sed -i \'s/"status": ".*"/"status": "RUNNING"/\' {EXP_DIR}/metadata.json',
      'fi',
    ]

    # Handle common environment variables.
    env_vars = ["\n# Environment variables"]
    if envs := kwargs.get("env"):
      for env_var in envs:
        env_vars.append(f"export {env_var}")
    
    # Run the command and update the status
    footer = [
      "\n# User command",
      'echo "Running command: {CMD}"',
      'eval {CMD}',
      'EXIT_CODE=$?',
      "\n# Update status to FINISHED or FAILED",
      'if [ -f "{EXP_DIR}/metadata.json" ]; then',
      '  if [ $EXIT_CODE -eq 0 ]; then',
      '    sed -i \'s/"status": ".*"/"status": "FINISHED"/\' {EXP_DIR}/metadata.json',
      '  else',
      '    sed -i \'s/"status": ".*"/"status": "FAILED"/\' {EXP_DIR}/metadata.json',
      '  fi',
      'fi',
      'exit $EXIT_CODE',
    ]

    all_lines = header + scheduler_directives + working_dir_setup + env_vars + footer
    return "\n".join(all_lines)

  @abstractmethod
  def _generate_scheduler_directives(self, name: str, **kwargs) -> List[str]:
    """
    (Abstract) Generates the list of scheduler-specific directive lines.
    This must be implemented by subclasses.
    """
    pass

  @abstractmethod
  def submit(self, script_path: Path, exp_dir: Path) -> str:
    """
    Submits the job to the scheduler and returns the job ID.
    This method contains all logic for running the submission command.

    Args:
      script_path: The path to the executable bash script.
      exp_dir: The directory for the experiment's logs.

    Returns:
      The job ID as a string.
    """
    pass