"""Base module for model prediction endpoint query."""

from __future__ import annotations

import json
import logging
from abc import ABC, abstractmethod
from pathlib import Path
from typing import Any

from folding_studio_data_models import FoldingModel
from pydantic import BaseModel


class Query(ABC):
    """Interface to define folding job queries."""

    MODEL: FoldingModel | None = None

    @classmethod
    @abstractmethod
    def from_file(cls, path: str | Path, **kwargs) -> Query:
        """Instantiates a Query object from a file."""
        ...

    @classmethod
    @abstractmethod
    def from_directory(cls, path: str | Path, **kwargs) -> Query:
        """Instantiates a Query object from a directory."""
        ...

    @classmethod
    @abstractmethod
    def from_protein_sequence(cls, protein: str, **kwargs) -> Query:
        """Instantiates a Query object from string representation of a protein."""
        ...

    @property
    @abstractmethod
    def payload(self) -> dict[str, Any]:
        """Returns the payload to be sent in the POST request."""
        ...

    @property
    @abstractmethod
    def parameters(self) -> BaseModel:
        """Parameters of the query."""
        ...

    def save_parameters(self, output_dir: Path) -> None:
        """Writes the input parameters to a JSON file inside the output directory.

        Args:
            output_dir (Path): The directory where the inference parameters JSON file will be saved.
        """
        inference_parameters_path = output_dir / "query_parameters.json"
        output_dir.mkdir(parents=True, exist_ok=True)
        with inference_parameters_path.open("w", encoding="utf-8") as f:
            json.dump(self.parameters.model_dump(mode="json"), f, indent=4)
        logging.info(f"Input parameters written to {inference_parameters_path}")