#!/usr/bin/env python
# -*- coding: utf-8 -*-
#
# Copyright 2024 Confluent Inc.
#
# 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.
#

import logging
from typing import List, Optional, Set, Dict, Any

from confluent_kafka.schema_registry import RegisteredSchema
from confluent_kafka.schema_registry.common.schema_registry_client import \
    RulePhase
from confluent_kafka.schema_registry.common.serde import ErrorAction, \
    FieldTransformer, Migration, NoneAction, RuleAction, \
    RuleConditionError, RuleContext, RuleError, SchemaId
from confluent_kafka.schema_registry.schema_registry_client import RuleMode, \
    Rule, RuleKind, Schema, RuleSet
from confluent_kafka.serialization import Serializer, Deserializer, \
    SerializationContext, SerializationError

__all__ = [
    'BaseSerde',
    'BaseSerializer',
    'BaseDeserializer',
]

log = logging.getLogger(__name__)


class BaseSerde(object):
    __slots__ = ['_use_schema_id', '_use_latest_version', '_use_latest_with_metadata',
                 '_registry', '_rule_registry', '_subject_name_func',
                 '_field_transformer']

    def _get_reader_schema(self, subject: str, fmt: Optional[str] = None) -> Optional[RegisteredSchema]:
        if self._use_schema_id is not None:
            schema = self._registry.get_schema(self._use_schema_id, subject, fmt)
            return self._registry.lookup_schema(
                subject, schema, normalize_schemas=False, deleted=True)
        if self._use_latest_with_metadata is not None:
            return self._registry.get_latest_with_metadata(
                subject, self._use_latest_with_metadata, deleted=True, fmt=fmt)
        if self._use_latest_version:
            return self._registry.get_latest_version(subject, fmt)
        return None

    def _execute_rules(
        self, ser_ctx: SerializationContext, subject: str,
        rule_mode: RuleMode,
        source: Optional[Schema], target: Optional[Schema],
        message: Any, inline_tags: Optional[Dict[str, Set[str]]],
        field_transformer: Optional[FieldTransformer]
    ) -> Any:
        return self._execute_rules_with_phase(
            ser_ctx, subject, RulePhase.DOMAIN, rule_mode,
            source, target, message, inline_tags, field_transformer)

    def _execute_rules_with_phase(
        self, ser_ctx: SerializationContext, subject: str,
        rule_phase: RulePhase, rule_mode: RuleMode,
        source: Optional[Schema], target: Optional[Schema],
        message: Any, inline_tags: Optional[Dict[str, Set[str]]],
        field_transformer: Optional[FieldTransformer]
    ) -> Any:
        if message is None or target is None:
            return message
        rules: Optional[List[Rule]] = None
        if rule_mode == RuleMode.UPGRADE:
            if target is not None and target.rule_set is not None:
                rules = target.rule_set.migration_rules
        elif rule_mode == RuleMode.DOWNGRADE:
            if source is not None and source.rule_set is not None:
                rules = source.rule_set.migration_rules
                rules = rules[:] if rules else []
                rules.reverse()
        else:
            if target is not None and target.rule_set is not None:
                if rule_phase == RulePhase.ENCODING:
                    rules = target.rule_set.encoding_rules
                else:
                    rules = target.rule_set.domain_rules
                if rule_mode == RuleMode.READ:
                    # Execute read rules in reverse order for symmetry
                    rules = rules[:] if rules else []
                    rules.reverse()

        if not rules:
            return message

        for index in range(len(rules)):
            rule = rules[index]
            if self._is_disabled(rule):
                continue
            if rule.mode == RuleMode.WRITEREAD:
                if rule_mode != RuleMode.READ and rule_mode != RuleMode.WRITE:
                    continue
            elif rule.mode == RuleMode.UPDOWN:
                if rule_mode != RuleMode.UPGRADE and rule_mode != RuleMode.DOWNGRADE:
                    continue
            elif rule.mode != rule_mode:
                continue

            ctx = RuleContext(ser_ctx, source, target, subject, rule_mode, rule,
                              index, rules, inline_tags, field_transformer)
            rule_executor = self._rule_registry.get_executor(rule.type.upper())
            if rule_executor is None:
                self._run_action(ctx, rule_mode, rule, self._get_on_failure(rule), message,
                                 RuleError(f"Could not find rule executor of type {rule.type}"),
                                 'ERROR')
                return message
            try:
                result = rule_executor.transform(ctx, message)
                if rule.kind == RuleKind.CONDITION:
                    if not result:
                        raise RuleConditionError(rule)
                elif rule.kind == RuleKind.TRANSFORM:
                    message = result
                self._run_action(
                    ctx, rule_mode, rule,
                    self._get_on_failure(rule) if message is None else self._get_on_success(rule),
                    message, None,
                    'ERROR' if message is None else 'NONE')
            except SerializationError:
                raise
            except Exception as e:
                self._run_action(ctx, rule_mode, rule, self._get_on_failure(rule),
                                 message, e, 'ERROR')
        return message

    def _get_on_success(self, rule: Rule) -> Optional[str]:
        override = self._rule_registry.get_override(rule.type)
        if override is not None and override.on_success is not None:
            return override.on_success
        return rule.on_success

    def _get_on_failure(self, rule: Rule) -> Optional[str]:
        override = self._rule_registry.get_override(rule.type)
        if override is not None and override.on_failure is not None:
            return override.on_failure
        return rule.on_failure

    def _is_disabled(self, rule: Rule) -> Optional[bool]:
        override = self._rule_registry.get_override(rule.type)
        if override is not None and override.disabled is not None:
            return override.disabled
        return rule.disabled

    def _run_action(
        self, ctx: RuleContext, rule_mode: RuleMode, rule: Rule,
        action: Optional[str], message: Any,
        ex: Optional[Exception], default_action: str
    ):
        action_name = self._get_rule_action_name(rule, rule_mode, action)
        if action_name is None:
            action_name = default_action
        rule_action = self._get_rule_action(ctx, action_name)
        if rule_action is None:
            log.error("Could not find rule action of type %s", action_name)
            raise RuleError(f"Could not find rule action of type {action_name}")
        try:
            rule_action.run(ctx, message, ex)
        except SerializationError:
            raise
        except Exception as e:
            log.warning("Could not run post-rule action %s: %s", action_name, e)

    def _get_rule_action_name(
        self, rule: Rule, rule_mode: RuleMode, action_name: Optional[str]
    ) -> Optional[str]:
        if action_name is None or action_name == "":
            return None
        if rule.mode in (RuleMode.WRITEREAD, RuleMode.UPDOWN) and ',' in action_name:
            parts = action_name.split(',')
            if rule_mode in (RuleMode.WRITE, RuleMode.UPGRADE):
                return parts[0]
            elif rule_mode in (RuleMode.READ, RuleMode.DOWNGRADE):
                return parts[1]
        return action_name

    def _get_rule_action(self, ctx: RuleContext, action_name: str) -> Optional[RuleAction]:
        if action_name == 'ERROR':
            return ErrorAction()
        elif action_name == 'NONE':
            return NoneAction()
        return self._rule_registry.get_action(action_name)


