import json
import logging.config
import os
import threading
import time
from contextlib import contextmanager
from typing import Optional

import yaml
from requests.exceptions import RequestException

from load.projects.cloud.tank_client import utils
from load.projects.cloud.tank_client.ammo import Ammo
from load.projects.cloud.tank_client.client import TankApiClient, TankapiException
from load.projects.cloud.tank_client.exceptions import LoadTestingFailedPreconditionError, LoadTestingInternalError, \
    LoadTestingNotFoundError, LoadTestingConnectionError, ObjectStorageError, LoadTestingServerError
from load.projects.cloud.tank_client.job import AdditionalJobStatus
from load.projects.cloud.tank_client.job import Job
from load.projects.cloud.tank_client.send_logs import send_log
from load.projects.cloud.tank_client.utils import TankStatus, LogType

LOGGER = logging.getLogger('tank_client')

STOPPING_SLEEP_TIME = 15

LOCK_DIR = '/var/lock'
FINISHED_FILE = 'finish_status.yaml'


class TankClientService:
    config_file_name = 'config'

    def __init__(self, loadtesting_client, tank_api_client, jobs_storage, workdir, sleep_time, logging_host,
                 logging_port):
        os.makedirs(workdir, exist_ok=True)
        self.workdir = workdir

        self.loadtesting_client = loadtesting_client
        self.tank_api_client: TankApiClient = tank_api_client
        self.jobs_storage = jobs_storage
        self.sleep_time = sleep_time
        self.tank_status_report_delay = sleep_time
        self.job: Optional[Job] = None
        self.logging_host = logging_host
        self.logging_port = logging_port
        self.job_pooling_delay = sleep_time

    def report_tank_status(self):
        tank_status_response = self.tank_api_client.get_tank_status()
        tank_status = self.tank_status(tank_status_response)
        LOGGER.debug('Tank status %s', tank_status.name)
        self.loadtesting_client.claim_tank_status(tank_status)

    @contextmanager
    def reporting_tank_status(self):
        stop = threading.Event()

        def worker():
            while not stop.is_set():
                try:
                    self.report_tank_status()
                except Exception:
                    LOGGER.exception("failed to report tank status")
                finally:
                    stop.wait(self.tank_status_report_delay)
            self.report_tank_status()

        t = threading.Thread(target=worker)
        t.start()
        try:
            yield
        finally:
            stop.set()
            t.join()

    @staticmethod
    def get_current_job():
        locks = list(filter(lambda x: x.startswith('lunapark'), os.listdir(LOCK_DIR)))
        if len(locks) != 1:
            return
        with open(os.path.join(LOCK_DIR, locks[0])) as f:
            return yaml.safe_load(f)

    def _extract_ammo(self, job_message):
        if job_message.HasField('test_data'):
            if not job_message.test_data.object_storage_filename:
                LOGGER.warn('Test data specified with no name.')
                return None
            name = job_message.test_data.object_storage_filename
            ammo_file_path = os.path.join(self.workdir, name)
            self.loadtesting_client.download_ammo(job_message.test_data, ammo_file_path)
            LOGGER.debug('Ammo file from test data saved to %s', ammo_file_path)
            return Ammo(name, ammo_file_path)

        elif job_message.HasField('ammo'):
            if not job_message.ammo.name:
                LOGGER.warn('Ammo file specified with no name.')
                return None
            name = job_message.ammo.name
            ammo_file_path = os.path.join(self.workdir, name)
            LOGGER.debug('Ammo file from ammo form saved to %s', ammo_file_path)
            return Ammo.from_content(name, ammo_file_path, job_message.ammo.content)
        return None

    def get_job(self):
        # sending tank status and waiting for job
        tank_status_response = self.tank_api_client.get_tank_status()
        tank_status = self.tank_status(tank_status_response)
        # TODO check code
        if tank_status == TankStatus.READY_FOR_TEST:
            try:
                job_message = self.loadtesting_client.get_job()
            except (LoadTestingInternalError, LoadTestingNotFoundError, LoadTestingConnectionError,
                    LoadTestingFailedPreconditionError) as error:
                LOGGER.error(f'Error with getting job: {error}')
                return

            if job_message is None or not job_message.id:
                return

            try:
                ammo = self._extract_ammo(job_message)
            except (ObjectStorageError, RequestException) as error:
                err_message = f'Error loading test data: {error}'
                LOGGER.error(err_message)
                self.claim_job_failed(err_message)
                return

            return Job(
                id_=job_message.id,
                origin=Job.Origin.LT_SERVER,
                ammo=ammo,
                log_group_id=job_message.logging_log_group_id,
                job_config=json.loads(job_message.config),
            )
        elif tank_status in (TankStatus.TESTING, TankStatus.PREPARING_TEST) and (
                current_job := self.get_current_job()) is not None:
            tank_job_id = current_job.get('test_id')
            job_id = self.jobs_storage.get_cloud_job_id(tank_job_id)
            if job_id:
                origin = Job.Origin.LT_SERVER
                LOGGER.info(f"Cloud job id from storage {job_id}")
            else:
                origin = Job.Origin.LOCAL
                LOGGER.info('Cloud job id is not found in storage.')

            return Job(
                id_=job_id,
                origin=origin,
                tank_job_id=tank_job_id,
                is_running=True,
            )
        return

    def prepare_lt_job(self):
        if self.job.is_running:
            return
        LOGGER.info('going to run job %s', self.job.id)
        # prepare data and job
        tank_config_path = os.path.join(self.workdir, self.config_file_name)
        utils.write_config_file(self.job.job_config, tank_config_path)
        tankapi_response = self.tank_api_client.prepare_job(tank_config_path, self.job.ammo)
        self.claim_and_raise_on_error(tankapi_response, 'Could not run job: {error}')
        self.claim_and_raise_on_unsuccessful(tankapi_response, 'Could not run job.')

        tank_job_id = tankapi_response['id']
        self.job.tank_job_id = tank_job_id

        # FIXME there is some delay between starting the job and adding it to the storage.
        #  Hypothetically, DataUploader can try to create its own job
        self.jobs_storage.push_job(self.job.id, tank_job_id)

        self.job.is_running = True
        LOGGER.info(f'Prepared job {self.job.id=}, {tank_job_id=}')

    def serve_lt_job(self):
        LOGGER.info(f'waiting for job {self.job.id} to finish')
        while True:
            self.serve_lt_signal()
            self.report_lt_status()
            if self.job.status.finished:
                break
            time.sleep(self.sleep_time)
        LOGGER.info(f'The job {self.job.id} is finished. Tank job id {self.job.tank_job_id}')

    def serve_lt_signal(self):
        signal = self.loadtesting_client.get_job_signal(self.job.id)
        match name := signal.Signal.Name(signal.signal):
            case 'STOP':
                self.serve_stop_signal()
            case 'RUN_IN':
                self.serve_run_signal(signal.run_in)
            case 'WAIT':
                pass
            case 'SIGNAL_UNSPECIFIED':
                pass
            case _:
                raise LoadTestingServerError(f'Unknown signal {name} returned from server')

    def claim_job_failed(self, error, error_type=None):
        assert self.job
        return self.loadtesting_client.claim_job_status(
            self.job.id, AdditionalJobStatus.FAILED.name, error, error_type
        )

    def claim_and_raise_on_error(self, tankapi_response, message_template='{error}'):
        error, error_type = TankApiClient.extract_error(tankapi_response)
        if error or error_type:
            message = message_template.format(error=error, error_type=error_type)
            self.claim_job_failed(message, error_type)
            raise TankapiException(f'{message} {error_type=}')

    def claim_and_raise_on_unsuccessful(self, tankapi_response, message):
        if not tankapi_response.get('success'):
            self.claim_job_failed(message)
            raise TankapiException(message)

    def serve_stop_signal(self, stopping_sleep_time=None) -> None:
        self.job.stop_signal_received = True
        stop_response = self.tank_api_client.stop_job()
        self.claim_and_raise_on_error(stop_response, 'Could not stop job: {error}')
        self.claim_and_raise_on_unsuccessful(stop_response, 'Could not stop job')

        sleep_time = stopping_sleep_time if stopping_sleep_time is not None else STOPPING_SLEEP_TIME
        time.sleep(sleep_time)

    def serve_run_signal(self, run_in) -> None:
        if run_in > self.sleep_time:
            return
        if run_in > 0:
            time.sleep(run_in)
        run_response = self.tank_api_client.run_job()
        self.claim_and_raise_on_error(run_response)

    def report_lt_status(self):
        job_status_response = self.tank_api_client.get_job_status(self.job.tank_job_id)
        self.job.update_status(job_status_response)

        code = self.loadtesting_client.claim_job_status(self.job.id,
                                                        self.job.status.name,
                                                        self.job.status.error,
                                                        self.job.status.error_type)
        if code:
            raise LoadTestingServerError(f"Loadtesting server return non-zero code {code} on claim status")

    def serve_local_job(self):
        # sending job status and waiting for job finishing
        job_id = self.job.id
        tank_job_id = self.job.tank_job_id
        test_dir = self.job.test_dir

        LOGGER.info(f"Got a job started manually: {job_id}")
        LOGGER.info(f'Waiting for the job {job_id=}, {tank_job_id=} started manually')

        tank_status = TankStatus.TESTING
        while tank_status == TankStatus.TESTING:
            # claim job status RUNNING (because we know nothing about state of the job)
            try:
                _ = self.loadtesting_client.claim_job_status(job_id, 'RUNNING')
            except (LoadTestingFailedPreconditionError, LoadTestingNotFoundError):
                # do not claim status any more, if loadtesting service knows nothing about current job
                break
            time.sleep(self.sleep_time)

        try:
            with open(os.path.join(test_dir, FINISHED_FILE)) as f:
                job_finish_status = yaml.safe_load(f)
            error, error_type = self.tank_api_client.extract_error(job_finish_status)
            if error or (job_status := job_finish_status.get('status_code')) is None:
                job_status = AdditionalJobStatus.FAILED.name
        except (PermissionError, FileNotFoundError):
            error = 'Could not find file with final status'
            error_type = None
            job_status = AdditionalJobStatus.FAILED.name
        _ = self.loadtesting_client.claim_job_status(job_id, job_status, error, error_type)

        LOGGER.info(f'The job {job_id} is finished. Tank job id {tank_job_id}')

    @staticmethod
    def tank_status(response):
        if response.get("success") is not True:
            return TankStatus.TANK_FAILED
        if not {"is_preparing", "is_testing"}.issubset(response):
            return TankStatus.TANK_FAILED
        if all([not response["is_preparing"], not response["is_testing"], response["success"]]):
            return TankStatus.READY_FOR_TEST
        if response["is_preparing"] is True:
            return TankStatus.PREPARING_TEST
        if response["is_testing"] is True:
            return TankStatus.TESTING
        return TankStatus.TANK_FAILED

    def serve(self):
        with self.reporting_tank_status():
            while True:
                with self.cleanup_after_job():
                    with self.send_report_to_logging():
                        self.wait_for_a_job()
                        match self.job.origin:
                            case Job.Origin.LT_SERVER:
                                self.prepare_lt_job()
                                self.serve_lt_job()
                            case Job.Origin.LOCAL:
                                self.serve_local_job()
                            case _:
                                raise RuntimeError("Improbability Drive generates unreal job "
                                                   f"with origin {self.job.origin}.")

    @contextmanager
    def cleanup_after_job(self):
        assert self.job is None
        try:
            yield
        except Exception as error:
            LOGGER.exception("Unexpected error: %s", error)
            raise error
        finally:
            # TODO:  try to stop job here???
            self.job = None

    def _send_log(self, log_type: LogType):
        send_log(self.job.test_dir, log_type, self.job.log_group_id, self.job.id,
                 self.loadtesting_client.token, self.logging_host, self.logging_port)

    @contextmanager
    def send_report_to_logging(self):
        try:
            yield
        finally:
            if self.job.log_group_id:
                LOGGER.info('Sending logs...')
                self._send_log(LogType.TANK)
                self._send_log(LogType[self.job.generator.name])

    def wait_for_a_job(self):
        while True:
            job = self.get_job()
            if job is not None and job.id is not None:
                LOGGER.info(f'Get job: {job}')
                self.job = job
                return
            time.sleep(self.job_pooling_delay)
