import requests
import threading
import warnings
from datetime import datetime, timezone
from typing import Callable
from time import sleep, time
from queue import PriorityQueue

import binascii
import dill

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

BURLA_SERVICE_URL = "https://burla-webservice-gxc7eo4ala-uc.a.run.app"

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

TZ = timezone.utc
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, last_timestamp: str, attempt=0):
    try:
        response = requests.get(f"{BURLA_SERVICE_URL}/v1/jobs/{job_id}/{last_timestamp}")
        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, last_timestamp, attempt=attempt + 1)
        else:
            raise e


def remote_parallel_map(
    function_: Callable,
    inputs: list,
    cpus: int = 0,
    image: str = "us-east4-docker.pkg.dev/burla-389321/burla-images/burla:1",
):
    if not function_.__annotations__.get("return") is str:
        raise AttributeError("Please add the return type annotation 'str' to your input function.")

    if cpus > 100:
        raise AttributeError("Currently unable to satisfy requests for >100 computers")

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

    func_pkl_hex_ascii = binascii.hexlify(dill.dumps(function_.__code__)).decode("ascii")
    data = {
        "key": api_key,
        "func_pkl_hex_ascii": func_pkl_hex_ascii,
        "inputs": inputs,
        "image": image,
        "num_computers": cpus,
    }
    response = requests.post(f"{BURLA_SERVICE_URL}/v1/jobs/{job_id}", json=data)
    response.raise_for_status()

    job_start_time = time()
    last_timestamp = "2023-07-01T00:00:00.000001Z"
    epoch_time = 1688169600.000001  # same-ish as above
    job_is_running = True
    job_timed_out = False

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

    # loop until job finishes, or times out
    while job_is_running and (not job_timed_out):
        sleep(JOB_STATUS_POLL_RATE_SEC)

        job = _get_job_info(job_id, last_timestamp)

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

        last_timestamp = datetime.fromtimestamp(epoch_time, TZ).isoformat().replace("+00:00", "Z")
        job_is_running = job["job_is_done"] == False
        job_timed_out = (time() - job_start_time) > (TIMEOUT_MIN * 60)

    stop_event.set()
    log_thread.join()

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

    return job["result"]
