from os import path

import yaml
from pydantic import BaseModel

_PROJECT_ROOT = path.dirname(path.dirname(path.abspath(__file__)))


class ApiConfig(BaseModel):
    url: str = ''
    email: str = ''
    password: str = ''


class QueueConfig(BaseModel):
    host: str = 'localhost'
    port: int = 5552
    consumer_user: str = ''
    consumer_pw: str = ''
    name: str = ''


class DirsConfig(BaseModel):
    root: str = '/azaion'
    data: str = 'data'
    data_seed: str = 'data-seed'
    data_deleted: str = 'data_deleted'


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


class ExportConfig(BaseModel):
    onnx_imgsz: int = 1280


class Config(BaseModel):
    dirs: DirsConfig = DirsConfig()
    api: ApiConfig = ApiConfig()
    queue: QueueConfig = QueueConfig()
    training: TrainingConfig = TrainingConfig()
    export: ExportConfig = ExportConfig()

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

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

    @property
    def images_dir(self) -> str:
        return path.join(self.data_dir, IMAGES_DIR)

    @property
    def labels_dir(self) -> str:
        return path.join(self.data_dir, LABELS_DIR)

    @property
    def seed_dir(self) -> str:
        return path.join(self.dirs.root, self.dirs.data_seed)

    @property
    def seed_images_dir(self) -> str:
        return path.join(self.seed_dir, IMAGES_DIR)

    @property
    def seed_labels_dir(self) -> str:
        return path.join(self.seed_dir, LABELS_DIR)

    @property
    def del_dir(self) -> str:
        return path.join(self.dirs.root, self.dirs.data_deleted)

    @property
    def del_images_dir(self) -> str:
        return path.join(self.del_dir, IMAGES_DIR)

    @property
    def del_labels_dir(self) -> str:
        return path.join(self.del_dir, LABELS_DIR)

    @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_DIR)

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

    @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'

IMAGES_DIR = 'images'
LABELS_DIR = 'labels'

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(path.join(_PROJECT_ROOT, CONFIG_FILE))
