from pathlib import Path
from rdkit import Chem
import subprocess

from qspoc.converter import xyz2gjf

# Assuming your io module exists and has xyz2gjf
# from your_package.io import xyz2gjf


class G16Computer:
    def __init__(self, g16_path: str = "g16"):
        """
        Initialize the Gaussian 16 computer.
        :param g16_path: Command to run Gaussian (usually 'g16').
        """
        self.g16_path = g16_path

    def struct_opt(self, coord_file: str, calc_command: str = "B3LYP/6-31G*", charge: int = 0, multiplicity: int = 1, **settings) -> str:
        """
        Perform geometry optimization using Gaussian 16.
        """
        # 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
        gjf_path = work_dir / f"{input_path.stem}.gjf"
        log_path = work_dir / f"{input_path.stem}.log"
        output_filename = f"{input_path.stem}-opt_with_g16.xyz"
        output_path = work_dir / output_filename

        xyz2gjf(
            str(input_path),
            gjf_filepath=str(gjf_path),
            tag=input_path.stem,
            calc_command=f"opt {calc_command}",
            charge=charge,
            multiplicity=multiplicity,
            **settings,
        )

        try:
            with open(gjf_path, "r") as fin, open(log_path, "w") as fout:
                subprocess.run([self.g16_path], stdin=fin, stdout=fout, stderr=subprocess.PIPE, cwd=work_dir, check=True, text=True)

            if log_path.exists():
                self._extract_xyz_from_log(log_path, output_path)
                self._cleanup_temp_files(work_dir, input_path.stem)
                return str(output_path)
            else:
                raise RuntimeError("Gaussian run finished but no log file was found.")

        except subprocess.CalledProcessError as e:
            print(f"Gaussian 16 calculation failed! Error:\n{e.stderr}")
            raise e

    def singlepoint_energy(
        self, tag: str, coord_file: str, calc_command: str = "B3LYP/6-31G*", charge: int = 0, multiplicity: int = 1, **settings
    ) -> str:
        """
        Perform single point energy calculation using Gaussian 16.

        Args:
            coord_file: Path to the coordinate file (XYZ format)
            calc_command: Calculation method (e.g., "B3LYP/6-31G*")
            charge: Molecular charge
            multiplicity: Spin multiplicity
            **settings: Additional Gaussian settings

        Returns:
            Path to the fchk (formatted checkpoint) 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
        gjf_path = work_dir / f"{input_path.stem}.gjf"
        log_path = work_dir / f"{input_path.stem}.log"
        chk_path = work_dir / f"{input_path.stem}.chk"
        output_filename = f"{tag}-sp_with_g16.fchk"
        output_path = work_dir / output_filename

        # 2. Generate Gaussian input file (without optimization, with output to checkpoint file)
        xyz2gjf(
            str(input_path),
            gjf_filepath=str(gjf_path),
            tag=input_path.stem,
            calc_command=calc_command,  # Just the method, no "opt"
            charge=charge,
            multiplicity=multiplicity,
            chk_path=chk_path,
            **settings,
        )

        try:
            # 3. Run Gaussian
            with open(gjf_path, "r") as fin, open(log_path, "w") as fout:
                subprocess.run([self.g16_path], stdin=fin, stdout=fout, stderr=subprocess.PIPE, cwd=work_dir, check=True, text=True)
            if chk_path.exists():
                # 4. Run formchk to convert .chk to .fchk
                subprocess.run(["formchk", str(chk_path), str(output_path)], cwd=work_dir, capture_output=True, text=True, check=True)

                # 5. Cleanup temporary files
                self._cleanup_temp_files(work_dir, input_path.stem)

                return str(output_path)
            else:
                raise RuntimeError("Gaussian run finished but log file or checkpoint file was not found.")

        except subprocess.CalledProcessError as e:
            print(f"Gaussian 16 calculation failed! Error:\n{e.stderr}")
            raise e

    def _extract_xyz_from_log(self, log_path: Path, output_xyz: Path):
        """
        Helper method to extract the final optimized geometry from the Gaussian log.
        This is necessary because Gaussian does not output an .xyz file by default.
        Uses a simple parser logic.
        """
        with open(log_path, "r") as f:
            lines = f.readlines()

        # Find the last "Standard orientation" or "Input orientation"
        start_idx = -1
        for i, line in enumerate(lines):
            if "Standard orientation:" in line:
                start_idx = i

        if start_idx == -1:
            raise RuntimeError("Could not find optimized coordinates in the Gaussian log.")

        # Gaussian Standard orientation table has a 4-line header
        # ---------------------------------------------------------------------
        # Center     Atomic      Atomic             Coordinates (Angstroms)
        # Number     Number       Type             X           Y           Z
        # ---------------------------------------------------------------------
        atoms = []
        for line in lines[start_idx + 5 :]:
            if "---" in line:
                break
            parts = line.split()
            # Atom mapping: Gaussian uses atomic numbers
            # For simplicity, this example just extracts coords;
            # real implementations usually map atomic numbers to symbols.
            atom_num = parts[1]
            x, y, z = parts[3], parts[4], parts[5]
            atoms.append((atom_num, x, y, z))

        # Write to XYZ format

        num2symbol = lambda x: Chem.GetPeriodicTable().GetElementSymbol(int(x))

        with open(output_xyz, "w") as f:
            f.write(f"{len(atoms)}\n")
            f.write("Optimized by G16Computer\n")
            for a in atoms:
                f.write(f"{num2symbol(a[0])} {a[1]} {a[2]} {a[3]}\n")

    def _cleanup_temp_files(self, directory: Path, stem: str):
        """Clean up Gaussian temporary files like .chk, .rwf, etc."""
        # Common Gaussian extensions
        temp_extensions = [".chk", ".rwf", ".skr", ".inp", ".gjf"]
        for ext in temp_extensions:
            p = directory / f"{stem}{ext}"
            if p.exists():
                p.unlink()

        # Also cleanup common Fortran units if they appear
        for f in ["fort.7", "gau-temp.chk"]:
            p = directory / f
            if p.exists():
                p.unlink()
