import base64
import sys
import time
from urllib.parse import urlparse

import requests

from biolib import api, utils
from biolib.biolib_api_client import JobDict, StartedCloudJobReturnedOnJob
from biolib.biolib_binary_format.stdout_and_stderr import StdoutAndStderr
from biolib.biolib_errors import BioLibError
from biolib.biolib_logging import logger, logger_no_user_data
from biolib.compute_node.utils import SystemExceptionCodeMap, SystemExceptionCodes
from biolib.jobs.job_result import JobResult
from biolib.typing_utils import Optional, List, cast
from biolib.utils import IS_RUNNING_IN_NOTEBOOK


class Job:

    def __init__(self, uuid: str, auth_token: Optional[str] = None):
        self._uuid: str = uuid
        self._auth_token: Optional[str] = auth_token

        self._job_dict: Optional[JobDict] = None
        self._result: Optional[JobResult] = None

    def __str__(self) -> str:
        return self._uuid

    @property
    def result(self) -> JobResult:
        if not self._result:
            self._result = JobResult(job_uuid=self._uuid, job_auth_token=self._auth_token)

        return self._result

    def get_stdout(self) -> bytes:
        return self.result.get_stdout()

    def get_stderr(self) -> bytes:
        return self.result.get_stderr()

    def get_exit_code(self) -> int:
        return self.result.get_exit_code()

    def save_files(self, output_dir: str) -> None:
        self.result.save_files(output_dir=output_dir)

    def get_status(self) -> str:
        return self._get_job_dict(force_refetch=True)['state']

    @staticmethod
    def print_logs_packages(stdout_and_stderr_packages_b64):
        for stdout_and_stderr_package_b64 in stdout_and_stderr_packages_b64:
            stdout_and_stderr_package = base64.b64decode(stdout_and_stderr_package_b64)
            stdout_and_stderr = StdoutAndStderr(stdout_and_stderr_package).deserialize()

            sys.stdout.write(stdout_and_stderr.decode())
            if not IS_RUNNING_IN_NOTEBOOK:  # for some reason flushing in jupyter notebooks breaks \r handling
                sys.stdout.flush()
        # flush after having processed all packages
        sys.stdout.flush()

    def stream_logs(self) -> None:
        self._stream_logs()

    def _stream_logs(self, enable_print: bool = True) -> None:
        cloud_job = self._get_cloud_job_awaiting_started()
        compute_node_url = cloud_job['compute_node_url']
        logger_no_user_data.debug(f'Using compute node URL "{compute_node_url}"')

        if utils.BIOLIB_CLOUD_BASE_URL:
            compute_node_url = utils.BIOLIB_CLOUD_BASE_URL + str(urlparse(compute_node_url).path)
            logger_no_user_data.debug(f'Using cloud proxy URL from env var BIOLIB_CLOUD_BASE_URL: {compute_node_url}')

        if enable_print:
            self._print_full_logs(node_url=compute_node_url)

        final_status_messages: List[str] = []
        while True:
            time.sleep(2)
            status_json = self._get_job_status_from_compute_node(compute_node_url)
            job_is_completed = status_json['is_completed']
            for status_update in status_json['status_updates']:
                # If the job is completed, print the log messages after all stdout and stderr has been written
                if job_is_completed:
                    final_status_messages.append(status_update['log_message'])
                else:
                    # Print the status before writing stdout and stderr
                    logger.info(f'Cloud: {status_update["log_message"]}')

            if 'stdout_and_stderr_packages_b64' and enable_print:
                self.print_logs_packages(status_json['stdout_and_stderr_packages_b64'])

            if 'error_code' in status_json:
                error_code = status_json['error_code']
                error_message = SystemExceptionCodeMap.get(error_code, f'Unknown error code {error_code}')

                raise BioLibError(f'Cloud: {error_message}')

            if job_is_completed:
                break

        # Print the final log messages after stdout and stderr has been written
        for message in final_status_messages:
            logger.info(f'Cloud: {message}')

    def _get_job_dict(self, force_refetch: bool = False) -> JobDict:
        if not self._job_dict or force_refetch:
            job_dict: JobDict = api.client.get(
                url=f'/jobs/{self._uuid}/',
                headers={'Job-Auth-Token': self._auth_token} if self._auth_token else None,
            ).json()
            self._job_dict = job_dict

        return self._job_dict

    def _print_full_logs(self, node_url: str) -> None:
        try:
            response = requests.get(f'{node_url}/v1/job/{self._uuid}/status/?logs=full')
            if not response.ok:
                raise BioLibError(response.content)
        except Exception as error:
            logger.error(f'Could not get full streamed logs due to: {error}')
            raise BioLibError from error

        self.print_logs_packages(response.json()['streamed_logs_packages_b64'])

    def _get_cloud_job_awaiting_started(self) -> StartedCloudJobReturnedOnJob:
        max_retry_attempts = 100
        retry_interval_seconds = 10
        for _ in range(max_retry_attempts):
            cloud_job = self._get_job_dict(force_refetch=True)['cloud_job']
            if not cloud_job:
                raise BioLibError(f'Job {self._uuid} did not register correctly. Try creating a new job')

            if cloud_job['finished_at']:
                raise BioLibError(f'The job {self._uuid} is already finished. Get its output by calling `.result()`')

            if cloud_job and cloud_job['started_at']:
                return cast(StartedCloudJobReturnedOnJob, cloud_job)

            logger.info('Cloud: The job has been queued. Please wait...')
            time.sleep(retry_interval_seconds)

        raise BioLibError('Cloud: Timed out waiting for the job to start.')

    def _get_job_status_from_compute_node(self, compute_node_url):
        for _ in range(15):
            try:
                response = requests.get(f'{compute_node_url}/v1/job/{self._uuid}/status/')
                if not response.ok:
                    raise BioLibError(response.content)
                return response.json()

            except Exception:  # pylint: disable=broad-except
                cloud_job = self._get_job_dict(force_refetch=True)['cloud_job']
                logger.debug("Failed to get status from compute node, retrying...")
                if cloud_job['finished_at']:
                    logger.debug("Job no longer exists on compute node, checking for error...")
                    if cloud_job['error_code'] != SystemExceptionCodes.COMPLETED_SUCCESSFULLY.value:
                        error_message = SystemExceptionCodeMap.get(
                            cloud_job['error_code'],
                            f'Unknown error code {cloud_job["error_code"]}'
                        )
                        raise BioLibError(f'Cloud: {error_message}') from None
                    else:
                        logger.info(f'The job {self._uuid} is finished. Get its output by calling `.result()`')
                        return

                time.sleep(2)

        raise BioLibError(
            'Failed to stream logs, did you lose internet connection?\n'
             'Call `.stream_logs()` on your job to resume streaming logs.'
        )
