import os
import sys
import requests
import warnings
import pickle
from six import reraise
from threading import Thread, Event
from typing import Callable, Optional
from time import sleep, time
from queue import PriorityQueue

import cloudpickle
from yaspin import yaspin, Spinner
from tblib import pickling_support

from burla import _BURLA_SERVICE_URL
from burla._logstream import print_logs_from_queue
from burla._env_inspection import get_pip_packages, get_function_dependencies
from burla._auth import auth_headers_from_local_config, AuthException, login_required
from burla._helpers import (
    nopath_warning,
    JobTimeoutError,
    InstallError,
    ServerError,
    StatusMessage,
    upload_inputs,
    download_outputs,
)

pickling_support.install()
warnings.formatwarning = nopath_warning


MAX_CPUS = 2000  # please ask before you increase this <3
MAX_GPUS = 300  # please ask before you increase this <3

JOB_STATUS_POLL_RATE_SEC = 6  # how often to check job status
TIMEOUT_MIN = 60 * 12  # max time a Burla job can run for
IN_COLAB = os.getenv("COLAB_RELEASE_TAG") is not None

BYTES_HEADER = {"Content-Type": "application/octet-stream"}


def _get_job_info_since(epoch: int, job_id: str, headers: dict):
    """Raises: ServerError, InstallError, HTTPError, or re-raised error from UDF"""
    response = requests.get(f"{_BURLA_SERVICE_URL}/v1/jobs/{job_id}/{epoch}", headers=headers)
    if response.status_code == 401:
        AuthException()
    elif response.status_code >= 500:
        raise ServerError()
    elif response.status_code >= 300:
        response.raise_for_status()

    job = response.json()
    if job.get("udf_error"):
        error = pickle.loads(bytes.fromhex(job.get("udf_error")))
        reraise(*error)
    elif job.get("install_error"):
        raise InstallError(job["install_error"])

    if job.get("output_urls"):
        return job["udf_started"], job["logs"], job["output_urls"]
    else:
        return job["udf_started"], job["logs"], None


def _start_job(
    function_: Callable,
    inputs: list,
    auth_headers: dict,
    verbose: bool,
    spinner: Spinner,
    func_cpu: int,
    func_ram: int,
    func_gpu: int,
    parallelism: int,
    image: Optional[str] = None,
) -> str:
    if verbose:
        spinner.text = StatusMessage.preparing()
        spinner.start()

    installed_packages = list(get_pip_packages())
    imported_modules = list(get_function_dependencies(function_))
    required_packages = [pkg for pkg in installed_packages if pkg["name"] in imported_modules]

    payload = {
        "n_inputs": len(inputs),
        "image": image,
        "func_cpu": func_cpu,
        "func_ram": func_ram,
        "func_gpu": func_gpu,
        "parallelism": parallelism,
        "python_version": f"3.8" if sys.version_info.minor < 8 else f"3.{sys.version_info.minor}",
        "packages": None if image else required_packages,
    }
    response = requests.post(f"{_BURLA_SERVICE_URL}/v1/jobs/", json=payload, headers=auth_headers)
    if response.status_code == 401:
        raise AuthException()
    else:
        response.raise_for_status()
        response_json = response.json()
        job_id = response_json["job_id"]

    # download list of urls to upload to, these urls are big and sometimes there are a lot of them.
    response = requests.get(response_json["input_urls_url"], headers=BYTES_HEADER)
    response.raise_for_status()
    input_urls = pickle.loads(response.content)

    spinner.text = StatusMessage.uploading_inputs
    upload_inputs(input_urls, inputs)

    spinner.text = StatusMessage.uploading_function
    function_pkl = cloudpickle.dumps(function_)
    requests.put(response_json["function_url"], headers=BYTES_HEADER, data=function_pkl)

    spinner.text = StatusMessage.preparing()
    response = requests.post(f"{_BURLA_SERVICE_URL}/v1/jobs/{job_id}", headers=auth_headers)
    response.raise_for_status()
    return job_id


