from abc import ABC, abstractmethod
from pathlib import Path
from dataclasses import asdict, dataclass
from typing import List, Optional, Union

import yaml

from sbatchman.core.status import Status
from sbatchman.exceptions import ConfigurationError, SchedulerMismatchError
from sbatchman.config.project_config import get_project_config_dir, get_project_configs_file_path

@dataclass
class BaseConfig(ABC):
  """Abstract base class for all scheduler configs."""

  name: str
  cluster_name: str
  template_path: Optional[Path] = None
  env: Optional[List[str]] = None

  _schedulers = {}

  def __post_init__(self):
    self.template_path = self._get_config_template_path()

  def _generate_jobid_update_line(self) -> str:
    """Returns the shell command to update job_id in metadata.yaml for this scheduler."""
    scheduler = self.get_scheduler_name()
    if scheduler == "slurm":
      return 'perl -i -pe"s/job_id: [0-9]*/job_id: $SLURM_JOB_ID/" "{EXP_DIR}/metadata.yaml"'
    elif scheduler == "pbs":
      return 'perl -i -pe"s/job_id: [0-9]*/job_id: $PBS_JOBID/" "{EXP_DIR}/metadata.yaml"'
    elif scheduler == "local":
      return 'perl -i -pe"s/job_id: [0-9]*/job_id: $$/" "{EXP_DIR}/metadata.yaml"'
    else:
      return "# No job_id update for unknown scheduler"


  def _generate_script(self) -> str:
    """
    Generates the full submission script using a template method pattern.
    """
    header = [
      "#!/bin/bash",
      "# ======================================================================",
      "# This file was automatically generated by SbatchMan.",
      "# Do not edit this file manually (unless you know what you are doing).",
      "# ======================================================================\n",
    ]

    start_timestamp = ['echo "start_timestamp: \'$(date +%Y%m%d_%H%M%S.%N)\'" >> "{EXP_DIR}/metadata.yaml"\n']

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

    # Add a command to CD into the submission directory
    working_dir_setup = [
      "\n# Change to the submission directory",
      'cd "{CWD}"',
      "\n# Update job_id in metadata.yaml",
      self._generate_jobid_update_line(),
      "\n# Update status to RUNNING",
      'if [ -f "{EXP_DIR}/metadata.yaml" ]; then',
      ' perl -i -pe\'s/status: .*/status: RUNNING/\' {EXP_DIR}/metadata.yaml',
      'fi',
    ]

    # Set environment variables
    env_vars = []
    if envs := self.env:
      env_vars = ["\n# Environment variables"]
      if (type(envs) is not list):
        raise ConfigurationError("The 'env' parameter must be a list.")
      for env_var in envs:
        env_vars.append(f"export {env_var}")
    
    # Run the command and update the status
    footer = [
      "\n# (Optional) Preprocess command",
      '{PREPROCESS}',
      "\n# User command",
      '{CMD}',
      'EXIT_CODE=$?',
      "\n# Update status to COMPLETED or FAILED",
      'if [ -f "{EXP_DIR}/metadata.yaml" ]; then',
      '  if [ $EXIT_CODE -eq 0 ]; then',
      '    perl -i -pe\'s/status: .*/status: COMPLETED/\' {EXP_DIR}/metadata.yaml',
      '  else',
      '    perl -i -pe\'s/status: .*/status: FAILED/\' {EXP_DIR}/metadata.yaml',
      '  fi',
      '  perl -i -pe"s/exitcode: .*/exitcode: $EXIT_CODE/" {EXP_DIR}/metadata.yaml',
      'fi',
      "\n# (Optional) Postprocess command",
      '{POSTPROCESS}',
      '\necho "end_timestamp: \'$(date +%Y%m%d_%H%M%S.%N)\'" >> "{EXP_DIR}/metadata.yaml"\n'
      '\nexit $EXIT_CODE',
    ]

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

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

  @staticmethod
  @abstractmethod
  def get_scheduler_name() -> str:
    """
    Returns the name of the scheduler this parameters class is associated with.
    This must be implemented by subclasses.
    """
    pass

  def _update_main_config(self, overwrite: bool = False):
    """Reads, updates, and writes to the central configurations.yaml file."""
    config_path = get_project_configs_file_path()

    try:
      with open(config_path, 'r') as f:
        data = yaml.safe_load(f) or {}
    except FileNotFoundError:
      data = {}

    scheduler_name = self.get_scheduler_name()
    if self.cluster_name in data:
      if data[self.cluster_name].get('scheduler') != scheduler_name:
        raise SchedulerMismatchError(f"Cluster '{self.cluster_name}' is already configured with scheduler '{data[self.cluster_name]['scheduler']}'. Cannot add a '{scheduler_name}' configuration.")
    
    if not overwrite and self.cluster_name in data and self.name in data[self.cluster_name].get('configs', {}):
      raise ConfigurationError(f"Configuration '{self.name}' for cluster '{self.cluster_name}' already exists. Use '--overwrite' to update it.")
    
    data.setdefault(self.cluster_name, {})['scheduler'] = scheduler_name
    data[self.cluster_name].setdefault('configs', {})

    # Use asdict for clean conversion and filter keys
    param_dict = asdict(self)
    # Convert Path objects to strings
    for k, v in param_dict.items():
      if isinstance(v, Path):
        param_dict[k] = str(v)
    clean_config_params = {
      k: v for k, v in param_dict.items() 
      if v is not None and k not in ['name', 'cluster_name', 'scheduler']
    }
    clean_config_params['scheduler'] = scheduler_name
    data[self.cluster_name]['configs'][self.name] = clean_config_params

    with open(config_path, 'w') as f:
      yaml.dump(data, f, default_flow_style=False, sort_keys=False)

  def _get_config_template_path(self) -> Path:
    config_dir = get_project_config_dir() / self.cluster_name
    config_dir.mkdir(parents=True, exist_ok=True)
    return config_dir / f"{self.name}.sh"

  def _write_script(self) -> Path:
    """Saves the configuration script to a file inside a scheduler-specific folder."""
    script_content = self._generate_script()

    if self.template_path is None:
      raise ConfigurationError("Template path is not set. This should not happen.")

    with open(self.template_path, "w") as f:
      f.write(script_content)
    return self.template_path
  
  def save_config(self, overwrite: bool = False) -> Path:
    """
    Creates a configuration script and updates the configuration file.

    Returns:
      The path to the created configuration script file.
    """
    self._update_main_config(overwrite)
    return self._write_script()
  
  @staticmethod
  @abstractmethod
  def get_job_status(job_id: Union[str, int]) -> Status:
    """
    Returns the status of a job for this scheduler.
    This must be implemented by subclasses.
    """
    pass