# Copyright European Organization for Nuclear Research (CERN) since 2012
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""
Core tracer module
"""

import ipaddress
import json
import logging.handlers
import random
import socket
from typing import TYPE_CHECKING, Any, Union, overload

import stomp
from jsonschema import Draft7Validator, ValidationError, validate

from rucio.common.config import config_get, config_get_int, config_get_list
from rucio.common.exception import InvalidObject, TraceValidationSchemaNotFound
from rucio.common.logging import rucio_log_formatter
from rucio.common.schema.generic import TIME_ENTRY, UUID, IPv4orIPv6
from rucio.core.monitor import MetricManager

if TYPE_CHECKING:
    from datetime import datetime

    from referencing.jsonschema import ObjectSchema, Schema

    from rucio.common.types import TraceSchemaDict

METRICS = MetricManager(module=__name__)

CONFIG_COMMON_LOGLEVEL = getattr(logging, config_get('common', 'loglevel', raise_exception=False, default='DEBUG').upper())

CONFIG_TRACE_LOGLEVEL = getattr(logging, config_get('trace', 'loglevel', raise_exception=False, default='DEBUG').upper())
CONFIG_TRACE_LOGFORMAT = config_get('trace', 'logformat', raise_exception=False, default='%(message)s')
CONFIG_TRACE_TRACEDIR = config_get('trace', 'tracedir', raise_exception=False, default='/var/log/rucio/trace')
CONFIG_TRACE_MAXBYTES = config_get_int('trace', 'maxbytes', raise_exception=False, default=1000000000)
CONFIG_TRACE_BACKUPCOUNT = config_get_int('trace', 'backupCount', raise_exception=False, default=10)

# reset root logger handlers. Otherwise everything from ROTATING_LOGGER will also end up in the apache logs.
logging.getLogger().handlers = []

LOGGER = logging.getLogger('trace')
LOGGER.setLevel(CONFIG_COMMON_LOGLEVEL)

ROTATING_LOGGER = logging.getLogger('trace_buffer')
ROTATING_LOGGER.setLevel(CONFIG_TRACE_LOGLEVEL)

HANDLER = logging.StreamHandler()
FORMATTER = rucio_log_formatter()
HANDLER.setFormatter(FORMATTER)
LOGGER.addHandler(HANDLER)

ROTATING_HANDLER = logging.handlers.RotatingFileHandler(filename='%s/trace' % CONFIG_TRACE_TRACEDIR, maxBytes=CONFIG_TRACE_MAXBYTES, backupCount=CONFIG_TRACE_BACKUPCOUNT)
ROTATING_LOGFORMATTER = logging.Formatter(CONFIG_TRACE_LOGFORMAT)
ROTATING_HANDLER.setFormatter(ROTATING_LOGFORMATTER)
ROTATING_LOGGER.addHandler(ROTATING_HANDLER)

BROKERS_ALIAS, BROKERS_RESOLVED = [], []
try:
    BROKERS_ALIAS = config_get_list('trace', 'brokers')
except:
    raise Exception('Could not load brokers from configuration')

PORT = config_get_int('trace', 'port')
TOPIC = config_get('trace', 'topic')
USERNAME = config_get('trace', 'username')
PASSWORD = config_get('trace', 'password')
VHOST = config_get('trace', 'broker_virtual_host', raise_exception=False)

TOUCH_SCHEMA: 'ObjectSchema' = {
    "description": "touch one or more DIDs",
    "type": "object",
    "properties": {
        "eventType": {"enum": ["touch"]},
        "clientState": {"type": "string"},
        "account": {"type": "string"},
        "scope": {"type": "string"},
        "filename": {"type": "string"},
        "datasetScope": {"type": ["string", "null"]},
        "dataset": {"type": ["string", "null"]},
        "traceTimeentry": TIME_ENTRY,
        "traceTimeentryUnix": {"type": "number"},
        "traceIp": IPv4orIPv6,
        "traceId": UUID,
        "localSite": {"type": "string"},
        "remoteSite": {"type": "string"},
        "usrdn": {"type": "string"},
    },
    "required": ['eventType', 'clientState', 'account', 'traceTimeentry', 'traceTimeentryUnix', 'traceIp', 'traceId']
}

UPLOAD_SCHEMA: 'ObjectSchema' = {
    "description": "upload method",
    "type": "object",
    "properties": {
        "eventType": {"enum": ["upload"]},
        "hostname": {"type": "string"},
        "eventVersion": {"type": "string"},
        "clientState": {"type": "string"},
        "account": {"type": "string"},
        "uuid": UUID,
        "scope": {"type": "string"},
        "datasetScope": {"type": ["string", "null"]},
        "dataset": {"type": ["string", "null"]},
        "remoteSite": {"type": "string"},
        "filesize": {"type": "number"},
        "protocol": {"type": "string"},
        "transferStart": {"type": "number"},
        "transferEnd": {"type": "number"},
        "traceTimeentry": TIME_ENTRY,
        "traceTimeentryUnix": {"type": "number"},
        "traceIp": IPv4orIPv6,
        "traceId": UUID,
        "vo": {"type": "string"},
        "stateReason": {"type": "string"},
        "filename": {"type": "string"},
        "name": {"type": "string"},
        "usrdn": {"type": "string"},
    },
    "required": ['hostname', 'account', 'eventType', 'eventVersion', 'uuid', 'scope', 'dataset',
                 'remoteSite', 'filesize', 'protocol', 'transferStart', 'traceTimeentry', 'traceTimeentryUnix',
                 'traceIp', 'traceId']
}

DOWNLOAD_SCHEMA: 'ObjectSchema' = {
    "description": "download method",
    "type": "object",
    "properties": {
        "eventType": {"enum": ["download"]},
        "hostname": {"type": "string"},
        "eventVersion": {"type": "string"},
        "localSite": {"type": "string"},
        "remoteSite": {"type": "string"},
        "account": {"type": "string"},
        "uuid": UUID,
        "scope": {"type": "string"},
        "filename": {"type": "string"},
        "datasetScope": {"type": ["string", "null"]},
        "dataset": {"type": ["string", "null"]},
        "filesize": {"type": ["number", "null"]},
        "clientState": {"type": "string"},
        "stateReason": {"type": "string"},
        "protocol": {"type": "string"},
        "transferStart": {"type": "number"},
        "transferEnd": {"type": "number"},
        "traceTimeentry": TIME_ENTRY,
        "traceTimeentryUnix": {"type": "number"},
        "traceIp": IPv4orIPv6,
        "traceId": UUID,
        "vo": {"type": "string"},
        "usrdn": {"type": "string"},
        "name": {"type": "string"},
    },
    "required": ['hostname', 'eventType', 'localSite', 'account', 'eventVersion', 'uuid', 'scope',
                 'filename', 'datasetScope', 'dataset', 'filesize', 'clientState', 'stateReason']
}

GET_SCHEMA: 'ObjectSchema' = {
    "description": "get method, mainly sent by pilots",
    "type": "object",
    "properties": {
        "eventType": {"enum": ["get", "get_sm", "sm_get", "get_sm_a", "sm_get_a"]},
        "clientState": {"type": "string"},
        "stateReason": {"type": "string"},
        "url": {"type": ["string", "null"]},
        "vo": {"type": "string"},
        "scope": {"type": "string"},
        "eventVersion": {"type": "string"},
        "remoteSite": {"type": "string"},
        "datasetScope": {"type": "string"},
        "dataset": {"type": "string"},
        "filename": {"type": "string"},
        "name": {"type": "string"},
        "traceTimeentry": TIME_ENTRY,
        "traceTimeentryUnix": {"type": "number"},
        "traceIp": IPv4orIPv6,
        "traceId": UUID,
        "usrdn": {"type": "string"},
    },
    "required": ['eventType', 'localSite', 'eventVersion', 'uuid', 'scope',
                 'filename', 'dataset']
}

PUT_SCHEMA: 'ObjectSchema' = {
    "description": "get method, mainly sent by pilots",
    "type": "object",
    "properties": {
        "eventType": {"enum": ["put_sm", "put_sm_a"]},
        "clientState": {"type": "string"},
        "stateReason": {"type": "string"},
        "url": {"type": ["string", "null"]},
        "vo": {"type": "string"},
        "scope": {"type": "string"},
        "eventVersion": {"type": "string"},
        "remoteSite": {"type": "string"},
        "datasetScope": {"type": "string"},
        "dataset": {"type": "string"},
        "filename": {"type": "string"},
        "name": {"type": "string"},
        "traceTimeentry": TIME_ENTRY,
        "traceTimeentryUnix": {"type": "number"},
        "traceIp": IPv4orIPv6,
        "traceId": UUID,
        "usrdn": {"type": "string"},
        "pq": {"type": "string"},
        "localSite": {"type": "string"}
    },
    "required": ['eventType', 'localSite', 'eventVersion', 'uuid',
                 'filename', 'dataset']
}

SPECIAL_SCHEMA: 'ObjectSchema' = {
    "description": "A special schema to capture most unsupported eventTypes",
    "type": "object",
    "properties": {
        "eventType": {"enum": ["sfo2eos"]},
        "clientState": {"type": "string"},
        "account": {"type": "string"},
        "scope": {"type": "string"},
        "filename": {"type": "string"},
        "datasetScope": {"type": ["string", "null"]},
        "dataset": {"type": ["string", "null"]},
        "traceTimeentry": TIME_ENTRY,
        "traceTimeentryUnix": {"type": "number"},
        "traceIp": IPv4orIPv6,
        "traceId": UUID,
        "localSite": {"type": "string"},
        "remoteSite": {"type": "string"},
        "usrdn": {"type": "string"},
    },
    "required": ['eventType', 'clientState', 'account', 'traceTimeentry', 'traceTimeentryUnix', 'traceIp', 'traceId']
}

SCHEMAS: dict[str, 'Schema'] = {
    'touch': TOUCH_SCHEMA,
    'upload': UPLOAD_SCHEMA,
    'download': DOWNLOAD_SCHEMA,
    'get': GET_SCHEMA,
    'get_sm': GET_SCHEMA,
    'sm_get': GET_SCHEMA,
    'get_sm_a': GET_SCHEMA,
    'sm_get_a': GET_SCHEMA,
    'put': PUT_SCHEMA,
    'put_sm': PUT_SCHEMA,
    'put_sm_a': PUT_SCHEMA,
    'sm_put': PUT_SCHEMA,
    'sm_put_a': PUT_SCHEMA,
    'sfo2eos': SPECIAL_SCHEMA
}

FORMAT_CHECKER = Draft7Validator.FORMAT_CHECKER


@FORMAT_CHECKER.checks(format="ipv4_or_ipv6")
def ip_format_checker(value: str) -> bool:
    """
    Validates IPv4 or IPv6 string values. json schemas can use `ipv4_or_ipv6` as a valid `format` argument
    """
    try:
        ipaddress.ip_address(value)
    except ValueError:
        LOGGER.debug(f"{value} is not a valid IPv4 or IPv6 address and raises an errors upon validation.")
        result = False
    else:
        result = True
    return result


logging.getLogger("stomp").setLevel(logging.CRITICAL)

for broker in BROKERS_ALIAS:
    try:
        addrinfos = socket.getaddrinfo(broker, 0, socket.AF_INET, 0, socket.IPPROTO_TCP)
        BROKERS_RESOLVED = [ai[4][0] for ai in addrinfos]
    except:
        pass

CONNS = []

for broker in BROKERS_RESOLVED:
    CONNS.append(stomp.Connection(host_and_ports=[(broker, PORT)], vhost=VHOST, reconnect_attempts_max=3))


@overload
def date_handler(obj: "datetime") -> str:
    ...


@overload
def date_handler(obj: object) -> object:
    ...


def date_handler(obj: Any) -> Union[str, object]:
    """ format dates to ISO format """
    return obj.isoformat() if hasattr(obj, 'isoformat') else obj


@METRICS.count_it
def trace(payload: dict[str, Any]) -> None:
    """
    Write a trace to the buffer log file and send it to active mq.

    :param payload: Python dictionary with trace report.
    """

    report = json.dumps(payload, default=date_handler)
    ROTATING_LOGGER.debug(report)
    t_conns = CONNS[:]

    try:
        validate_schema(report)
    except InvalidObject as error:
        ROTATING_LOGGER.warning("Problem validating schema: %s" % error)
        LOGGER.warning("Problem validating schema: %s" % error)

    try:
        for i in range(len(t_conns)):
            try:
                conn = random.sample(t_conns, 1)[0]
                if not conn.is_connected():
                    LOGGER.info('reconnect to ' + conn.transport._Transport__host_and_ports[0][0])
                    conn.connect(USERNAME, PASSWORD)
            except stomp.exception.NotConnectedException:
                LOGGER.warning('Could not connect to broker %s, try another one' %
                               conn.transport._Transport__host_and_ports[0][0])
                t_conns.remove(conn)
                continue
            except stomp.exception.ConnectFailedException:
                LOGGER.warning('Could not connect to broker %s, try another one' %
                               conn.transport._Transport__host_and_ports[0][0])
                t_conns.remove(conn)
                continue

        if conn.is_connected:
            conn.send(body=report, destination=TOPIC, headers={'persistent': 'true', 'appversion': 'rucio'})
        else:
            LOGGER.error("Unable to connect to broker. Could not send trace: %s" % report)
    except Exception as error:
        LOGGER.error(error)


def validate_schema(obj: str) -> None:
    """
    Validate object against json schema

    :param obj: The object to validate.

    :raises: InvalidObject, TraceValidationSchemaNotFound
    """
    loaded_obj: 'TraceSchemaDict' = json.loads(obj)

    if loaded_obj and 'eventType' in loaded_obj:
        event_type = loaded_obj['eventType'].lower()
        schema = SCHEMAS.get(event_type)
        if schema is not None:
            try:
                validate(loaded_obj, schema, format_checker=FORMAT_CHECKER)
            except ValidationError as error:
                raise InvalidObject(error)
        else:
            raise TraceValidationSchemaNotFound("Trace schema for eventType %s not found. This event type might not be supported." % event_type)