class BaseSerializer(BaseSerde, Serializer):
    __slots__ = ['_auto_register', '_normalize_schemas', '_schema_id_serializer']


class BaseDeserializer(BaseSerde, Deserializer):
    __slots__ = ['_schema_id_deserializer']

    def _get_writer_schema(
            self, schema_id: SchemaId, subject: Optional[str] = None,
            fmt: Optional[str] = None) -> Schema:
        if schema_id.id is not None:
            return self._registry.get_schema(schema_id.id, subject, fmt)
        elif schema_id.guid is not None:
            return self._registry.get_schema_by_guid(str(schema_id.guid), fmt)
        else:
            raise SerializationError("Schema ID or GUID is not set")

    def _has_rules(self, rule_set: RuleSet, phase: RulePhase, mode: RuleMode) -> bool:
        if rule_set is None:
            return False
        if phase == RulePhase.MIGRATION:
            rules = rule_set.migration_rules
        elif phase == RulePhase.DOMAIN:
            rules = rule_set.domain_rules
        elif phase == RulePhase.ENCODING:
            rules = rule_set.encoding_rules
        if mode in (RuleMode.UPGRADE, RuleMode.DOWNGRADE):
            return any(rule.mode == mode or rule.mode == RuleMode.UPDOWN
                       for rule in rules or [])
        elif mode == RuleMode.UPDOWN:
            return any(rule.mode == mode for rule in rules or [])
        elif mode in (RuleMode.WRITE, RuleMode.READ):
            return any(rule.mode == mode or rule.mode == RuleMode.WRITEREAD
                       for rule in rules or [])
        elif mode == RuleMode.WRITEREAD:
            return any(rule.mode == mode for rule in rules or [])
        return False

    def _get_migrations(
        self, subject: str, source_info: Schema,
        target: RegisteredSchema, fmt: Optional[str]
    ) -> List[Migration]:
        source = self._registry.lookup_schema(
            subject, source_info, normalize_schemas=False, deleted=True)
        migrations = []
        if source.version < target.version:
            migration_mode = RuleMode.UPGRADE
            first = source
            last = target
        elif source.version > target.version:
            migration_mode = RuleMode.DOWNGRADE
            first = target
            last = source
        else:
            return migrations
        previous: Optional[RegisteredSchema] = None
        versions = self._get_schemas_between(subject, first, last, fmt)
        for i in range(len(versions)):
            version = versions[i]
            if i == 0:
                previous = version
                continue
            if version.schema.rule_set is not None and self._has_rules(
                  version.schema.rule_set, RulePhase.MIGRATION, migration_mode):
                if migration_mode == RuleMode.UPGRADE:
                    migration = Migration(migration_mode, previous, version)
                else:
                    migration = Migration(migration_mode, version, previous)
                migrations.append(migration)
            previous = version
        if migration_mode == RuleMode.DOWNGRADE:
            migrations.reverse()
        return migrations

    def _get_schemas_between(
        self, subject: str, first: RegisteredSchema,
        last: RegisteredSchema, fmt: Optional[str] = None
    ) -> List[RegisteredSchema]:
        if last.version - first.version <= 1:
            return [first, last]
        version1 = first.version
        version2 = last.version
        result = [first]
        for i in range(version1 + 1, version2):
            result.append(self._registry.get_version(subject, i, True, fmt))
        result.append(last)
        return result

    def _execute_migrations(
        self, ser_ctx: SerializationContext, subject: str,
        migrations: List[Migration], message: Any
    ) -> Any:
        for migration in migrations:
            message = self._execute_rules_with_phase(
                ser_ctx, subject, RulePhase.MIGRATION, migration.rule_mode,
                migration.source.schema, migration.target.schema,
                message, None, None)
        return message