# coding: utf-8
import os
import random
import time
import json

from collections import OrderedDict
from subprocess import check_output, STDOUT, Popen, PIPE, run

from . import data_store
from . import instance_sessions_data
from . import license_docker
from .docker_and_languages import DOCKER_URL, DOCKER_TRANSCRIPTION_REPO, get_available_docker_languages
from .eeg_python_utils import utils

SPEECHMATICS_IP = os.environ['SPEECHMATICS_IP']
BASE_SPEECHMATICS_URL = os.environ['BASE_SPEECHMATICS_URL']
SPEECHMATICS_DOCKER_LOGS_DIR = os.environ['SPEECHMATICS_DOCKER_LOGS_DIR']

# SPEECHMATICS LANGUAGES
# Get the available SPEECHMATICS VERSIONS from the images downloaded
SPEECHMATICS_VERSIONS = {}
available_languages = get_available_docker_languages()
if len(available_languages) > 0:
    for language in available_languages:
        base_model_name = language.get('base_model_name')
        lang_version = language.get('version')

        if not base_model_name or not lang_version:
            print(f'Error with the language {language}')
            pass

        # Check if the image actually exists
        docker_image = f'{DOCKER_URL}/{DOCKER_TRANSCRIPTION_REPO}{base_model_name}:{lang_version}'
        args = ['docker', 'inspect', docker_image]
        result = Popen(args, stdout=PIPE).communicate()[0].decode("utf-8").strip()

        # If the image exists, the reult is a JSON-like object with the image information
        # Otherwise, it's []
        if len(result) == 2 or result == '[]':
            print(f'Error: the image {docker_image} does not exist')
            pass

        SPEECHMATICS_VERSIONS[base_model_name] = docker_image
# GPU RECIPE LANGUAGES
RECIPES_LANGUAGES = {
    'recipe1': ['eu', 'ba', 'uk', 'cy', 'mn', 'eo', 'ta', 'be', 'ug', 'tr', 'mr', 'en'],
    'recipe2': ['id', 'sk', 'hr', 'lt', 'gl', 'bg', 'sl', 'ia', 'lv', 'et', 'fa', 'es', 'ro'],
    'recipe3': ['ko', 'de', 'it', 'ms', 'sv', 'da', 'hu', 'hi', 'cs', 'ca', 'el', 'fi'],
    'recipe4': ['yue', 'ar', 'th', 'fr', 'no', 'pl', 'vi', 'pt', 'ja', 'ru', 'cmn', 'nl']
}

USERNAME = 'lexilocal'
EMAIL = "support@eegent.com"
COMPANY = 'EEG'

print(os.environ)
DOCKER_NETWORK_BRIDGE = os.environ['NETWORK_BRIDGE_NAME']
USAGE_CONTAINER_EATS_PORT = os.environ['EATS_PORT']
USAGE_CONTAINER_DATA_PORT = os.environ['DATA_PORT']
USAGE_CONTAINER_NAME = os.environ['USAGE_CONTAINER_NAME']
GPU_CONTAINER_NAME = os.environ['GPU_CONTAINER_NAME']

BASE_MODELS = {
    'en': 'Global English',
    'en-US': 'American English',
    'en-GB': 'British English',
    'en-AU': 'Australian English',
    'de': 'German',
    'es': 'Spanish',
    'fr': 'French',
    'ja': 'Japanese',
    'hi': 'Hindi',
    'nl': 'Dutch',
    'ru': 'Russian',
    'it': 'Italian',
    'pt': 'Portuguese',
}

ordered_base_model_names = ['en', 'en-US', 'en-GB', 'en-AU', 'de', 'es', 'fr', 'ja', 'hi', 'nl', 'ru', 'it', 'pt']
ORDERED_BASE_MODELS = OrderedDict()

for base_model in ordered_base_model_names:
    ORDERED_BASE_MODELS[base_model] = {
        'display_name': BASE_MODELS[base_model],
        'speechmatics_name': base_model,
    }


def list_base_models():
    return OrderedDict(ORDERED_BASE_MODELS)


# Management
# Returns a list of strings containing the ids of all docker containers.
def get_docker_containers():
    # Get list of currently running docker containers - select speechmatics only - select their id as column #1 from the table.
    try:
        # search for any running docker container with speechmatics in the image name.
        # do not try to search for the containers by name since shell=True could lead to shell injection as container names are user provided.
        running_dockers = check_output('docker ps | grep speechmatics.*rt-asr-transcriber | awk \'{print($1)}\'', stderr=STDOUT, shell=True)
        running_dockers_list_of_string = running_dockers.decode("utf-8").splitlines()
        return running_dockers_list_of_string
    except Exception:
        return []


