import os
import re
import io
import base64
import requests
from os.path import join, dirname, isfile
from dotenv import load_dotenv, set_key

import torch
from torch.utils import model_zoo

from .loggers import Logger
from .settings import TEMP_DIR, ME_URL, CLIENTS_URL, MODELS_URL, JOBS_URL, SUBMISSIONS_URL, DATA_URL
from .settings import system_info

# Load environment variables
load_dotenv()

# Create a Logger
logger = Logger()
logger.info(system_info())


def login(username, password):
    """
    Loging in to the FEDBIONET API.
    :param str username: Username to log in. 
    :param password: Password
    :return dict: JSON response. 
    """
    TOKEN = base64.b64encode(f'{username}:{password}'.encode()).decode()
    env_file = join(dirname(os.path.realpath(__file__)), '.env')  # New .env file
    if not isfile(env_file):
        open(env_file, 'w')

    set_key(env_file, 'FEDBIONET_TOKEN', TOKEN)


def get_auth():
    load_dotenv()
    TOKEN = os.environ['FEDBIONET_TOKEN']
    return tuple(base64.b64decode(TOKEN.encode()).decode().split(':'))


def get_uuid_from_url(url):
    """
    Extract the uuid of an object from the URL.
    :param str url: URL where the UUID is contained.
    :return str: UUID as s string.
    """
    return re.findall(r'[0-9a-fA-F]{8}\-[0-9a-fA-F]{4}\-[0-9a-fA-F]{4}\-[0-9a-fA-F]{4}\-[0-9a-fA-F]{12}', url)[0]


def handle_response(res):
    # logger.info(f'Response content: {res.headers}')
    if res.status_code is 204:
        logger.info('Request successful. Status code: 204.')
        return {'detail': 'Request successful. Status code: 204.'}
    elif res.status_code in range(200, 300):
        logger.info('Request successful')
        logger.info(f'Response: {res.json()}')
        return res.json()
    elif res.status_code in range(400, 500):
        logger.error(f'Bad request. Status code: {res.status_code}')
        logger.error(f'Response: {res.json()}')
    else:
        logger.error(f'Bad request. Status code: {res.status_code}')


def get(url, **kwargs):
    logger.debug(f'Sending GET request to: {url}')
    res = requests.get(url, auth=get_auth(), **kwargs)
    return handle_response(res)


def post(url, **kwargs):
    logger.debug(f'Sending POST request to: {url}')
    res = requests.post(url, auth=get_auth(), **kwargs)
    return handle_response(res)


def put(url, **kwargs):
    logger.debug(f'Sending PUT request to: {url}')
    res = requests.put(url, auth=get_auth(), **kwargs)
    return handle_response(res)


def delete(url, **kwargs):
    logger.debug(f'Sending DELETE request to: {url}')
    res = requests.delete(url, auth=get_auth(), **kwargs)
    return handle_response(res)


# ============================================================
# USER DATA
# ============================================================
def me():
    logger.info('Getting profile information...')
    return get(ME_URL)[0]


# ============================================================
# DATA MANAGEMENT
# ============================================================
def get_available_data():
    logger.info(f'Getting available data...')
    return get(DATA_URL)


# ============================================================
# JOBS AND SUBMISSION MANAGEMENT
# ============================================================
def get_jobs():
    logger.info('Getting jobs...')
    return get(JOBS_URL)


def update_job_status(job_url, status):
    logger.info(f'Updating job status: f{job_url} ...')
    job_json = get(job_url)
    job_json['status'] = status
    return put(job_url, json=job_json)


def get_submissions(job_url=None):
    """
    Gets a list of submissions made by the user requesting the list. Include job_url if you want to
    filter the list of submissions by url.

    :param str job_url: Job URL. This will filter the submissions by Job.
    :return list: List of dict-like objects containing the submissions associated to the job_url.
    """
    # Filter submissions by job if job URL is provided
    params = None
    if job_url:
        params = {'job': get_uuid_from_url(job_url)}

    # Get list of submissions
    return get(SUBMISSIONS_URL, params=params)


def post_submission(model, job_url, current_round):
    """
    Submit a new version of your model.
    :param torch.Module model: Torch model object.
    :param str job_url: Job URL associated to the model trained.
    :param int current_round: Version of the trained model (rounds).
    :return:
    """
    logger.info(f'Submitting model...')

    # Save model to buffer
    model_io = io.BytesIO()
    torch.save(model, model_io)
    model_file = ('model.pth', model_io.getvalue())

    # Append files to upload
    files = {
        'model': model_file,
        'log': open(logger.file_path, 'rb')}

    data = {
        'job': job_url,
        'round': current_round
    }
    return post(SUBMISSIONS_URL, files=files, data=data)


# ============================================================
# MODEL MANAGEMENT
# ============================================================
def load_model_from_url(model_url):
    """
    Download torch model and load it as a torch Module.
    :param str model_url: Model .pt or .pth URL
    :return torch.Module: Model loaded
    """
    return model_zoo.load_url(model_url, progress=False)


def load_base_model(job_url):
    """
    Downloads base torch model.
    :param str job_url: URL of the job containing the model.
    :return str: File path to the downloaded model.
    """
    logger.info('Downloading base model...')
    json = get(job_url)  # Get job data
    json = get(json['model'])  # Get model url

    # Download file and return model file path
    return load_model_from_url(json['model_file'])
