import shutil
from datetime import datetime
from os import path

from constants import models_dir, prefix, date_format, MODELS_FOLDER
from api_client import ApiClient
from augmentation import Augmentator
from exports import export_onnx
from security import Security

# Augmentator().augment_annotations()
#train_dataset(from_scratch=True)

# resume_training('/azaion/dev/ai-training/runs/detect/train12/weights/last.pt')

result_dir = '/azaion/dev/ai-training/runs/detect/train12'
model_dir = path.join(models_dir, f'{prefix}2025-05-18')
shutil.copytree(result_dir, model_dir, dirs_exist_ok=True)

model_path = path.join(models_dir, f'{prefix[:-1]}.pt')
shutil.copy(path.join(model_dir, 'weights', 'best.pt'), model_path)

api_client = ApiClient()
onnx_path = export_onnx(model_path)
print(f'Conversion done: onnx path: {onnx_path}')

with open(onnx_path, 'rb') as binary_file:
    onnx_bytes = binary_file.read()

key = Security.get_model_encryption_key()
api_client.upload_big_small_resource(onnx_bytes, onnx_path, MODELS_FOLDER, key)