# Allow spaces on the front-end but convert them into underscores (_) when working with the containers
def parse_container_name(container_name):
    return container_name.replace(' ', '_')


# Check if a given container is running
def get_running_container_data(container_name):
    container_name = parse_container_name(container_name)
    p = Popen(['docker', 'container', 'inspect', container_name], start_new_session=True, stdout=PIPE)
    out_byte_str, err_byte_str = p.communicate()
    out_str = out_byte_str.decode('utf-8')
    raw_json_data = json.loads(out_str)

    if (len(raw_json_data) == 0):
        return {}

    return raw_json_data[0]


def get_sesion_id_from_log(container_name):
    container_name = parse_container_name(container_name)

    try:
        log_path = f'{SPEECHMATICS_DOCKER_LOGS_DIR}/{container_name}.txt'

        if not utils.check_if_file(log_path):
            return None

        with open(log_path) as file:
            lines = [line.rstrip() for line in file]

        session_id_lines = set(line.split('Session-ID=')[1].split()[0] for line in lines if 'Session-ID=' in line)

        if len(session_id_lines) != 1:
            return None

        return session_id_lines.pop()
    except Exception as e:
        print(f'Error!!! {e}')
        return None


def update_speechmatics_session_id(instance_id, session_id, speechmatics_session_id):
    instance_sessions = instance_sessions_data.get_instance_sessions(instance_id)

    session = [session for session in instance_sessions if session.get('session_id') == session_id]

    if len(session) != 1:
        return None

    session = session[0]

    session['speechmatics_session_id'] = speechmatics_session_id

    data_store.set_instance_sessions(instance_id, instance_sessions)\

    return speechmatics_session_id


def save_log(container_name):
    container_name = parse_container_name(container_name)
    with open(SPEECHMATICS_DOCKER_LOGS_DIR + '/' + container_name + '.txt', 'w') as log_file:
        run(['docker', 'logs', container_name], stdout=log_file, stderr=STDOUT)
        log_file.flush()
        log_file.close()
        print('Saved ASR log')


# Stop the container using its ID to avoid stopping a new container sapwned after
# the stopping action was completed
def stop_docker_container(container_name, container_id=None):
    container_name = parse_container_name(container_name)
    print(f'Before kill, saving log for: {container_name}')
    save_log(container_name)
    print(f'Killing container: {container_name}')

    if container_id is None:
        container_id = get_container_id(container_name)
    if container_id != container_name:
        print(f'Container ID: {container_id}')

    try:
        # Killing a container doesn't fully clean it up.
        # But, we used --rm as a flag when creating the container,
        # meaning it will automatically clean itself up when killed.
        res = check_output(['docker', 'stop', container_id])
        print('Docker stopped')
        print(res)
    except Exception as e:
        print(f'Error stopping container {container_id}: {e}')


# Returns the container's ID (if there is an error, returns the container's name)
def get_container_id(container_name):
    container_name = parse_container_name(container_name)
    try:
        container_id = check_output(f'docker ps -aqf "name=^{container_name}$"', shell=True, stderr=STDOUT).decode('utf-8')
        # Container's ID is only alphanumeric. Remove any whitespace
        container_id = ''.join(letter for letter in container_id if letter.isalnum())

        return container_id if len(container_id) > 0 else container_name
    except Exception as e:
        print(f'Error getting the ID of the contaier {container_name}: {e}')
        return container_name


# Check if the instance could use the GPU container
def instance_use_gpu(instance_lang):
    gpu_container_data = get_running_container_data(GPU_CONTAINER_NAME)

    if len(gpu_container_data) == 0:
        print('No GPU')
        return False

    gpu_container_status = gpu_container_data.get('State', {}).get('Status')
    gpu_container_repo = gpu_container_data.get('Config', {}).get('Image')

    if None in [gpu_container_status, gpu_container_repo]:
        print('No GPU container')
        return False
    if gpu_container_status != 'running':
        print('GPU container is not running')
        return False

    gpu_container_version = gpu_container_repo.split('/')[1].split(':')[0]

    # Single-Language
    if 'recipe' not in gpu_container_version:
        gpu_container_lang = gpu_container_version[-2:]

        return instance_lang == gpu_container_lang

    # Recipes
    recipe_with_lang = [recipe for recipe in RECIPES_LANGUAGES if instance_lang in RECIPES_LANGUAGES[recipe]]

    if len(recipe_with_lang) != 1:
        print('No recipe found')
        return False

    gpu_container_recipe = gpu_container_version[-7:]

    return gpu_container_recipe == recipe_with_lang[0]


