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 tblib import pickling_support

from burla._logstream import print_logs_from_queue
from burla.config import load_api_key_from_local_config

pickling_support.install()

BURLA_SERVICE_URL = "https://burla-webservice-zqhes3whbq-uc.a.run.app"  # prod
JOB_ENV_REPO = "us-docker.pkg.dev/burla-prod/burla-job-environments"

# BURLA_SERVICE_URL = "https://burla-webservice-y66ufvpuua-uc.a.run.app"  # test
# JOB_ENV_REPO = "us-docker.pkg.dev/burla-test/burla-job-environments"

# BURLA_SERVICE_URL = "http://127.0.0.1:5000"

MAX_CONCURRENCY = 100  # If your snooping on my code and want to increase this, don't, it wont work.
JOB_STATUS_POLL_RATE_SEC = 6  # how often to check for job completion
TIMEOUT_MIN = 60 * 12  # max time a Burla job can run for


class JobTimeoutError(Exception):
    def __init__(self, job_id, timeout):
        super().__init__(f"Burla job with id: '{job_id}' timed out after {timeout} seconds.")


def _get_job_info(job_id: str, epoch: int, headers: dict, attempt=0):
    try:
        response = requests.get(f"{BURLA_SERVICE_URL}/v1/jobs/{job_id}/{epoch}", headers=headers)
        response.raise_for_status()
        return response.json()
    except requests.exceptions.HTTPError as e:
        if str(response.status_code).startswith("5") and attempt != 3:
            warnings.warn(
                (
                    f"Received {response.status_code} response from server checking status of job: "
                    f"{job_id} Retrying..."
                )
            )
            sleep(2)
            return _get_job_info(job_id, epoch, headers, attempt=attempt + 1)
        else:
            raise e


def start_job(function_: Callable, inputs: list, image: str, job_id: str, headers: dict):
    function_pkl_hex = dill.dumps(function_).hex()
    data = {"function_pkl_hex": function_pkl_hex, "inputs": inputs, "image": image}
    response = requests.post(f"{BURLA_SERVICE_URL}/v1/jobs/{job_id}", json=data, headers=headers)
    response.raise_for_status()


def remote_parallel_map(
    function_: Callable,
    inputs: list,
    image: Optional[str] = None,
    api_key: Optional[str] = None,
):
    n_batches = math.ceil(len(inputs) / MAX_CONCURRENCY)
    if len(inputs) > MAX_CONCURRENCY:
        warnings.warn(
            (
                f"Because the current maximum concurrency is {MAX_CONCURRENCY} and you submitted "
                f"{len(inputs)} inputs, these inputs will be processed in {n_batches} separate "
                "batches. We are working hard to increase this concurrency limit."
            )
        )
    outputs = []
    input_batches = [
        inputs[i : i + MAX_CONCURRENCY] for i in range(0, len(inputs), MAX_CONCURRENCY)
    ]
    for input_batch in input_batches:
        outputs.extend(_remote_parallel_map(function_, input_batch, image, api_key))
    return outputs


def _remote_parallel_map(
    function_: Callable,
    inputs: list,
    image: Optional[str] = None,
    api_key: Optional[str] = None,
):
    if not function_.__annotations__.get("return") is str:
        raise AttributeError("Please add the return type annotation 'str' to your input function.")

    # https://www.rfc-editor.org/rfc/rfc7235 <- specifies "correct" api key location
    headers = {"Authorization": f"Bearer {api_key or load_api_key_from_local_config()}"}

    response = requests.post(f"{BURLA_SERVICE_URL}/v1/jobs/", json={}, headers=headers)
    response.raise_for_status()
    job_id = response.json()["job_id"]

    user_python_version = f"{sys.version_info.major}.{sys.version_info.minor}"
    default_docker_image = f"{JOB_ENV_REPO}/python{user_python_version}/burla_job_env:latest"
    image = image or default_docker_image  # sets image to default_docker_image if image is None

    # This is on a separate thread so we can print logs before all sub_jobs are running.
    start_job_thread = Thread(target=start_job, args=(function_, inputs, image, job_id, headers))
    job_started_at_epoch = int(time())
    start_job_thread.start()

    last_epoch = job_started_at_epoch
    epoch = last_epoch
    job_is_running = True
    job_timed_out = False

    # Start printing logs generated by this job using a separate thread.
    print_queue = PriorityQueue()
    stop_event = Event()
    log_thread = Thread(target=print_logs_from_queue, args=(print_queue, stop_event))
    log_thread.start()

    # loop until job finishes, or times out
    while job_is_running and (not job_timed_out):
        # TODO: record and subtract time so this loop always takes exactly JOB_STATUS_POLL_RATE_SEC
        # sec to run? This might fix late logs?
        sleep(JOB_STATUS_POLL_RATE_SEC)

        job = _get_job_info(job_id, last_epoch, headers)

        # add all logs to print queue
        for epoch, log_message in job["logs"]:
            print_queue.put((epoch, log_message))

        last_epoch = epoch
        job_is_running = job["job_is_done"] == False
        job_failed = job.get("error") is not None
        job_timed_out = (time() - job_started_at_epoch) > (TIMEOUT_MIN * 60)

    stop_event.set()
    log_thread.join()
    start_job_thread.join()

    if job_failed:
        error = pickle.loads(bytes.fromhex(job.get("error")))
        reraise(*error)

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

    return job["return_values"]
