import os
import shutil
from os import path, scandir, makedirs
from pathlib import Path
import random

import netron
import yaml
from ultralytics import YOLO

import constants
from api_client import ApiClient, ApiCredentials
from cdn_manager import CDNManager, CDNCredentials
from constants import datasets_dir, processed_images_dir
from security import Security
from utils import Dotdict


def export_rknn(model_path):
    model = YOLO(model_path)
    model.export(format="rknn", name="rk3588", simplify=True)
    model_stem = Path(model_path).stem
    folder_name = f'{model_stem}_rknn_model'
    shutil.move(path.join(folder_name, f'{Path(model_path).stem}-rk3588.rknn'), f'{model_stem}.rknn')
    shutil.rmtree(folder_name)
    pass


def export_onnx(model_path, batch_size=4):
    model = YOLO(model_path)
    onnx_path = Path(model_path).stem + '.onnx'
    if path.exists(onnx_path):
        os.remove(onnx_path)

    model.export(
        format="onnx",
        imgsz=1280,
        batch=batch_size,
        simplify=True,
        nms=True,
        device=0
    )
    return onnx_path


def export_tensorrt(model_path):
    YOLO(model_path).export(
        format='engine',
        batch=4,
        half=True,
        simplify=True,
        nms=True
    )


def form_data_sample(destination_path, size=500, write_txt_log=False):
    images = []
    with scandir(processed_images_dir) as imd:
        for image_file in imd:
            if not image_file.is_file():
                continue
            images.append(image_file)
    print('shuffling images')
    random.shuffle(images)
    images = images[:size]

    shutil.rmtree(destination_path, ignore_errors=True)
    makedirs(destination_path, exist_ok=True)

    lines = []
    for image in images:
        shutil.copy(image.path, path.join(destination_path, image.name))
        lines.append(f'./{image.name}')
    if write_txt_log:
        with open(path.join(destination_path, 'azaion_subset.txt'), 'w', encoding='utf-8') as f:
            f.writelines([f'{line}\n' for line in lines])


def show_model(model: str = None):
    netron.start(model)


def upload_model(model_path: str, filename: str, size_small_in_kb: int=3):
    with open(model_path, 'rb') as f_in:
        model_bytes = f_in.read()

    key = Security.get_model_encryption_key()
    model_encrypted = Security.encrypt_to(model_bytes, key)

    part1_size = min(size_small_in_kb * 1024, int(0.3 * len(model_encrypted)))
    model_part_small = model_encrypted[:part1_size]  # slice bytes for part1
    model_part_big = model_encrypted[part1_size:]

    with open(constants.CONFIG_FILE, "r") as f:
        config_dict = yaml.safe_load(f)
        d_config = Dotdict(config_dict)
        api_c = Dotdict(d_config.api)
    api = ApiClient(ApiCredentials(api_c.url, api_c.user, api_c.pw, api_c.folder))

    yaml_bytes = api.load_bytes(constants.CDN_CONFIG, '')
    data = yaml.safe_load(yaml_bytes)
    creds = CDNCredentials(data["host"],
                           data["downloader_access_key"],
                           data["downloader_access_secret"],
                           data["uploader_access_key"],
                           data["uploader_access_secret"])

    cdn_manager = CDNManager(creds)

    api.upload_file(f'{filename}.small', model_part_small, constants.MODELS_FOLDER)
    cdn_manager.upload(constants.MODELS_FOLDER, f'{filename}.big', model_part_big)