"""The jobs module defines the core data structures for representing structure prediction jobs. This includes the PredictionJob class, which encapsulates all the necessary information about a prediction task, such as the input sequences, configuration parameters, and metadata. This module serves as the foundation for defining what a structure prediction job looks like and how it should be executed by the runners.

A PredictionJob contains all the information needed to execute a structure prediction, including:
- The input protein sequences and their corresponding chain IDs.
- The output directory where results should be saved.
- Optional ligand information (SMILES, ID, CCD code).
- Optional paths to MSA files or AF3 output directories for loading MSAs.
- A random seed for reproducibility.
The PredictionJob class also includes factory methods for creating instances from FASTA files or from an existing JSON input, providing flexibility in how jobs are defined and initialized.
"""
# varidock/jobs/jobs.py
from dataclasses import dataclass
from pathlib import Path
from typing import Optional, Sequence

from varidock.io.fasta import _read_single_fasta

@dataclass(frozen=True)
class PredictionJob:
    """Represents a structure prediction job, encapsulating all necessary information about the task."""

    name: str
    protein_sequences: Sequence[str]
    protein_chain_ids: Sequence[str]
    output_dir: Path
    input_json_path: Path | None = None
    ligand_smiles: Optional[str] = None
    ligand_id: Optional[str] = None
    ligand_ccd_code: Optional[str] = None
    msa_paths: Optional[Sequence[Path]] = None
    seed: Optional[int] = 42
    af3_output_dir: Optional[Path] = None

    @classmethod
    def from_fasta_files(
        cls,
        *,
        name: str,
        seed: int =42,
        fasta_paths: Sequence[Path],
        protein_chain_ids: Sequence[str],
        output_dir: Path,
        ligand_smiles: Optional[str] = None,
        ligand_id: Optional[str] = None,
        msa_paths: Optional[Sequence[Path]] = None,
    ) -> "PredictionJob":
        """Generate a PredictionJob instance from a set of FASTA files and corresponding chain IDs.
        
        :param cls: A reference to the PredictionJob class, used for creating an instance of the class.
        :param name: A string representing the name of the prediction job. This should be unique for each job and is used for identification purposes.
        :type name: str
        :param seed: An integer representing the random seed for reproducibility of the prediction results. This ensures that the same inputs will yield the same outputs across different runs.
        :type seed: int
        :param fasta_paths: A sequence of Path objects, each pointing to a FASTA file containing the amino acid sequence of a protein chain. Each FASTA file should contain exactly one sequence.
        :type fasta_paths: Sequence[Path]
        :param protein_chain_ids: A sequence of strings representing the chain IDs for each protein in the job. The length of this sequence must match the length of fasta_paths.
        :type protein_chain_ids: Sequence[str]
        :param output_dir: The path to the output directory where results should be saved.
        :type output_dir: Path
        :param ligand_smiles: The SMILES representation of a ligand, if applicable.
        :type ligand_smiles: Optional[str]
        :param ligand_id: A unique identifier for the ligand, which is required if a ligand is provided. This ID is used in the AF3 input JSON to reference the ligand.
        :type ligand_id: Optional[str]
        :param msa_paths: An optional sequence of Path objects pointing to pre-computed MSA files for each protein chain. If provided, the length of this sequence must match the number of protein chains. If not provided, MSAs will be generated by the pipeline.
        :type msa_paths: Optional[Sequence[Path]]
        :return: An instance of PredictionJob initialized with the provided parameters, ready to be executed by a runner.
        :rtype: PredictionJob
        """
        if len(fasta_paths) != len(protein_chain_ids):
            raise ValueError("fasta_paths and protein_chain_ids must have the same length")

        sequences = [_read_single_fasta(Path(p)) for p in fasta_paths]

        return cls(
            name=name,
            protein_sequences=sequences,
            protein_chain_ids=protein_chain_ids,
            output_dir=output_dir,
            ligand_smiles=ligand_smiles,
            ligand_id=ligand_id,
            msa_paths=msa_paths,
            seed=seed,
        )

    @classmethod
    def from_json(cls, name:str | None, json_path: Path, output_dir: Path) -> "PredictionJob":
        """Create a PredictionJob instance from an existing JSON input file. This is useful for initializing a job based on a previously prepared AF3 input JSON, allowing for flexibility in how jobs are defined and executed.
        
        :param cls: A reference to the PredictionJob class, used for creating an instance of the class.
        :param name: A string representing the name of the prediction job. If None, the name is derived from the JSON file stem.
        :type name: str | None
        :param json_path: The path to a JSON file containing input data for the prediction job.
        :type json_path: Path
        :param output_dir: The path to the output directory where results should be saved.
        :type output_dir: Path
        :return: An instance of PredictionJob initialized from the JSON file and output directory.
        :rtype: PredictionJob
        """
        return cls(
            name=name or json_path.stem,
            output_dir=output_dir,
            protein_sequences=[],
            protein_chain_ids=[],
            input_json_path=json_path,
        )
    def __post_init__(self):
        """Validate the consistency of the PredictionJob parameters after initialization. This includes checks to ensure that the number of protein sequences matches the number of chain IDs, that MSA paths (if provided) match the number of protein sequences, and that both msa_paths and af3_output_dir are not set simultaneously, as they represent mutually exclusive ways of providing MSA data.
        
        :param self: The PredictionJob instance being initialized.
        """
        if len(self.protein_chain_ids) != len(self.protein_sequences):
            raise ValueError("protein_chain_ids and protein_sequences must have the same length")
        if self.msa_paths is not None and len(self.msa_paths) != len(self.protein_sequences):
            raise ValueError("msa_paths must match protein_sequences length (or be None)")
        object.__setattr__(self, "output_dir", self.output_dir.resolve())

        if self.msa_paths is not None and self.af3_output_dir is not None:
            raise ValueError("Cannot set both msa_paths and af3_output_dir. Please provide your own MSA path, or the AF3 output dir to load MSAs from.")
