"""Structure prediction module using ColabFold."""

import logging
import os
import subprocess
from pathlib import Path
from typing import List, Optional, Union

from joblib import Parallel, delayed

logger = logging.getLogger(__name__)


class ColabFoldPredictor:
    """Wrapper for ColabFold batch structure prediction."""

    def __init__(
        self,
        colabfold_command: str = "colabfold_batch",
        num_recycles: int = 3,
        num_models: int = 5,
        mutation_name: Optional[str] = None,
    ):
        """Initialize ColabFold predictor.

        Parameters
        ----------
        colabfold_command : str, optional
            Command to run ColabFold, by default "colabfold_batch".
        num_recycles : int, optional
            Number of recycles for structure prediction, by default 3.
        num_models : int, optional
            Number of models to generate, by default 5.
        mutation_name : Optional[str], optional
            Name/identifier for the mutation, by default None.

        """
        self.colabfold_command = colabfold_command
        self.num_recycles = num_recycles
        self.num_models = num_models
        self.mutation_name: Optional[str] = mutation_name

        # Check ColabFold availability at init, do fast check without --help output.
        try:
            subprocess.run(
                [self.colabfold_command, "--help"],
                stdout=subprocess.DEVNULL,
                stderr=subprocess.DEVNULL,
                timeout=10,
                check=True,
            )
        except Exception:
            raise FileNotFoundError(
                f"ColabFold command '{self.colabfold_command}' not found or not working in current environment."
            )

    def predict_structure(
        self,
        fasta_path: Path,
        output_dir: Path,
    ) -> bool:
        """Predict structure for a single FASTA file.

        Parameters
        ----------
        fasta_path : Path
            Path to input FASTA file.
        output_dir : Path
            Directory to save structure predictions.
        mutation_name : str
            Name/identifier for the mutation.

        Returns
        -------
        bool
            True if prediction succeeded, False otherwise.
        """
        fasta_path = Path(fasta_path).expanduser().resolve()
        output_dir = Path(output_dir).expanduser().resolve()

        # use mutation_name if provided, else derive from fasta_path
        mutation_name = self.mutation_name or fasta_path.stem

        output_path = output_dir / mutation_name

        cmd = [
            self.colabfold_command,
            str(fasta_path),
            str(output_path),
            "--num-recycle",
            str(self.num_recycles),
            "--num-models",
            str(self.num_models),
        ]

        logger.info(f"Running ColabFold for {self.mutation_name}...")
        logger.info(f"Command: {' '.join(cmd)}")

        try:
            subprocess.run(cmd, capture_output=True, text=True, check=True)
            logger.info(f"Successfully predicted structure for {self.mutation_name}")
            return True
        except subprocess.CalledProcessError as e:
            logger.error(f"ColabFold failed for {self.mutation_name}: {e.stderr}")
            return False
        except Exception as e:
            logger.error(f"Unexpected error during ColabFold prediction: {e}")
            return False

    def predict_structures(self, fasta_paths: List[Path], output_dir: Path) -> List[bool]:
        """Predict structures for multiple FASTA files.

        Parameters
        ----------
        fasta_paths : List[Path]
            List of paths to FASTA files.
        output_dir : Path
            Directory to save structure predictions.

        Returns
        -------
        List[bool]
            List of success flags (True/False) for each prediction.

        """
        fasta_paths = [Path(fp).expanduser().resolve() for fp in fasta_paths]
        output_dir = Path(output_dir).expanduser().resolve()
        os.makedirs(output_dir, exist_ok=True)

        # Run structure predictions in parallel using joblib
        results: List[bool] = Parallel(n_jobs=-1)(
            delayed(self.predict_structure)(
                fasta_path=fasta_path, output_dir=output_dir, mutation_name=fasta_path.stem
            )
            for fasta_path in fasta_paths
        )

        return results

    def get_pdb_files(self, output_dir: Union[Path, str]) -> List[Path]:
        """Get all PDB files generated for a mutation.

        Parameters
        ----------
        output_dir : Path
            Directory containing structure predictions.

        Returns
        -------
        List[Path]
            List of paths to PDB files.

        """
        if self.mutation_name is None:
            logger.error("mutation_name is not set; cannot locate PDB files.")
            return []
        mutation_dir = Path(output_dir).expanduser().resolve() / self.mutation_name
        if not mutation_dir.exists():
            return []

        pdb_fpaths = sorted(mutation_dir.rglob("*.pdb"))
        return pdb_fpaths
