import ast
import inspect
import logging.config
from functools import cached_property
from pathlib import Path

import grpc
import requests
import yaml
from enum import Enum
from requests_toolbelt.multipart.encoder import MultipartEncoder
# from yandex.cloud.priv.loadtesting.agent.v1 import agent_service_pb2, agent_service_pb2_grpc, job_service_pb2, \
#     job_service_pb2_grpc, agent_registration_service_pb2_grpc, agent_registration_service_pb2
from tankapi.server import APIServer
from yandex.cloud.loadtesting.agent.v1 import agent_service_pb2, agent_service_pb2_grpc, job_service_pb2, \
    job_service_pb2_grpc, agent_registration_service_pb2_grpc, agent_registration_service_pb2
from yandextank.common.interfaces import AbstractCriterion

from load.projects.cloud.cloud_helper import metadata_compute as compute, aws
from load.projects.cloud.cloud_helper.grpc_options import COMMON_CHANNEL_OPTIONS
from load.projects.cloud.cloud_helper.metadata_compute import SaToken
from load.projects.cloud.tank_client.ammo import Ammo
from load.projects.cloud.tank_client.exceptions import LoadTestingConnectionError, JWTError, \
    AgentTypeError
from load.projects.cloud.tank_client.jwt import JWTToken
from load.projects.cloud.tank_client.utils import retry, catch_exceptions, get_creds

LOGGER = logging.getLogger('tank_client')

METADATA_AGENT_VERSION_ATTR = 'agent-version'

EMPTY_FILE_POINTER = None
EMPTY_FILE_TYPE = None

INTERNAL_ERROR_TYPE = 'internal'

AUTOSTOP_EXIT_CODES = [
    value for attr, value in
    inspect.getmembers(AbstractCriterion, lambda a: not (inspect.isroutine(a)))
    if attr.startswith('RC')
]


class AgentType(Enum):
    UNKNOWN = 0
    COMPUTE_LT_CREATED = 1
    COMPUTE_EXTERNAL = 2
    EXTERNAL = 3


