import binascii
import logging
import os
from pathlib import Path
from typing import Callable, Dict, List

import aiofiles
import yaml
from cryptography.fernet import Fernet, InvalidToken
from suzieq.poller.worker.inventory.inventory import Inventory
from suzieq.shared.exceptions import InventorySourceError

logger = logging.getLogger(__name__)


class StaticManagerInventory(Inventory):
    """StaticManager allows to import the inventory from files
    generated by the controller
    """

    def __init__(self, add_task_fn: Callable, **kwargs) -> None:
        worker_id = kwargs.pop('worker-id', '0')

        # Retrieve the key for credential file decryption
        cred_key_str = os.environ.get('SQ_CONTROLLER_POLLER_CRED', None)
        if not cred_key_str:
            raise InventorySourceError(
                'Unable to retrieve the key for decrypting the cred file'
            )
        self._cred_key = cred_key_str.encode('utf-8')

        # Get the paths of both credential and inventory files
        inv_path = os.environ.get('SQ_INVENTORY_PATH', None)
        if not inv_path:
            raise InventorySourceError(
                'Unable to get the inventory path'
            )

        self._inventory_file = Path(inv_path).joinpath(
            f'inv_{worker_id}.yml').resolve()

        self._cred_file = Path(inv_path).joinpath(
            f'cred_{worker_id}').resolve()

        if not self._inventory_file.is_file():
            raise InventorySourceError(
                f'No inventory found at {self._inventory_file}')

        if not self._cred_file.is_file():
            raise InventorySourceError(
                f'No credential file found at {self._cred_file}')

        super().__init__(add_task_fn, **kwargs)

        self._max_outstanding_cmd = int(
            os.environ.get('SQ_MAX_OUTSTANDING_CMD', 0))

    async def _get_device_list(self) -> List[Dict]:

        async with aiofiles.open(str(self._inventory_file), "r") as out_file:
            inventory_data = await out_file.read()
            inventory = yaml.safe_load(inventory_data)

        if not isinstance(inventory, dict):
            raise InventorySourceError('Invalid inventory format. Expected'
                                       f'dict, found {type(inventory)}')

        async with aiofiles.open(str(self._cred_file), "r") as out_file:
            cred_enc_data = await out_file.read()

        # Decrypt the credentials
        try:
            decryptor = Fernet(self._cred_key)
            cred_data = decryptor.decrypt(cred_enc_data.encode('utf-8'))
        except (InvalidToken, ValueError, binascii.Error):
            raise InventorySourceError(
                'The credential decryption key is not valid')
        credentials = yaml.safe_load(cred_data.decode())

        if not isinstance(credentials, dict):
            raise InventorySourceError('Invalid credential file format.')

        # Add credentials to the inventory file
        for k, v in credentials.items():
            if k in inventory:
                inventory[k].update(v)

        return inventory.values()
