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

import dill
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._auth import auth_headers_from_local_config
from burla._env_inspection import get_pip_packages, get_imported_modules
from burla._helpers import (
    nopath_warning,
    concurrency_warning_message,
    JobTimeoutError,
    InstallError,
    ServerError,
    StatusMessage,
)

pickling_support.install()
warnings.formatwarning = nopath_warning

MAX_CONCURRENCY = 1000
JOB_STATUS_POLL_RATE_SEC = 6  # how often to check job status
TIMEOUT_MIN = 60 * 12  # max time a Burla job can run for


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 >= 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("return_values"):
        return_values = [dill.loads(bytes.fromhex(value)) for value in job["return_values"]]
        return job["udf_started"], job["logs"], return_values
    else:
        return job["udf_started"], job["logs"], None


def _start_job(
    function_: Callable,
    inputs: list,
    headers: dict,
    verbose: bool,
    spinner: Spinner,
    image: Optional[str] = None,
) -> str:
    if verbose:
        spinner.text = StatusMessage.PREPARING
        spinner.start()

    pickled_function = dill.dumps(function_, recurse=True)
    installed_packages = list(get_pip_packages())
    imported_modules = get_imported_modules(pickled_function)
    required_packages = [pkg for pkg in installed_packages if pkg["name"] in imported_modules]

    payload = {
        "function_pkl_hex": pickled_function.hex(),
        "inputs": [dill.dumps(input_).hex() for input_ in inputs],
        "image": image,
        "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=headers)
    response.raise_for_status()
    return response.json()["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
    return_values = 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 return_values) and (not job_timed_out):
        sleep(JOB_STATUS_POLL_RATE_SEC)
        udf_started, logs, return_values = _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 return_values


def remote_parallel_map(
    function_: Callable,
    inputs: list,
    verbose: bool = True,
    image: Optional[str] = None,
):
    n_inputs = len(inputs)
    n_batches = math.ceil(n_inputs / MAX_CONCURRENCY)
    if n_inputs > MAX_CONCURRENCY:
        warnings.warn(concurrency_warning_message(n_batches, n_inputs, MAX_CONCURRENCY))

    spinner = yaspin()
    StatusMessage.PREPARING = f"Preparing to run `{function_.__name__}` on {n_inputs} CPUs"
    StatusMessage.RUNNING = f"Running `{function_.__name__}` on {n_inputs} CPUs"

    headers = auth_headers_from_local_config()
    input_batches = [inputs[i : i + MAX_CONCURRENCY] for i in range(0, n_inputs, MAX_CONCURRENCY)]

    udf_return_values = []
    for input_batch in input_batches:
        try:
            start_time = int(time())
            job_id = _start_job(function_, input_batch, headers, verbose, spinner, image)
            udf_return_values = _watch_job(job_id, start_time, headers, verbose, spinner)
        except Exception as e:
            spinner.stop()
            raise e
        udf_return_values.extend(udf_return_values)

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

    all_values_are_none = all(value is None for value in udf_return_values)
    if not all_values_are_none:
        return udf_return_values
