import logging
import os
import subprocess
from dataclasses import dataclass
from textwrap import dedent

from autodft.utils.gaussian_output import move_redundant
from autodft.helpers.crest_job import CrestJob
from autodft.helpers.gaussian_inputfile import MolInfo, GaussianInputWriter


logger = logging.getLogger(__name__)


# CONSTANTS
GAUSSIAN_SCRIPT = 'run_gaussian.sh'


@dataclass(kw_only=True)
class GaussianManager:
    """A class that manages the submission of Gaussian jobs"""

    molname: str
    charge: int
    multiplicity: int
    gwriter: GaussianInputWriter

    version: str = '16'
    maxjobs: int = 10
    nprocshared: int = 12
    mem: str = '12gb'
    time: str = '24:00:00'
    write_chk: bool = False

    submitted_comfiles: list[str] = None

    def submit_from_crest_confs(self, crest_job: CrestJob) -> None:
        """Submits Gaussian jobs for conformers generated by CREST.
        Up to maxjobs conformers are submitted. The class also stores
        a list of the .com files submitted."""

        # Write a .com file for each conformer
        conf_cartesians = crest_job.conf_cartesians
        total = len(conf_cartesians)
        comfiles = [f'{self.molname}-conf{i+1}.com' for i in range(total)]
        conf_infos = [MolInfo(charge=self.charge,
                              multiplicity=self.multiplicity,
                              cartesians=x)
                      for x in conf_cartesians]
        for info, name in zip(conf_infos, comfiles):
            self.gwriter.write(info, name)

        # Determine number of jobs to submit
        n_to_submit = min([total, self.maxjobs])
        logger.info(f'Number of conformers generated by CREST: {total}')
        logger.info('Submitting Gaussian calculations to SLURM...'
               f'({n_to_submit} out of maximum {self.maxjobs} jobs)')

        # Write shell script for job
        self.write_script(n_to_submit=n_to_submit)
        self.submitted_comfiles = comfiles[:n_to_submit]

        # Submit job
        subprocess.run(['sbatch', '-W', GAUSSIAN_SCRIPT])

        # Handle extraneous files
        logger.info('The submitted Gaussian jobs have completed.')
        logfiles = [x.replace('.com', '.log') for x in self.submitted_comfiles]
        move_redundant(logfiles)
        self.delete_chk_files()

    def write_script(self, n_to_submit: str) -> None:
        """Writes a shell script to submit an array job to SLURM"""

        with open(GAUSSIAN_SCRIPT, 'w') as f:
            f.write(dedent(f"""\
                #!/bin/sh
                
                #SBATCH --nodes=1
                #SBATCH --ntasks=1
                #SBATCH --cpus-per-task={self.nprocshared}
                #SBATCH --mem={self.mem}
                #SBATCH --time={self.time}
                #SBATCH --output=/dev/null
                #SBATCH --error=/dev/null
                #SBATCH --job-name={self.molname[-8:]}
                #SBATCH --array=1-{n_to_submit}
                
                cd $SLURM_SUBMIT_DIR
                module load gaussian/{self.version}
                g{self.version} {self.molname}-conf${{SLURM_ARRAY_TASK_ID}}.com
                """))

    def delete_chk_files(self) -> None:
        """Deletes .chk files if requested"""

        if self.write_chk:
            return

        chk_files = [x.replace('.com', '.chk')
                     for x in self.submitted_comfiles]
        for f in chk_files:
            try:
                os.remove(f)
            except OSError:
                print(f'Unable to delete: {f}')
