"""Main Multiwfn descriptor generator module."""

import os
from pathlib import Path
import subprocess
from typing import List, Tuple, Optional
import pandas as pd

from .operators import MultiwfnOperators
from .descriptors.esp import ESPDescriptorCalculator
from .descriptors.cdft import CDFTDescriptorCalculator
from .descriptors.shape import ShapeDescriptorCalculator
from .descriptors.atom import AtomDescriptorCalculator
from .descriptors.bond import BondDescriptorCalculator
from .descriptors.orbit import OrbitDescriptorCalculator
from .descriptors.internal import InternalCoordinateCalculator
from .descriptors.steric import StericHindranceCalculator


class MultiwfnDescriptorGenerator:
    """Use Multiwfn to generate quantum descriptors."""

    def __init__(
        self,
        multiwfn_path: str,
        temp_file_dir: str,
        temp_coords_dir: str,
        kill: bool = False,
        file_mode: int = 1,
        acc: str = "normal",
        version: str = "stem",
    ):
        """
        Initialize the Multiwfn descriptor generator.

        Args:
            multiwfn_path: Path to Multiwfn executable
            temp_file_dir: Directory for temporary files
            kill: Continue calculation after Gaussian not converged
            file_mode: 0=don't read temp, 1=read if exist, 2=only read temp
        """
        self.multiwfn_path = multiwfn_path
        self.temp_file_dir = Path(temp_file_dir)
        self.temp_coords_dir = Path(temp_coords_dir)
        self.kill = kill
        self.file_mode = file_mode
        self.operators = MultiwfnOperators()
        self.acc = acc
        self.version = version

        self.env = None

    def load_calc_data_df(self, calc_data_df: pd.DataFrame) -> None:
        """
        Load calculation data DataFrame.

        Args:
            calc_data_df: DataFrame containing calculation data with columns:
                - name_list: Unique identifiers for molecules
                - atom_lists: List of atom indices for each molecule
                - sp_file: Path to single-point calculation files
                - prop_lists: Charge and spin states (optional)
                - ...
        """
        self.calc_data_df = calc_data_df
        self.n_molecules = len(calc_data_df)

        assert "sp_file" in self.calc_data_df.columns, "calc_data_df must contain 'sp_file' column.Did you run singlepoint_energy()?"

    def set_file_mode(self, file_mode: int) -> None:
        self.file_mode = file_mode

    def get_multiwfn_env(self, sh_filepath: str) -> None:
        sh_filepath = Path(sh_filepath)
        assert sh_filepath.is_file(), f"Shell file {sh_filepath} does not exist."
        cmd = f"source {sh_filepath} && env"

        proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, shell=True, executable="/bin/bash")
        stdout, _ = proc.communicate()

        new_env = os.environ.copy()

        for line in stdout.decode("utf-8").splitlines():
            if "=" in line:
                key, value = line.split("=", 1)
                new_env[key] = value
        self.env = new_env

    def get_ESP_descriptors(self) -> pd.DataFrame:
        """
        Get ESP descriptors from fchk files.

        Args:
            version: Descriptor version ('all' or 'stem')

        Returns:
            DataFrame with ESP descriptors
        """
        calculator = ESPDescriptorCalculator(
            multiwfn_path=self.multiwfn_path, temp_file_dir=self.temp_file_dir, file_mode=self.file_mode, acc=self.acc, version=self.version
        )
        return calculator.calculate(self.calc_data_df, self.n_molecules)

    def get_CDFT_descriptors(
        self,
        additional_prop_lists: Optional[List[List[Tuple]]] = None,
        cdft_command: str = "b3lyp/def2tzvp",
    ) -> pd.DataFrame:
        """
        Get CDFT descriptors from fchk files.

        Args:
            version: Descriptor version ('all' or 'stem')
            additional_prop_lists: Additional property lists for CDFT calculation
            calc_command: Quantum chemistry command for CDFT

        Returns:
            DataFrame with CDFT descriptors
        """
        calculator = CDFTDescriptorCalculator(
            multiwfn_path=self.multiwfn_path,
            temp_file_dir=self.temp_file_dir,
            file_mode=self.file_mode,
            kill=self.kill,
            acc=self.acc,
            version=self.version,
            env=self.env,
        )
        return calculator.calculate(self.calc_data_df, self.n_molecules, additional_prop_lists, cdft_command)

    def get_shape_descriptors(self) -> pd.DataFrame:
        """
        Get molecular shape descriptors from fchk files.

        Args:
            version: Descriptor version ('all' or 'stem')

        Returns:
            DataFrame with shape descriptors
        """
        calculator = ShapeDescriptorCalculator(
            multiwfn_path=self.multiwfn_path,
            temp_file_dir=self.temp_file_dir,
            file_mode=self.file_mode,
            acc=self.acc,
            version=self.version,
            env=self.env,
        )
        return calculator.calculate(self.calc_data_df)

    def get_atom_descriptors(self) -> pd.DataFrame:
        """
        Get atom properties descriptors from fchk files.

        Args:
            version: Descriptor version ('all' or 'stem')

        Returns:
            DataFrame with atom descriptors
        """
        calculator = AtomDescriptorCalculator(
            multiwfn_path=self.multiwfn_path,
            temp_file_dir=self.temp_file_dir,
            file_mode=self.file_mode,
            version=self.version,
            env=self.env,
        )
        return calculator.calculate(self.calc_data_df, self.n_molecules)

    def get_bond_descriptors(self) -> pd.DataFrame:
        """
        Get bond properties descriptors from fchk files.

        Args:
            version: Descriptor version ('all' or 'stem')

        Returns:
            DataFrame with bond descriptors
        """
        calculator = BondDescriptorCalculator(
            multiwfn_path=self.multiwfn_path,
            temp_file_dir=self.temp_file_dir,
            file_mode=self.file_mode,
            version=self.version,
            env=self.env,
        )
        return calculator.calculate(
            self.calc_data_df,
            self.n_molecules,
        )

    def get_orbit_descriptors(self) -> pd.DataFrame:
        """
        Get molecular orbital descriptors from fchk files.

        Args:
            version: Descriptor version ('all' or 'stem')

        Returns:
            DataFrame with orbital descriptors
        """
        calculator = OrbitDescriptorCalculator(
            multiwfn_path=self.multiwfn_path, temp_file_dir=self.temp_file_dir, file_mode=self.file_mode, version=self.version, env=self.env
        )
        return calculator.calculate(self.calc_data_df, self.n_molecules)

    def get_internal_coordinate(self) -> pd.DataFrame:
        """
        Get internal coordinates from fchk files.

        Returns:
            DataFrame with internal coordinate descriptors
        """
        calculator = InternalCoordinateCalculator()
        return calculator.calculate(self.calc_data_df, self.temp_coords_dir)

    def get_steric_parameter(self) -> pd.DataFrame:
        """
        Get steric parameters from fchk files.

        Returns:
            DataFrame with steric parameters
        """
        calculator = StericHindranceCalculator()
        return calculator.calculate(self.calc_data_df, self.temp_coords_dir)

    def _cleanup_temp_files(self) -> None:
        """Clean up temporary files generated during calculation."""
        trash_files = [
            "opt_struct.txt",
            "orbit.txt",
            "multipole.txt",
            "bndmat.txt",
            "atom_moment.txt",
            "AtomESP.txt",
            "CDFT.txt",
            "bond_order.txt",
            "charge.txt",
            "energy_index.txt",
            "oxidation.txt",
            "fort.7",
        ]

        for trash_file in trash_files:
            path = Path.cwd() / Path(trash_file)
            path.unlink() if path.exists() else None

        # Clean .chg and new.fch files in temp directory
        for f in Path.cwd().glob("*.chg"):
            f.unlink() if f.exists() else None

        new_fch_path = Path.cwd() / Path("new.fch")
        new_fch_path.unlink() if new_fch_path.exists() else None