def _watch_job(job_id: str, job_started_time: int, headers: dict, verbose: bool, spinner: Spinner):
    last_epoch = job_started_time
    epoch = last_epoch
    output_urls = None
    job_timed_out = False

    # Start printing logs generated by this job using a separate thread.
    print_queue = PriorityQueue()
    stop_event = Event()
    args = (print_queue, stop_event, spinner)
    log_thread = Thread(target=print_logs_from_queue, args=args, daemon=True)
    log_thread.start()
    while (not output_urls) and (not job_timed_out):
        sleep(JOB_STATUS_POLL_RATE_SEC)
        udf_started, logs, output_urls = _get_job_info_since(last_epoch, job_id, headers)

        for epoch, log_message in logs:
            print_queue.put((epoch, log_message))

        if verbose and udf_started:
            spinner.text = StatusMessage.running()

        last_epoch = epoch
        job_timed_out = (time() - job_started_time) > (TIMEOUT_MIN * 60)

    stop_event.set()
    log_thread.join()

    if job_timed_out:
        raise JobTimeoutError(job_id=job_id, timeout=TIMEOUT_MIN)
    return output_urls


@login_required
def remote_parallel_map(
    function_: Callable,
    inputs: list,
    verbose: bool = True,
    image: Optional[str] = None,
    func_cpu: int = 1,
    func_ram: int = 1,
    func_gpu: int = 0,
    parallelism: int = -1,
):
    if (func_cpu > 96) or (func_cpu < 1):
        raise ValueError("CPU per function call must be one of [1.. 96]")
    if (func_ram > 624) or (func_ram < 1):
        raise ValueError("RAM per function call must be one of [1.. 624]")
    if (func_gpu > 4) or (func_gpu < 0):
        raise ValueError("GPU per function call must be one of [0.. 4]")

    n_inputs = len(inputs)
    parallelism = n_inputs if parallelism == -1 else parallelism
    requested_cpu = parallelism * func_cpu
    requested_gpu = parallelism * func_gpu

    if (requested_cpu > MAX_CPUS) and not (requested_gpu > MAX_GPUS):
        parallelism = MAX_CPUS // func_cpu
        warnings.warn(
            f"Limiting parallelism to {parallelism} to stay under current limit of {MAX_CPUS} CPUs."
        )
    if requested_gpu > MAX_GPUS:
        parallelism = MAX_GPUS // func_gpu
        warnings.warn(
            f"Limiting parallelism to {parallelism} to stay under current limit of {MAX_GPUS} GPUs."
        )

    max_concurrency = 10000
    input_batches = [inputs[i : i + max_concurrency] for i in range(0, n_inputs, max_concurrency)]
    if len(input_batches) > 1:
        msg = f"Because a maximum of {max_concurrency} inputs can be processed at a time, "
        msg += f"these inputs will be processed in {len(input_batches)} separate batches. "
        msg += "Consider using a lower number of larger inputs, this may reduce runtime."
        warnings.warn(msg)

    spinner = yaspin()
    StatusMessage.function_name = function_.__name__
    StatusMessage.n_inputs = len(inputs) if len(inputs) < max_concurrency else max_concurrency
    StatusMessage.total_cpus = parallelism * func_cpu
    StatusMessage.total_gpus = parallelism * func_gpu

    auth_headers = auth_headers_from_local_config()

    return_values = []
    for input_batch in input_batches:
        try:
            start_time = int(time())
            job_id = _start_job(
                function_=function_,
                inputs=input_batch,
                auth_headers=auth_headers,
                verbose=verbose,
                spinner=spinner,
                func_cpu=func_cpu,
                func_ram=func_ram,
                func_gpu=func_gpu,
                parallelism=parallelism,
                image=image,
            )
            output_urls = _watch_job(job_id, start_time, auth_headers, verbose, spinner)

            spinner.text = StatusMessage.downloading
            for output_pkl in download_outputs(output_urls):
                output = cloudpickle.loads(output_pkl)
                if output["pickled_with_cloudpickle"]:
                    return_values.append(cloudpickle.loads(output["return_value_pkl"]))
                elif output["return_type"] == "<class 'spacy.tokens.doc.Doc'>":
                    if (not globals().get("Doc")) and (not globals().get("Vocab")):
                        from spacy.tokens.doc import Doc
                        from spacy.vocab import Vocab
                    return_values.append(Doc(Vocab()).from_bytes(output["return_value_pkl"]))

        except Exception as e:
            spinner.stop()
            raise e

    if verbose:
        spinner.text = "Done!"
        spinner.ok("✔")

    all_return_values_are_none = all(value is None for value in return_values)
    if not all_return_values_are_none:
        return return_values
