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

import requests

import biolib.api

from biolib import utils
from biolib.biolib_api_client.auth import BearerAuth
from biolib.biolib_api_client import BiolibApiClient, CreatedJobDict, CloudJob
from biolib.biolib_binary_format.stdout_and_stderr import StdoutAndStderr
from biolib.biolib_errors import BioLibError, RetryLimitException, StorageDownloadFailed
from biolib.compute_node.utils import SystemExceptionCodeMap
from biolib.biolib_logging import logger
from biolib.utils import BIOLIB_PACKAGE_VERSION, IS_RUNNING_IN_NOTEBOOK
from biolib.typing_utils import List, TypedDict, Optional, Literal


class PresignedS3UploadLinkResponse(TypedDict):
    presigned_upload_url: str


class PresignedS3DownloadLinkResponse(TypedDict):
    presigned_download_url: str


def _get_user_info() -> Optional[str]:
    if utils.BASE_URL_IS_PUBLIC_BIOLIB:
        return None

    enterprise_agent_info_opt_env_vars = ['DOMINO_STARTING_USERNAME', 'USER']

    for env_var in enterprise_agent_info_opt_env_vars:
        env_var_value = os.getenv(env_var)
        if env_var_value:
            return env_var_value

    return None


class BiolibJobApi:

    @staticmethod
    def create(
            app_version_id,
            override_command=False,
            caller_job=None,
            machine='',
            experiment_uuid: Optional[str] = None
    ):
        data = {
            'app_version_id': app_version_id,
            'client_type': 'biolib-python',
            'client_version': BIOLIB_PACKAGE_VERSION,
            'client_opt_user_info': _get_user_info(),
        }

        if override_command:
            data.update({
                'arguments_override_command': override_command
            })

        if caller_job:
            data['caller_job'] = caller_job

        if machine:
            data.update({
                'requested_machine': machine
            })

        if experiment_uuid:
            data['experiment_uuid'] = experiment_uuid

        response = biolib.api._client.post(path='/jobs/', data=data)  # pylint: disable=protected-access

        return response.json()

    @staticmethod
    def update_state(job_id, state):
        response = requests.patch(
            f'{BiolibApiClient.get().base_url}/api/jobs/{job_id}/',
            json={'state': state},
            auth=BearerAuth(BiolibApiClient.get().access_token)
        )

        # TODO: Error handling with response object
        if not response.ok:
            raise BioLibError(response.content)

        return response.json()

    @staticmethod
    def create_cloud_job(job_id: str, result_name_prefix: Optional[str]) -> CloudJob:
        response = None
        data = {'job_id': job_id}
        if result_name_prefix:
            data['result_name_prefix'] = result_name_prefix

        for retry in range(4):
            try:
                response = requests.post(
                    f'{BiolibApiClient.get().base_url}/api/jobs/cloud/',
                    json=data,
                    auth=BearerAuth(BiolibApiClient.get().access_token)
                )

                if response.status_code == 503:
                    raise RetryLimitException(response.content)
                # Handle possible validation errors from backend
                elif not response.ok:
                    raise BioLibError(response.text)

                break

            except RetryLimitException as retry_exception:  # pylint: disable=broad-except
                if retry > 3:
                    raise BioLibError('Reached retry limit for cloud job creation') from retry_exception
                time.sleep(1)

        if not response:
            raise BioLibError('Could not create new cloud job')

        cloud_job: CloudJob = response.json()
        return cloud_job

    @staticmethod
    def save_compute_node_job(job, module_name, access_token, node_url):
        response = requests.post(
            f'{node_url}/v1/job/',
            json={'module_name': module_name, 'job': job, 'access_token': access_token}
        )

        if not response.ok:
            raise BioLibError(response.content)

        return response.json

    @staticmethod
    def start_local_job(job_id, module_input_serialized, node_url):
        response = requests.post(
            f'{node_url}/v1/job/{job_id}/start/',
            data=module_input_serialized,
        )

        if not response.ok:
            raise BioLibError(response.content)

    # TODO: Deprecate this
    @staticmethod
    def await_local_compute_node_status(
            retry_interval_seconds: float,
            retry_limit_minutes: float,
            status_to_await: str,
            compute_type: str,
            node_url: str,
            job: CreatedJobDict,
    ):
        status_max_retry_attempts = int(retry_limit_minutes * 60 / retry_interval_seconds)
        status_is_reached = False
        final_status_messages: List[str] = []

        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()

        for _ in range(status_max_retry_attempts):
            response = requests.get(f'{node_url}/v1/job/{job["public_id"]}/status/')
            if not response.ok:
                raise Exception(response.content)

            status_json = response.json()

            for status_update in status_json['status_updates']:
                if status_update.get('log_message') == status_to_await:
                    status_is_reached = True

                # If the status is reached print the log messages after all stdout and stderr has been written
                if status_is_reached:
                    final_status_messages.append(status_update['log_message'])
                else:
                    # Print the status before writing stdout and stderr
                    logger.info(f'{compute_type}: {status_update["log_message"]}')

            app_version = job['app_version']
            if 'stdout_and_stderr_packages_b64' in status_json and utils.STREAM_STDOUT and \
                    (app_version.get('main_output_file') or app_version.get('stdout_render_type') == 'text'):
                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'{compute_type}: {error_message}')

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

                return

            time.sleep(retry_interval_seconds)

        raise BioLibError(f'{compute_type}: Failed to get results: Retry limit exceeded')

    @staticmethod
    def get_cloud_result(job_id: str, node_url: str) -> bytes:
        for _ in range(5):
            try:
                response = requests.get(f'{node_url}/v1/job/{job_id}/result/', timeout=600)
                if response.ok:
                    return response.content
                else:
                    logger.error(
                        f'Getting result for job {job_id} failed with status {response.status_code} '
                        f'and error: {response.content.decode()}\nRetrying...'
                    )
            except Exception as error:  # pylint: disable=broad-except
                logger.error(f'Getting result for job {job_id} failed with error {error}\nRetrying...')

            time.sleep(2)

        raise BioLibError(f'Max retries hit, when getting result for job {job_id}')

    @staticmethod
    def get_enclave_json(biolib_base_url):
        response = requests.get(
            f'{biolib_base_url}/info-files/biolib-enclave.json',
        )
        return response.json()

    @staticmethod
    def get_job_storage_download_url(
            job_uuid: str,
            job_auth_token: str,
            storage_type: Literal['input', 'results'],
    ) -> str:
        response = requests.get(
            f'{BiolibApiClient.get().base_url}/api/jobs/{job_uuid}/storage/{storage_type}/download/',
            auth=BearerAuth(BiolibApiClient.get().access_token),
            headers={
                'Job-Auth-Token': job_auth_token
            }
        )

        if not response.ok:
            if response.status_code == 404:
                raise StorageDownloadFailed(response.content)
            else:
                raise BioLibError(response.content)

        presigned_s3_download_link_response: PresignedS3DownloadLinkResponse = response.json()
        presigned_download_url = presigned_s3_download_link_response['presigned_download_url']

        app_caller_proxy_job_storage_base_url = os.getenv('BIOLIB_CLOUD_JOB_STORAGE_BASE_URL', '')
        if app_caller_proxy_job_storage_base_url:
            # Done to hit App Caller Proxy when downloading result from inside an app
            parsed_url = urlparse(presigned_download_url)
            presigned_download_url = f'{app_caller_proxy_job_storage_base_url}{parsed_url.path}?{parsed_url.query}'

        return presigned_download_url
