from enum import Enum
from typing import Any, Dict, Optional
from pydantic import BaseModel


class WebhookEventType(Enum):
    TRAINING_STATUS_UPDATE = 'training_jobs.status_update'
    TRAINING_TASK_TYPE_DETAIL = 'training_jobs.task_type_detail'
    TRAINING_PROGRESS_UPDATE = 'training_progress_update'
    TRAINING_PROGRESS_AND_LOSS_UPDATE= 'training_jobs.progress_and_loss_update'
    TRAINING_LOSS_UPDATE = 'training_loss_update'
    CHECKPOINT_UPLOADED = 'fine_tuned_models.uploaded'
    MODEL_TEST_RESULTS = 'fine_tuned_models.test_results'
    DATASET_UPLOADED = 'datasets.uploaded'


class TrainingProgress(BaseModel):
    current_epoch: int
    current_step: int
    total_epoch: int
    total_step: int
    estimated_remaining_time: int


class WebhookEvent(BaseModel):
    event_type: WebhookEventType
    training_id: Optional[str] = None   # Remove later platform refactoring is finish
    dataset_id: Optional[str] = None
    checkpoint: Optional[int] = None
    training_progress: Optional[TrainingProgress] = None
    training_loss: Optional[float] = None
    status: Optional[str] = None
    kill_instance: Optional[bool] = None
    test_results: Optional[dict] = None
    graphs: Optional[dict] = None


class WebhookPayload(BaseModel):
    service_id: str             # change it to project_id later
    training_job_id: str
    event_type: WebhookEventType
    event: WebhookEvent