class LoadTestingGRPCClient:

    def __init__(self, host, port, agent_version, agent_id_file, object_storage_url, iam_token_service_url,
                 instance_lt_created=True, external_agent_config_path=None, insecure_connection=False):
        self.timeout = 30.0
        self.agent_version = agent_version
        self.agent_id_file = agent_id_file
        self.object_storage_url = object_storage_url
        self.iam_token_service_url = iam_token_service_url
        self.instance_lt_created = instance_lt_created
        self.external_agent_config_path = external_agent_config_path
        self.agent_type = self._identify_agent_type()
        LOGGER.info(f'Agent type {self.agent_type}')
        self._agent_id = None
        self._host = host
        self._port = port
        self._insecure_connection = insecure_connection
        self._identify_agent_id()

    def _identify_agent_type(self):
        if self.instance_lt_created:
            return AgentType.COMPUTE_LT_CREATED
        try:
            SaToken.get()
            return AgentType.COMPUTE_EXTERNAL
        except compute.SaTokenError:
            LOGGER.info("Try to get token with JWT")
        try:
            JWTToken.get(self.iam_token_service_url, self.external_agent_config)
            return AgentType.EXTERNAL
        except JWTError as e:
            LOGGER.error(f"Couldn't get iam token with JWT: {e}")
        raise AgentTypeError("Couldn't identify agent type")

    @property
    def agent_id(self):
        if self._agent_id is None:
            try:
                with open(self.agent_id_file) as f:
                    self._agent_id = f.read() or None  # None if empty
                    LOGGER.info(f'Set agent_id from file {self._agent_id}')
            except FileNotFoundError:
                return
        return self._agent_id

    @agent_id.setter
    def agent_id(self, val):
        self._agent_id = val
        LOGGER.info(f'Set agent_id {self._agent_id}')
        with open(self.agent_id_file, 'w') as f:
            f.write(val)

    @cached_property
    def external_agent_config(self):
        try:
            with open(self.external_agent_config_path) as f:
                return yaml.safe_load(f)
        except (FileNotFoundError, yaml.scanner.ScannerError) as error:
            raise type(error)("Couldn't get agent config") from error

    @property
    @retry(LOGGER, max_delay=30, max_count=10)
    def token(self):
        if self.agent_type in (AgentType.COMPUTE_LT_CREATED, AgentType.COMPUTE_EXTERNAL):
            return SaToken.get()
        if self.agent_type == AgentType.EXTERNAL:
            return JWTToken.get(self.iam_token_service_url, self.external_agent_config)
        raise AgentTypeError('Unknown agent type')

    def _request_metadata(self, additional_meta=None):
        meta = [
            ('authorization', f'Bearer {self.token}'),
            (METADATA_AGENT_VERSION_ATTR, self.agent_version),
        ]
        if additional_meta:
            meta.extend(additional_meta)
        return meta

    @cached_property
    def channel(self):
        host = self._host
        port = self._port
        if ':' in host and not host.startswith('['):
            target = f'[{host}]:{port}'
        else:
            target = f'{host}:{port}'
        LOGGER.info('Connect to %s', target)
        if self._insecure_connection:
            return grpc.insecure_channel(target,
                                         options=COMMON_CHANNEL_OPTIONS + (('grpc.enable_http_proxy', 0),))
        else:
            creds = get_creds(host, port)
            return grpc.secure_channel(target, creds, options=COMMON_CHANNEL_OPTIONS)

    @cached_property
    def stub_agent(self):
        return agent_service_pb2_grpc.AgentServiceStub(self.channel)

    @cached_property
    def stub_job(self):
        return job_service_pb2_grpc.JobServiceStub(self.channel)

    @catch_exceptions(logger=LOGGER)
    def _identify_agent_id(self):
        if self.agent_id:
            LOGGER.info(f'The agent has been registered with id={self.agent_id}')
            return

        stub_register = agent_registration_service_pb2_grpc.AgentRegistrationServiceStub(self.channel)
        if self.agent_type == AgentType.COMPUTE_LT_CREATED:
            response = stub_register.Register(
                agent_registration_service_pb2.RegisterRequest(
                    compute_instance_id=compute.get_current_instance_id()),
                timeout=self.timeout,
                metadata=self._request_metadata()
            )
            self.agent_id = response.agent_instance_id
            return
        else:
            if self.agent_type == AgentType.COMPUTE_EXTERNAL:
                args = dict(compute_instance_id=compute.get_current_instance_id())
            elif self.agent_type == AgentType.EXTERNAL:
                args = dict(name=self.external_agent_config.get('name'),
                            folder_id=self.external_agent_config.get('folder_id'), )
            else:
                raise AgentTypeError('Unknown agent type')
            response = stub_register.ExternalAgentRegister(
                agent_registration_service_pb2.ExternalAgentRegisterRequest(
                    **args
                ),
                timeout=self.timeout,
                metadata=self._request_metadata(),
            )
            metadata = agent_registration_service_pb2.ExternalAgentRegisterMetadata()
            response.metadata.Unpack(metadata)
            self.agent_id = metadata.agent_instance_id
            return

        if not self.agent_id:
            LOGGER.error("Couldn't connect to Cloud Loadtesting service")
            raise LoadTestingConnectionError("Couldn't connect to Cloud Loadtesting service")

    def claim_tank_status(self, tank_status):
        # TODO return status and error message
        request = agent_service_pb2.ClaimAgentStatusRequest(
            agent_instance_id=self.agent_id,
            status=tank_status.name
        )
        result = self.stub_agent.ClaimStatus(
            request,
            timeout=self.timeout,
            metadata=self._request_metadata()
        )
        return result.code

    @retry(logger=LOGGER, max_delay=20, max_count=5)
    @catch_exceptions(logger=LOGGER, return_data=0)
    def claim_job_status(self, job_id, job_status, error='', error_type=None):
        request = job_service_pb2.ClaimJobStatusRequest(
            job_id=job_id,
            status=job_status,
            error=error
        )
        metadata = []
        if error_type is not None:
            metadata.append(('error-type', error_type))
        result = self.stub_job.ClaimStatus(
            request,
            timeout=self.timeout,
            metadata=self._request_metadata(metadata)
        )
        return result.code

    @catch_exceptions(logger=LOGGER, return_data=None)
    def get_job(self):
        request = job_service_pb2.GetJobRequest(
            agent_instance_id=self.agent_id
        )
        try:
            job = self.stub_job.Get(
                request,
                timeout=self.timeout,
                metadata=self._request_metadata(),
            )
            LOGGER.debug(f'Get job: {job}')
            return job
        except grpc.RpcError as error:
            if error.code() == grpc.StatusCode.NOT_FOUND:
                LOGGER.info("There is no job for the agent")
                return

    @retry(logger=LOGGER, max_delay=20, max_count=5)
    @catch_exceptions(logger=LOGGER, return_data=None)
    def get_job_signal(self, job_id):
        LOGGER.debug('Requesting job signal')
        request = job_service_pb2.JobSignalRequest(
            job_id=job_id
        )
        result = self.stub_job.GetSignal(
            request,
            timeout=self.timeout,
            metadata=self._request_metadata(),
        )
        LOGGER.info(f'Got signal {result.signal}')
        return result

    @retry(logger=LOGGER, max_delay=5, max_count=3)
    def download_ammo(self, test_data, path_to_download):
        aws.get_file(test_data.object_storage_bucket, test_data.object_storage_filename, self.token,
                     self.object_storage_url, path_to_download)


