from typing import Optional
from pydantic import BaseModel, model_validator
from emissary_ml_toolkit.config.storage import StorageConfigs
from emissary_ml_toolkit.config.dataset import DatasetConfigs
from emissary_ml_toolkit.config.train import TrainConfig, BaseModelConfig
from emissary_ml_toolkit.config.test import TestConfig
from emissary_ml_toolkit.config.webhook import WebhookConfig


class TrainingJobConfig(BaseModel):
    training_job_id: Optional[str] = None
    project_id: Optional[str] = None

    # Shared configuration for every pipeline
    webhook_config: Optional[WebhookConfig] = None
    storage_config: Optional[StorageConfigs] = None

    dataset_config: Optional[DatasetConfigs] = None

    base_model_config: Optional[BaseModelConfig] = None
    train_config: Optional[TrainConfig] = None

    test_config: Optional[TestConfig] = None

    @model_validator(mode="after")
    def set_webhook_config(self):
        if self.training_job_id and self.project_id:
            self.webhook_config = WebhookConfig(
                training_job_id=self.training_job_id,
                project_id=self.project_id,
            )
        return self

    class Config:
        use_enum_values = True
