# Copyright 2026 AlQuraishi Laboratory
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from typing import Annotated

from biotite.structure import AtomArray
from pydantic import BaseModel, BeforeValidator, DirectoryPath, FilePath

from openfold3.core.config.config_utils import (
    _convert_molecule_type,
    _ensure_list,
)
from openfold3.core.data.primitives.caches.format import DatasetChainData
from openfold3.core.data.resources.residues import MoleculeType
from openfold3.projects.of3_all_atom.config.inference_query_format import Query


# MSA sample processor input configs
class MsaChainDataTrain(BaseModel):
    """Training input for a single chain in the MSA sample processor pipeline."""

    molecule_type: Annotated[MoleculeType, BeforeValidator(_convert_molecule_type)]
    alignment_representative_id: str | None


class MsaChainDataInference(BaseModel):
    """Inference input for a single chain in the MSA sample processor pipeline."""

    molecule_type: MoleculeType
    paired_msa_file_paths: (
        Annotated[list[FilePath | DirectoryPath], BeforeValidator(_ensure_list)] | None
    ) = None
    main_msa_file_paths: (
        Annotated[list[FilePath | DirectoryPath], BeforeValidator(_ensure_list)] | None
    ) = None


class MsaSampleProcessorInputTrain(BaseModel):
    """Dict-based expanded view of inference_query_format.query containing MSA data."""

    msa_chain_data: dict[str, MsaChainDataTrain]

    @classmethod
    def create_from_dataset_cache_entry(
        cls,
        dataset_cache_entry: DatasetChainData,
        atom_array: AtomArray,
        default_moltype: MoleculeType | None = None,
        default_alignment_representative_id: str | None = None,
    ):
        msa_chain_data = {}
        for chain_id, chain_data in dataset_cache_entry.chains.items():
            if hasattr(chain_data, "molecule_type"):
                molecule_type = chain_data.molecule_type
            else:
                molecule_type = default_moltype
            if hasattr(chain_data, "alignment_representative_id"):
                alignment_representative_id = chain_data.alignment_representative_id
            else:
                alignment_representative_id = default_alignment_representative_id

            msa_chain_data[chain_id] = MsaChainDataTrain(
                molecule_type=molecule_type,
                alignment_representative_id=alignment_representative_id,
            )

        # Subset to chains present in the cropped atom array
        msa_chain_data = {
            cid: cdata
            for cid, cdata in msa_chain_data.items()
            if cid in sorted(set(atom_array.chain_id.tolist()))
        }
        return cls(msa_chain_data=msa_chain_data)


class MsaSampleProcessorInputInference(BaseModel):
    """Dict-based expanded view of inference_query_format.query containing MSA data."""

    query_name: str | None = None
    msa_chain_data: dict[str, MsaChainDataInference]
    use_msas: bool
    use_paired_msas: bool
    use_main_msas: bool

    @classmethod
    def create_from_inference_query_entry(cls, inference_query: Query):
        msa_chain_data = {}
        for chain in inference_query.chains:
            for chain_id in chain.chain_ids:
                msa_chain_data[chain_id] = MsaChainDataInference(
                    molecule_type=chain.molecule_type,
                    paired_msa_file_paths=chain.paired_msa_file_paths,
                    main_msa_file_paths=chain.main_msa_file_paths,
                )
        return cls(
            query_name=inference_query.query_name,
            msa_chain_data=msa_chain_data,
            use_msas=inference_query.use_msas,
            use_paired_msas=inference_query.use_paired_msas,
            use_main_msas=inference_query.use_main_msas,
        )


# Type alias for shared MSA sample processor input
MsaSampleProcessorInput = (
    MsaSampleProcessorInputTrain | MsaSampleProcessorInputInference
)
