from os import path

import yaml
from pydantic import BaseModel


class DirsConfig(BaseModel):
    root: str = '/azaion'


class TrainingConfig(BaseModel):
    model: str = 'yolo11m.yaml'
    epochs: int = 120
    batch: int = 11
    imgsz: int = 1280
    save_period: int = 1
    workers: int = 24


class ExportConfig(BaseModel):
    onnx_imgsz: int = 1280
    onnx_batch: int = 4


class Config(BaseModel):
    dirs: DirsConfig = DirsConfig()
    training: TrainingConfig = TrainingConfig()
    export: ExportConfig = ExportConfig()

    @property
    def azaion(self) -> str:
        return self.dirs.root

    @property
    def data_dir(self) -> str:
        return path.join(self.dirs.root, 'data')

    @property
    def data_images_dir(self) -> str:
        return path.join(self.data_dir, 'images')

    @property
    def data_labels_dir(self) -> str:
        return path.join(self.data_dir, 'labels')

    @property
    def processed_dir(self) -> str:
        return path.join(self.dirs.root, 'data-processed')

    @property
    def processed_images_dir(self) -> str:
        return path.join(self.processed_dir, 'images')

    @property
    def processed_labels_dir(self) -> str:
        return path.join(self.processed_dir, 'labels')

    @property
    def corrupted_dir(self) -> str:
        return path.join(self.dirs.root, 'data-corrupted')

    @property
    def corrupted_images_dir(self) -> str:
        return path.join(self.corrupted_dir, 'images')

    @property
    def corrupted_labels_dir(self) -> str:
        return path.join(self.corrupted_dir, 'labels')

    @property
    def sample_dir(self) -> str:
        return path.join(self.dirs.root, 'data-sample')

    @property
    def datasets_dir(self) -> str:
        return path.join(self.dirs.root, 'datasets')

    @property
    def models_dir(self) -> str:
        return path.join(self.dirs.root, 'models')

    @property
    def current_pt_model(self) -> str:
        return path.join(self.models_dir, f'{prefix[:-1]}.pt')

    @property
    def current_onnx_model(self) -> str:
        return path.join(self.models_dir, f'{prefix[:-1]}.onnx')

    @classmethod
    def from_yaml(cls, config_file: str, root: str = None) -> 'Config':
        try:
            with open(config_file) as f:
                data = yaml.safe_load(f) or {}
        except FileNotFoundError:
            data = {}
        if root is not None:
            data.setdefault('dirs', {})['root'] = root
        return cls(**data)


prefix = 'azaion-'
date_format = '%Y-%m-%d'
checkpoint_file = 'checkpoint.txt'
checkpoint_date_format = '%Y-%m-%d %H:%M:%S'

CONFIG_FILE = 'config.yaml'

JPG_EXT = '.jpg'
TXT_EXT = '.txt'

OFFSET_FILE = 'offset.yaml'

SMALL_SIZE_KB = 3
CDN_CONFIG = 'cdn.yaml'
MODELS_FOLDER = 'models'

config: Config = Config.from_yaml(CONFIG_FILE)