import os
import sys
import requests
import warnings
import pickle
import math
from six import reraise
from threading import Thread, Event
from typing import Callable, Optional, Literal
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._env_inspection import get_pip_packages, get_function_dependencies
from burla._auth import (
    auth_headers_from_local_config,
    login_credentials_missing,
    AuthException,
    login,
)
from burla._helpers import (
    nopath_warning,
    concurrency_warning_message,
    JobTimeoutError,
    InstallError,
    ServerError,
    StatusMessage,
)

pickling_support.install()
warnings.formatwarning = nopath_warning

MAX_CONCURRENCY_NOGPU = 1000
MAX_CONCURRENCY_GPU = 300
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


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("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,
    auth_headers: dict,
    verbose: bool,
    spinner: Spinner,
    gpu: bool,
    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 = list(get_function_dependencies(function_))
    required_packages = [pkg for pkg in installed_packages if pkg["name"] in imported_modules]

    payload = {
        "function_pkl_hex": pickled_function.hex(),
        "n_inputs": len(inputs),
        "image": image,
        "gpu": gpu,
        "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()
        job_id = response.json()["job_id"]

    input_upload_headers = {"Content-Type": "application/octet-stream"}
    for input_, url in zip(inputs, response.json()["input_upload_urls"]):
        response = requests.put(url, headers=input_upload_headers, data=dill.dumps(input_))
        response.raise_for_status()

    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
    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,
    gpu: bool = False,
):
    n_inputs = len(inputs)
    max_concurrency = MAX_CONCURRENCY_GPU if gpu else MAX_CONCURRENCY_NOGPU
    n_batches = math.ceil(n_inputs / max_concurrency)
    if n_inputs > max_concurrency:
        warnings.warn(concurrency_warning_message(n_batches, n_inputs, max_concurrency, gpu))

    if login_credentials_missing() and IN_COLAB:
        login()
    elif login_credentials_missing():
        raise AuthException()

    spinner = yaspin()
    StatusMessage.PREPARING = f"Preparing to run `{function_.__name__}` on {n_inputs} "
    StatusMessage.PREPARING += "GPUs, " if gpu else "CPUs, "
    StatusMessage.PREPARING += f"this may take {'several' if gpu else 'a few'} minutes."
    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, gpu, image)
            udf_return_values_batch = _watch_job(job_id, start_time, headers, verbose, spinner)
        except Exception as e:
            spinner.stop()
            raise e
        udf_return_values.extend(udf_return_values_batch)

    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