# Spawn a new docker container.
# Returns a tuple of
# 1. The subprocess from Popen or None
# 2. The port used for 9000
def spawn_docker_container(container_name, license_status=None, language='en'):
    container_name = parse_container_name(container_name)

    # Get an avaiable port to be used to spaww the container
    # Will try up to RETRIES times, to avoid an infinite loop
    def get_available_port(port_range, retries=5):
        # Check if the given port is available
        # If the port is already in use, the command will return a non-empty value
        # If the port is not in use, the command will either return an empty value or return a non-zero exit status 1
        def check_port_avaialble(port):
            try:
                res = check_output(f'sudo lsof -i:{port}', shell=True, stderr=STDOUT)

                return len(res) == 0
            except Exception:
                return True

        port = None
        for count in range(0, retries):
            port = str(random.randint(port_range[0], port_range[1]))
            print(f'Trying with port: {port}')

            if check_port_avaialble(port):
                break

            port = None

        return port

    if license_status is None:
        license_status = license_docker.get_license()
        if license_status is None:
            return None
    license_token = license_status.get('signedclaimstoken', '')
    if license_token == '':
        return None

    # Get available ports
    p9000 = get_available_port([9000, 9999])
    if p9000 is None:
        print('Not able to find available port p9000')
        raise Exception
    # We are reserving the port 8001 for the GPU container
    p8001 = get_available_port([8002, 8999])
    if p8001 is None:
        print('Not able to find available port p8001')
        raise Exception

    # Raise an exception if there is no image to the given language
    if language not in SPEECHMATICS_VERSIONS:
        raise Exception(f'Wrong language. There is no image for {language}')

    SPEECHMATICS_DOCKER_VERSION = SPEECHMATICS_VERSIONS[language]

    # --rm will fully remove the container when killed
    # --network=DOCKER_NETWORK_BRIDGE will set the container to use a bridged network to communicate with the usage container
    # SM_EATS_URL is the docker container address and port that the ASR container will send event data to
    basic_args = ['docker', 'run', '--rm', '--network=' + DOCKER_NETWORK_BRIDGE, '--name', container_name, '-p', p9000 + ':9000', '-p', p8001 + ':8001', '-e', 'SM_DO_SESSION_TRANSFER=true']
    usage_args = ['-e', 'SM_ENABLE_USAGE_REPORTING=true', '-e', 'SM_EATS_URL=' + USAGE_CONTAINER_NAME + ':' + USAGE_CONTAINER_EATS_PORT]
    gpu_args = ['-e', 'SM_INFERENCE_ENDPOINT=' + GPU_CONTAINER_NAME + ':8001']
    license_and_language_args = ['-e', 'LICENSE_TOKEN=' + license_token, SPEECHMATICS_DOCKER_VERSION]

    args = []

    args += basic_args
    args += usage_args

    # Check if the Instance could use the GPU container
    if instance_use_gpu(instance_lang=language):
        args += gpu_args

    args += license_and_language_args

    print(' '.join(args))
    return Popen(args, start_new_session=True), p9000


# Returns a tuple of
# 1. Whether or not a container was successfully spawned.
# 2. The port used for 9000
# Retries is 3 by default since we are using random to guess an open port.
def try_spawn_docker_container(container_name, retries=3, language='en'):
    container_name = parse_container_name(container_name)

    for count in range(0, retries):
        try:
            spawn_docker_process, p9000 = spawn_docker_container(container_name, language=language)
            if spawn_docker_process is not None:
                time.sleep(5)  # Sleep for 5 seconds.
                poll_response = spawn_docker_process.poll()
                # check for success - in this case we want the subprocess to still be running
                if poll_response is None:
                    # Successful run
                    print(f'Spawned container: {container_name}')

                    container_id = get_container_id(container_name)

                    save_log(container_name)
                    return True, p9000, container_id
                else:
                    print(f'Tried but failed to spawn container: {container_name} on port {p9000}')
                    print(f'Poll response: {poll_response}')
        except Exception as e:
            print('Failed to spawn doker container')
            print(f'Error: {e}')

    print(f'Tried spawning a container {retries} times but failed!')
    return False, -1, None
