"""
有监督 Keras 模型任务的固定流水线。

这个文件放分类、分割这类“数据源 + 模型构建器 + 训练规则”任务共用的主流程。
任务入口只负责提供数据源、模型构建器和训练规则，具体训练、导出和样例测试由这里统一调度。
"""

from dataclasses import dataclass
from pathlib import Path

from deep_learning.data.spec import SupervisedDataSource
from deep_learning.env.keras import enable_mixed_precision
from deep_learning.env.logger import log
from deep_learning.env.resolve import resolve_saved
from deep_learning.pipeline.specs import Pipeline
from deep_learning.pipeline.specs.configs import CheckpointLoadRules, TrainingRule
from deep_learning.pipeline.services.logging_config import log_config
from deep_learning.models.spec import SupervisedModelBuilder
from deep_learning.pipeline.context import InferenceArtifactState, TrainingArtifactState
from deep_learning.pipeline.env.const import ENV
from deep_learning.pipeline.services import CheckpointService, build_runtime
from deep_learning.pipeline.services.training_callbacks import build_common_callbacks
from deep_learning.pipeline.stages.export import SaveKerasModelStage


@dataclass
class SupervisedModelPipeline(Pipeline):
    """
    有监督 Keras 任务主入口。

    这个类把训练、导出和样例测试固定成显式动作，分类和分割任务只需要替换数据源、
    模型构建器。
    """

    name: str
    data_source: SupervisedDataSource
    model_builder: SupervisedModelBuilder
    training_rule: TrainingRule
    checkpoint_load_rules: CheckpointLoadRules = None
    task_dir: Path = None

    def __post_init__(self):
        runtime = build_runtime(self.name, self.task_dir)
        self.task_dir = runtime.task_dir
        self._runtime_env = runtime
        if self.checkpoint_load_rules is None:
            super().__setattr__("checkpoint_load_rules", CheckpointLoadRules())

    def __setattr__(self, name, value):
        if "_runtime_env" in self.__dict__ and not name.startswith("_"):
            public_fields = {"data_source", "model_builder"}
            if name not in public_fields:
                raise AttributeError(f"cannot assign to field '{name}'")
        super().__setattr__(name, value)

    def train(self):
        with log():
            enable_mixed_precision()

        with log():
            self.log_config()

        with log():
            self._runtime_env.checkpoint_dir.mkdir(parents=True, exist_ok=True)

        with log():
            training_state = self._load_or_build_training_state(
                checkpoint_rule=self.checkpoint_load_rules.resolve_train_rule(
                    default_dirs=[self._runtime_env.checkpoint_dir]
                ),
                checkpoint_must=False
            )

        with log("构建训练计划"):
            self.model_builder.compile_training_model(training_state.training_artifact.model)
            train_ds, validation_ds = self.data_source.training_ds()
            callbacks = build_common_callbacks(
                runtime=self._runtime_env,
                checkpoint_filename="model_epoch_{epoch:03d}.weights.h5",
                save_weights_only=True
            )
            if hasattr(training_state.training_artifact.model, "summary"):
                training_state.training_artifact.model.summary()

        with log("开始训练", "训练结束"):
            return training_state.training_artifact.model.fit(
                train_ds,
                validation_data=validation_ds,
                epochs=self.training_rule.epochs,
                steps_per_epoch=self.training_rule.steps_per_epoch,
                callbacks=callbacks
            ).history

    def export_model(self) -> Path:
        self.log_config()
        self._runtime_env.checkpoint_dir.mkdir(parents=True, exist_ok=True)
        training_state = self._load_or_build_training_state(
            checkpoint_rule=self.checkpoint_load_rules.resolve_export_rule(
                default_dirs=[self._runtime_env.checkpoint_dir]
            ),
            checkpoint_must=True
        )
        inference_artifact = self.model_builder.convert_to_inference_artifact(
            training_artifact=training_state.training_artifact
        )
        inference_state = InferenceArtifactState(
            inference_artifact=inference_artifact,
            inference_resource=None
        )
        export_stage = SaveKerasModelStage(resolve_saved(f"models/{self.name}"))
        return export_stage.run(
            self._runtime_env,
            inference_state,
            training_state.checkpoint_epoch
        )

    def test_examples(self, checkpoint: str = "test") -> None:
        self.log_config()
        if checkpoint == "test":
            checkpoint_rule = self.checkpoint_load_rules.resolve_test_rule(
                default_dirs=[resolve_saved(f"models/{self.name}")]
            )
            if checkpoint_rule["suffix"] is None:
                checkpoint_rule["suffix"] = ".keras"
            checkpoint_service = CheckpointService()
            checkpoint_path, _ = checkpoint_service.resolve_existing_checkpoint(
                checkpoint_rule=checkpoint_rule,
                not_found_message="未找到任何导出模型文件"
            )
            artifact = self.model_builder.load_inference_artifact(checkpoint_path)
            checkpoint_service.log_loaded_checkpoint(checkpoint_path)
            self.data_source.test_examples(artifact.model)
            return

        training_state = self._load_or_build_training_state(
            checkpoint_rule=self.checkpoint_load_rules.resolve_train_rule(
                default_dirs=[self._runtime_env.checkpoint_dir]
            ),
            checkpoint_must=True
        )
        inference_artifact = self.model_builder.convert_to_inference_artifact(
            training_state.training_artifact
        )
        self.data_source.test_examples(
            inference_artifact.model
        )

    def log_config(self):
        self._runtime_env.log_dir.mkdir(parents=True, exist_ok=True)
        config_path = self._runtime_env.log_dir / "config.txt"
        output = log_config(self, config_path, header=f"ENV[{ENV}]")
        print(output)
        print(f"配置已保存到: {config_path}")

    def _load_or_build_training_state(
        self,
        checkpoint_rule: dict,
        checkpoint_must: bool
    ) -> TrainingArtifactState:
        checkpoint_service = CheckpointService()
        checkpoint_path, checkpoint_epoch = checkpoint_service.resolve_checkpoint(**checkpoint_rule)
        if checkpoint_path is not None:
            checkpoint_service.log_loading_checkpoint(checkpoint_path, checkpoint_epoch)
            training_artifact = self.model_builder.build_training_artifact()
            training_artifact.model.load_weights(str(checkpoint_path))
            checkpoint_service.log_loaded_checkpoint(checkpoint_path)
            return TrainingArtifactState(
                training_artifact=training_artifact,
                checkpoint_epoch=checkpoint_epoch
            )

        if checkpoint_must:
            raise ValueError(f"目录 {self._runtime_env.checkpoint_dir} 中未找到检查点文件")

        print("未找到检查点，使用新模型")
        return TrainingArtifactState(
            training_artifact=self.model_builder.build_training_artifact()
        )