from pathlib import Path
import shutil
import subprocess


class XTBComputer:
    def __init__(self, xtb_path: str):
        self.xtb_path = xtb_path

    def _generate_input(self):
        pass

    def struct_opt(self, coord_file: str, calc_command: str = "2", charge: int = 0, multiplicity: int = 1) -> str:
        # 1. 路径准备
        input_path = Path(coord_file).resolve()
        if not input_path.exists():
            raise FileNotFoundError(f"There is no file: {coord_file}")

        work_dir = input_path.parent
        output_filename = f"{input_path.stem}-opt_with_xtb.xyz"
        output_path = work_dir / output_filename
        cmd = [self.xtb_path, input_path.name, "--opt", "--chrg", str(charge), "--uhf", str(multiplicity - 1)]

        try:
            result = subprocess.run(cmd, cwd=work_dir, capture_output=True, text=True, check=True)
            assert result.returncode == 0
            xtb_default_output = work_dir / "xtbopt.xyz"

            if xtb_default_output.exists():
                shutil.move(str(xtb_default_output), str(output_path))
                self._cleanup_temp_files(work_dir)

                return str(output_path)
            else:
                raise RuntimeError("xtb ran successfully but the xtbopt.xyz file was not found")
        except subprocess.CalledProcessError as e:
            print(f"xtb calculation error! Error message: \n{e.stderr}")
            raise e

    def singlepoint_energy(self, tag: str, coord_file: str, calc_command: str = "2", charge: int = 0, multiplicity: int = 1) -> str:
        """
        Perform single point energy calculation using xtb.

        Args:
            coord_file: Path to the coordinate file (XYZ format)
            calc_command: Calculation method (e.g., "2" for GFN2-xTB)
            charge: Molecular charge
            multiplicity: Spin multiplicity

        Returns:
            Path to the molden output file
        """
        # 1. Path preparation
        input_path = Path(coord_file).resolve()
        if not input_path.exists():
            raise FileNotFoundError(f"There is no file: {coord_file}")

        work_dir = input_path.parent
        output_filename = f"{tag}-sp_with_xtb.molden"
        output_path = work_dir / output_filename

        # 2. Build command (use --sp for single point calculation with molden output)
        cmd = [self.xtb_path, input_path.name, "--sp", "--chrg", str(charge), "--uhf", str(multiplicity - 1), "--molden", "--etemp 0"]

        try:
            result = subprocess.run(cmd, cwd=work_dir, capture_output=True, text=True, check=True)
            self.sp_out = result.stdout
            # 3. Move the molden file to the desired output location
            xtb_molden_output = work_dir / "molden.input"
            if xtb_molden_output.exists():
                shutil.move(str(xtb_molden_output), str(output_path))
                self._cleanup_temp_files(work_dir)
                return str(output_path)
            else:
                raise RuntimeError("xtb ran successfully but the molden file was not found")

        except subprocess.CalledProcessError as e:
            print(f"xtb calculation error! Error message: \n{e.stderr}")
            raise e

    def get_sp_output(self) -> str:
        return self.sp_out

    def _cleanup_temp_files(self, directory: Path):
        temp_files = ["xtbout", "xtbrestart", "wbo", "charges", "xtbopt.log", ".xtboptok", "gfnff_topo"]
        for f in temp_files:
            p = directory / f
            if p.exists():
                p.unlink()