class TankApiClient:
    class _Path:
        prepare_job = 'api/v1/tests/prepare.json'
        run_job = 'api/v1/tests/run.json'
        tank_status = 'api/v1/tank/status.json'
        job_status = 'api/v1/tests/{job_id}/status.json'
        stop_job = 'api/v1/tests/stop'

    def __init__(self, host, port):
        self._origin = f'http://{host}:{port}'

    def _url(self, path):
        return f'{self._origin}/{path}'

    @catch_exceptions(logger=LOGGER, return_data={})
    # TODO here needed retry
    def prepare_job(self, conf_file_path, ammo: Ammo):
        with open(conf_file_path, 'rb') as conf_file:
            config_file_name = Path(conf_file_path).name
            multipart_fields = [(APIServer.CONFIG_FIELD,
                                 (config_file_name, conf_file, 'application/octet-stream'))]
            if ammo is not None:
                multipart_fields.append((APIServer.LOCAL_FILE,
                                         (ammo.abs_path, EMPTY_FILE_POINTER, EMPTY_FILE_TYPE, {'method': 'move'})))

            me = MultipartEncoder(fields=multipart_fields)
            response = requests.post(self._url(self._Path.prepare_job), data=me,
                                     headers={'Content-Type': me.content_type})

        json_response = response.json()
        LOGGER.info(f'Test has been prepared. Response: {json_response}')
        return json_response

    @catch_exceptions(logger=LOGGER, return_data={})
    def run_job(self):
        resp = requests.post(self._url(self._Path.run_job))
        resp.raise_for_status()
        status = resp.json()
        LOGGER.info(f'Test is running. Status: {status}')
        return status

    @catch_exceptions(logger=LOGGER, return_data={})
    def get_tank_status(self):
        response = requests.get(self._url(self._Path.tank_status))
        json_response = response.json()
        LOGGER.debug(f'Tank status: {json_response}')
        return json_response

    @catch_exceptions(logger=LOGGER, return_data={})
    def get_job_status(self, job_id):
        response = requests.get(self._url(self._Path.job_status).format(job_id=job_id))
        json_response = response.json()
        LOGGER.debug(f'Job status: {json_response}')
        return json_response

    @catch_exceptions(logger=LOGGER, return_data={})
    def stop_job(self):
        response = requests.get(self._url(self._Path.stop_job))
        json_response = ast.literal_eval(response.text)
        LOGGER.info(f'Job was stopped {json_response}')
        return json_response

    @staticmethod
    def extract_error(response):
        error = response.get('error', '')
        error_type = None
        exit_code = response.get('exit_code')
        if not error:
            error = response.get('tank_msg', '')
            if error:
                error_type = INTERNAL_ERROR_TYPE
            elif exit_code and exit_code not in AUTOSTOP_EXIT_CODES:
                error = 'Unknown generator error'
        return error, error_type


class TankapiException(Exception):
    pass
