"""
模型层契约定义。

这个包集中放模型产物、生成上下文和模型构建器协议。
Pipeline、服务和具体模型实现都从这里引用模型相关契约。
"""

from dataclasses import dataclass
from pathlib import Path
from typing import Callable, Protocol

import keras


@dataclass
class GenerationContext:
    end_of_text: int
    max_length: int
    sample_fn: Callable


@dataclass
class GenerationResult:
    token_ids: list[int]
    stop_reason: str


GenerateFn = Callable[[GenerationContext, list[int]], GenerationResult]


@dataclass
class ModelArtifact:
    model: keras.Model


@dataclass
class TextGenerationModel:
    model: keras.Model
    generate: GenerateFn


class SupervisedModelBuilder(Protocol):
    """有监督任务模型构建器需要提供训练、推理和加载能力。"""

    def build_training_artifact(self) -> ModelArtifact:
        """构建训练产物。"""
        ...

    def compile_training_model(self, model: keras.Model) -> None:
        """编译训练模型。"""
        ...

    def convert_to_inference_artifact(
        self,
        training_artifact: ModelArtifact
    ) -> ModelArtifact:
        """把训练产物转换成推理产物。"""
        return training_artifact

    def load_inference_artifact(self, model_path: Path) -> ModelArtifact:
        """从完整模型文件加载推理产物。"""
        model = keras.models.load_model(str(model_path))
        return ModelArtifact(model=model)


class TextGenerationModelBuilder(Protocol):
    """文本生成模型构建器协议"""

    def build_training_artifact(
        self,
        vocab_size: int,
        sequence_length: int
    ) -> TextGenerationModel:
        """构建文本训练产物"""
        ...

    def compile_training_model(self, model: keras.Model) -> None:
        """编译文本训练模型"""
        ...

    def convert_to_inference_artifact(
        self,
        training_artifact: TextGenerationModel
    ) -> TextGenerationModel:
        """把训练产物转换成文本推理产物"""
        return training_artifact

    def load_inference_artifact(
        self,
        model_path: Path
    ) -> TextGenerationModel:
        """从完整模型文件加载文本推理产物"""
        ...


__all__ = [
    "GenerateFn",
    "GenerationContext",
    "GenerationResult",
    "ModelArtifact",
    "SupervisedModelBuilder",
    "TextGenerationModel",
    "TextGenerationModelBuilder"
]