import sys
import requests
import threading
import warnings
from binascii import hexlify
from datetime import datetime, timezone
from typing import Callable, Optional
from time import sleep, time
from queue import PriorityQueue

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"
JOB_ENV_REPO = "us-east4-docker.pkg.dev/burla-389321/burla-job-environments"

# 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: 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.")

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

    api_key = api_key or 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"]

    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"

    data = {
        "key": api_key,
        "func_pkl_hex_ascii": hexlify(dill.dumps(function_.__code__)).decode("ascii"),
        "inputs": inputs,
        "image": image or default_docker_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"]
