from dataclasses import dataclass
from typing import Any, List, Optional, Protocol

from envipy_additional_information import EnviPyModel, register
from pydantic import HttpUrl

from utilities.chem import FormatConverter, ProductSet


@dataclass(frozen=True, slots=True)
class Context:
    uuid: str
    url: str
    work_dir: str


class CompoundProto(Protocol):
    url: str | None
    name: str | None
    smiles: str


class RuleProto(Protocol):
    url: str
    name: str

    def apply(self, smiles, *args, **kwargs): ...


class ReactionProto(Protocol):
    url: str
    name: str
    rules: List[RuleProto]


class EnviPyDTO(Protocol):
    def get_context(self) -> Context: ...

    def get_compounds(self) -> List[CompoundProto]: ...

    def get_reactions(self) -> List[ReactionProto]: ...

    def get_rules(self) -> List[RuleProto]: ...

    @staticmethod
    def standardize(smiles, remove_stereo=False, canonicalize_tautomers=False): ...

    @staticmethod
    def apply(
        smiles: str,
        smirks: str,
        preprocess_smiles: bool = True,
        bracketize: bool = True,
        standardize: bool = True,
        kekulize: bool = True,
        remove_stereo: bool = True,
        reactant_filter_smarts: str | None = None,
        product_filter_smarts: str | None = None,
    ) -> List["ProductSet"]: ...


class PredictedProperty(EnviPyModel):
    pass


@register("buildresult")
class BuildResult(EnviPyModel):
    data: dict[str, Any] | List[dict[str, Any]] | None


@register("runresult")
class RunResult(EnviPyModel):
    producer: HttpUrl
    description: Optional[str] = None
    result: PredictedProperty | List[PredictedProperty]


@register("evaluationresult")
class EvaluationResult(EnviPyModel):
    data: dict[str, Any] | List[dict[str, Any]] | None


class BaseDTO(EnviPyDTO):
    def __init__(
        self,
        uuid: str,
        url: str,
        work_dir: str,
        compounds: List[CompoundProto],
        reactions: List[ReactionProto],
        rules: List[RuleProto],
    ):
        self.uuid = uuid
        self.url = url
        self.work_dir = work_dir
        self.compounds = compounds
        self.reactions = reactions
        self.rules = rules

    def get_context(self) -> Context:
        return Context(uuid=self.uuid, url=self.url, work_dir=self.work_dir)

    def get_compounds(self) -> List[CompoundProto]:
        return self.compounds

    def get_reactions(self) -> List[ReactionProto]:
        return self.reactions

    def get_rules(self) -> List[RuleProto]:
        return self.rules

    @staticmethod
    def standardize(smiles, remove_stereo=False, canonicalize_tautomers=False):
        return FormatConverter.standardize(
            smiles, remove_stereo=remove_stereo, canonicalize_tautomers=canonicalize_tautomers
        )

    @staticmethod
    def apply(
        smiles: str,
        smirks: str,
        preprocess_smiles: bool = True,
        bracketize: bool = True,
        standardize: bool = True,
        kekulize: bool = True,
        remove_stereo: bool = True,
        reactant_filter_smarts: str | None = None,
        product_filter_smarts: str | None = None,
    ) -> List["ProductSet"]:
        return FormatConverter.apply(
            smiles,
            smirks,
            preprocess_smiles,
            bracketize,
            standardize,
            kekulize,
            remove_stereo,
            reactant_filter_smarts,
            product_filter_smarts,
        )