import json

from torch.multiprocessing import ProcessExitedException

from projectcode.training import pickled_train_and_evaluate
from loguru import logger
import time
import main_config
import requests
from projectcode.server.database import (
    get_task,
    pickle_load_from_str,
    get_incomplete_task,
    Task,
    pickle_to_str,
    get_own_inprogress_task,
    connect,
    get_inprogress_task,
)
from projectcode.datasets.cifar10 import CifarDataset
import datetime
from google.api_core.exceptions import ServiceUnavailable


class Client:
    def __init__(self, client_id: str, client_stream: str = "tpu"):
        self.headers = {
            "User-Agent": "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_11_2) AppleWebKit/601.3.9 (KHTML, like Gecko) Version/9.0.2 Safari/601.3.9"
        }
        self.client_id = client_id
        self.ds = None
        if "GPU" in client_id:
            self.client_stream = "gpu"
        else:
            self.client_stream = client_stream

        if self.client_stream == 'tpu':
            import torch_xla

    def do_task(self, task: Task):

        spec = pickle_load_from_str(task.pkl_data)

        return pickled_train_and_evaluate(spec=spec, client_id=self.client_id)

    def report_task(self, task: Task, results: dict = None, success: bool = True):
        connect()

        if get_task(key = task.key) is None:
            logger.warning("Task completed or failed, but not longer exists on the server")
        else:
            if success:
                if task.state == "complete":
                    logger.warning("Task already completed by another")
                else:


                    task.state = "complete"
                    task.result_pkl_data = pickle_to_str(obj=results)
                    task.save()
            else:
                if task.extra_workers != 0:
                    logger.info(
                        "There are other workers on this task. Attempting to restart the same task."
                    )
                    logger.info("Successfully Reported Failure of Task")

                else:
                    task.state = "incomplete"
                    task.save()
                    logger.info("Successfully Reported Failure of Task")

        return

    def get_task(self) -> Task:
        connect()

        task = get_own_inprogress_task(self.client_id, client_stream=self.client_stream)
        if task is None:

            task = get_incomplete_task(client_stream=self.client_stream)
            if task is None:
                # here, we get a task that's already being worked on so that if the first worker fails, the second worker
                # has already started working on the task and has made progress
                logger.info("Trying to Double Up")
                task = get_inprogress_task(client_stream=self.client_stream)
                if task is not None:
                    task.extra_workers = 1
                    task.save()
                    logger.info("Doubling up on a task")

        else:
            logger.info("Resuming own incomplete task")

        if task is not None:

            task.state = "inprogress"
            task.assigned_to = self.client_id
            task.assigned_at = datetime.datetime.now()

            logger.info(task)

            task.save()

        return task

    def prepare(self):
        self.ds = CifarDataset()

    def run(self):

        self.prepare()

        logger.info("Beginning training loop")

        while True:
            logger.info("Getting Task")

            task = self.get_task()

            if task is None:
                logger.info(
                    "No Task Found. Retrying in %i seconds"
                    % main_config.client.wait_time
                )
                time.sleep(main_config.client.wait_time)
            else:

                # TODO: convert to context manager to guarantee reporting of task failure

                sucess = False

                try:
                    results = self.do_task(task)
                    logger.info("Completed Execution of Task")
                    self.report_task(task, results)
                    logger.info("Reported Task Completion")
                    sucess = True

                except Exception as e:

                    if type(e) == ProcessExitedException:
                        logger.warning("Task failed due to Torch Multiprocessing")
                    else:
                        logger.warning("Task failed")

                    self.report_task(task, success=False)

                    logger.exception("Logging Exception")

                    if type(e) == ProcessExitedException:
                        pass
                    elif type(e) == ServiceUnavailable:
                        pass
                    else:
                        raise

                finally:

                    if not sucess:
                        logger.warning("Task failed")
                        self.report_task(task, success=False)
