"""Confluent Schema Registry implementation for DataHub."""

import json
import logging
from dataclasses import dataclass
from hashlib import md5
from typing import Dict, List, Optional, Set, Tuple

import avro.schema
import jsonref
from confluent_kafka.schema_registry.schema_registry_client import (
    RegisteredSchema,
    Schema,
    SchemaReference,
    SchemaRegistryClient,
)

from datahub.ingestion.extractor import protobuf_util, schema_util
from datahub.ingestion.extractor.json_schema_util import JsonSchemaTranslator
from datahub.ingestion.extractor.protobuf_util import ProtobufSchema
from datahub.ingestion.source.kafka.kafka import KafkaSourceConfig, KafkaSourceReport
from datahub.ingestion.source.kafka.kafka_schema_registry_base import (
    KafkaSchemaRegistryBase,
    SchemaAndFields,
)
from datahub.metadata.com.linkedin.pegasus2avro.schema import (
    KafkaSchema,
    SchemaField,
    SchemaMetadata,
)
from datahub.metadata.schema_classes import OwnershipSourceTypeClass
from datahub.utilities.mapping import OperationProcessor

logger = logging.getLogger(__name__)


@dataclass
class JsonSchemaWrapper:
    name: str
    subject: str
    content: str
    references: List[Schema]


class ConfluentSchemaRegistry(KafkaSchemaRegistryBase):
    """
    Confluent Schema Registry implementation for DataHub.
    Handles schema retrieval and parsing from Confluent Schema Registry.
    """

    def __init__(
        self, source_config: KafkaSourceConfig, report: KafkaSourceReport
    ) -> None:
        self.source_config: KafkaSourceConfig = source_config
        self.report: KafkaSourceReport = report
        self.schema_registry_client = SchemaRegistryClient(
            {
                "url": source_config.connection.schema_registry_url,
                **source_config.connection.schema_registry_config,
            }
        )
        self.known_schema_registry_subjects: List[str] = []
        try:
            self.known_schema_registry_subjects.extend(
                self.schema_registry_client.get_subjects()
            )
        except OSError as e:
            self.report.schema_registry_connectivity_failures += 1
            self.report.warning(
                title="Failed to list schema registry subjects",
                message="Could not reach the schema registry; every topic may be treated as schemaless.",
                exc=e,
            )
        except Exception as e:
            self.report.warning(
                title="Failed to list schema registry subjects",
                message="Every topic may be treated as schemaless.",
                exc=e,
            )

        self.field_meta_processor = OperationProcessor(
            self.source_config.field_meta_mapping,
            self.source_config.tag_prefix,
            OwnershipSourceTypeClass.SERVICE,
            self.source_config.strip_user_ids_from_email,
            match_nested_props=True,
        )

    @classmethod
    def create(
        cls, source_config: KafkaSourceConfig, report: KafkaSourceReport
    ) -> "ConfluentSchemaRegistry":
        return cls(source_config, report)

    def _get_subject_for_topic(self, topic: str, is_key_schema: bool) -> Optional[str]:
        subject_key_suffix: str = "-key" if is_key_schema else "-value"
        # For details on schema registry subject name strategy,
        # see: https://docs.confluent.io/platform/current/schema-registry/serdes-develop/index.html#how-the-naming-strategies-work

        # User-provided subject for the topic overrides the rest, regardless of the subject name strategy.
        # However, it is a must when the RecordNameStrategy is used as the schema registry subject name strategy.
        # The subject name format for RecordNameStrategy is: <fully-qualified record name>-<key/value> (cannot be inferred from topic name).
        subject_key: str = topic + subject_key_suffix
        if subject_key in self.source_config.topic_subject_map:
            return self.source_config.topic_subject_map[subject_key]

        # Subject name format when the schema registry subject name strategy is
        #  (a) TopicNameStrategy(default strategy): <topic name>-<key/value>
        #  (b) TopicRecordNameStrategy: <topic name>-<fully-qualified record name>-<key/value>
        #  there's a third case
        #  (c) TopicNameStrategy differing by environment name suffixes.
        #       e.g "a.b.c.d-value" and "a.b.c.d.qa-value"
        #       For such instances, the wrong schema registry entries could picked by the previous logic.
        for subject in self.known_schema_registry_subjects:
            if (
                self.source_config.disable_topic_record_naming_strategy
                and subject == subject_key
            ):
                return subject
            if (
                (not self.source_config.disable_topic_record_naming_strategy)
                and subject.startswith(topic)
                and subject.endswith(subject_key_suffix)
            ):
                return subject
        return None

    @staticmethod
    def _compact_schema(schema_str: str) -> str:
        # Eliminate all white-spaces for a compact representation.
        return json.dumps(json.loads(schema_str), separators=(",", ":"))

    def get_schema_str_replace_confluent_ref_avro(
        self, schema: Schema, schema_seen: Optional[Set[str]] = None
    ) -> str:
        if not schema.references:
            return self._compact_schema(schema.schema_str)

        if schema_seen is None:
            schema_seen = set()
        schema_str = self._compact_schema(schema.schema_str)
        for schema_ref in schema.references:
            ref_subject = schema_ref.subject
            if ref_subject in schema_seen:
                continue

            if ref_subject not in self.known_schema_registry_subjects:
                logger.warning(
                    f"{ref_subject} is not present in the list of registered subjects with schema registry!"
                )

            reference_schema = self.schema_registry_client.get_latest_version(
                subject_name=ref_subject,
            )
            schema_seen.add(ref_subject)
            logger.debug(
                f"ref for {ref_subject} is {reference_schema.schema.schema_str}"
            )
            # Replace only external type references with the reference schema recursively.
            # NOTE: The type pattern is dependent on _compact_schema.
            avro_type_kwd = '"type"'
            ref_name = schema_ref.name
            # Try by name first
            pattern_to_replace = f'{avro_type_kwd}:"{ref_name}"'
            if pattern_to_replace not in schema_str:
                # Try by subject
                pattern_to_replace = f'{avro_type_kwd}:"{ref_subject}"'
                if pattern_to_replace not in schema_str:
                    logger.warning(
                        f"Not match for external schema type: {{name:{ref_name}, subject:{ref_subject}}} in schema:{schema_str}"
                    )
                else:
                    logger.debug(
                        f"External schema matches by subject, {pattern_to_replace}"
                    )
            else:
                logger.debug(f"External schema matches by name, {pattern_to_replace}")
            schema_str = schema_str.replace(
                pattern_to_replace,
                f"{avro_type_kwd}:{self.get_schema_str_replace_confluent_ref_avro(reference_schema.schema, schema_seen)}",
            )
        return schema_str

    def get_schemas_from_confluent_ref_protobuf(
        self, schema: Schema, schema_seen: Optional[Set[str]] = None
    ) -> List[ProtobufSchema]:
        all_schemas: List[ProtobufSchema] = []

        if schema_seen is None:
            schema_seen = set()

        schema_ref: SchemaReference
        for schema_ref in schema.references:
            ref_subject: str = schema_ref.subject
            if ref_subject in schema_seen:
                continue
            reference_schema: RegisteredSchema = (
                self.schema_registry_client.get_latest_version(ref_subject)
            )
            schema_seen.add(ref_subject)
            all_schemas.append(
                ProtobufSchema(
                    name=schema_ref.name, content=reference_schema.schema.schema_str
                )
            )
        return all_schemas

    def get_schemas_from_confluent_ref_json(
        self,
        schema: Schema,
        name: str,
        subject: str,
        schema_seen: Optional[Set[str]] = None,
    ) -> List[JsonSchemaWrapper]:
        """Recursively get all the referenced schemas and their references starting from this schema"""
        all_schemas: List[JsonSchemaWrapper] = []
        if schema_seen is None:
            schema_seen = set()

        schema_ref: SchemaReference
        for schema_ref in schema.references:
            ref_subject: str = schema_ref.subject
            if ref_subject in schema_seen:
                continue
            reference_schema: RegisteredSchema = (
                self.schema_registry_client.get_version(
                    subject_name=ref_subject, version=schema_ref.version
                )
            )
            schema_seen.add(ref_subject)
            all_schemas.extend(
                self.get_schemas_from_confluent_ref_json(
                    reference_schema.schema,
                    name=schema_ref.name,
                    subject=ref_subject,
                    schema_seen=schema_seen,
                )
            )
        all_schemas.append(
            JsonSchemaWrapper(
                name=name,
                subject=subject,
                content=schema.schema_str,
                references=schema.references,
            )
        )
        return all_schemas

    def _get_schema_and_fields(
        self, topic: str, is_key_schema: bool, is_subject: bool
    ) -> Tuple[Optional[Schema], List[SchemaField]]:
        schema: Optional[Schema] = None
        kafka_entity = "subject" if is_subject else "topic"

        # if provided schema as topic, assuming it as value subject
        schema_type_str: Optional[str] = "value"
        topic_subject: Optional[str] = None
        if not is_subject:
            schema_type_str = "key" if is_key_schema else "value"
            topic_subject = self._get_subject_for_topic(
                topic=topic, is_key_schema=is_key_schema
            )
        else:
            topic_subject = topic

        if topic_subject is not None:
            logger.debug(
                f"The {schema_type_str} schema subject:'{topic_subject}' is found for {kafka_entity}: '{topic}'."
            )
            try:
                registered_schema = self.schema_registry_client.get_latest_version(
                    subject_name=topic_subject
                )
                schema = registered_schema.schema
            except Exception as e:
                self.report.warning(
                    title="Failed to get subject schema from schema registry",
                    message="Failed to get schema from schema registry",
                    context=f"entity={kafka_entity}, schema_type={schema_type_str or ''}, "
                    f"subject={topic_subject if is_subject else f'{topic}: {topic_subject}'}",
                    exc=e,
                    log=False,
                )
        else:
            logger.debug(
                f"For {kafka_entity}: {topic}, the schema registry subject for the {schema_type_str} schema is not found."
            )
            if not is_key_schema:
                # Value schema is always expected. Check if we should warn.
                if not self.source_config.schema_resolution.enabled:
                    self.report.warning(
                        title="Unable to find a matching subject name for the topic in the schema registry",
                        message="The entity is either schema-less, or no messages have been written yet. "
                        "If this is unexpected, check the topic_subject_map and topic_naming related configs. "
                        "Consider enabling 'schema_resolution.enabled' to automatically infer schema from message data.",
                        context=f"entity={kafka_entity}, schema_type={schema_type_str or ''}, topic={topic}",
                    )

        # Obtain the schema fields from schema for the topic.
        fields: List[SchemaField] = []
        if schema is not None:
            fields = self._get_schema_fields(
                topic=topic,
                schema=schema,
                is_key_schema=is_key_schema,
            )
        return (schema, fields)

    def get_schema_and_fields_batch(
        self, topics: List[str], is_key_schema: bool = False
    ) -> Dict[str, SchemaAndFields]:
        results: Dict[str, SchemaAndFields] = {}
        for topic in topics:
            try:
                schema, fields = self._get_schema_and_fields(
                    topic, is_key_schema, is_subject=False
                )
                results[topic] = SchemaAndFields(schema=schema, fields=fields)
            except OSError as e:
                self.report.schema_registry_connectivity_failures += 1
                self.report.warning(
                    title="Failed to retrieve schema from schema registry",
                    message="Registry connectivity error; topic will be treated as schemaless.",
                    context=topic,
                    exc=e,
                )
                results[topic] = SchemaAndFields()
            except Exception as e:
                self.report.warning(
                    title="Failed to retrieve schema for topic",
                    message="Topic will be treated as schemaless.",
                    context=topic,
                    exc=e,
                )
                results[topic] = SchemaAndFields()
        return results

    def _load_json_schema_with_resolved_references(
        self, schema: Schema, name: str, subject: str
    ) -> dict:
        imported_json_schemas: List[JsonSchemaWrapper] = (
            self.get_schemas_from_confluent_ref_json(schema, name=name, subject=subject)
        )
        schema_dict = json.loads(schema.schema_str)
        reference_map = {}
        for imported_schema in imported_json_schemas:
            reference_schema = json.loads(imported_schema.content)
            if "title" not in reference_schema:
                reference_schema["title"] = imported_schema.subject
            reference_map[imported_schema.name] = reference_schema

        jsonref_schema = jsonref.loads(
            json.dumps(schema_dict), loader=lambda x: reference_map.get(x)
        )
        return jsonref_schema

    def _get_schema_fields(
        self, topic: str, schema: Schema, is_key_schema: bool
    ) -> List[SchemaField]:
        # Parse the schema and convert it to SchemaFields.
        fields: List[SchemaField] = []
        if schema.schema_type == "AVRO":
            cleaned_str: str = self.get_schema_str_replace_confluent_ref_avro(schema)
            avro_schema = avro.schema.parse(
                cleaned_str,
                validate_names=False,
            )

            # "value.id" or "value.[type=string]id"
            fields = schema_util.avro_schema_to_mce_fields(
                avro_schema,
                is_key_schema=is_key_schema,
                meta_mapping_processor=(
                    self.field_meta_processor
                    if self.source_config.enable_meta_mapping
                    else None
                ),
                schema_tags_field=self.source_config.schema_tags_field,
                tag_prefix=self.source_config.tag_prefix,
            )

        elif schema.schema_type == "PROTOBUF":
            try:
                imported_schemas: List[ProtobufSchema] = (
                    self.get_schemas_from_confluent_ref_protobuf(schema)
                )
                base_name: str = topic.replace(".", "_")
                fields = protobuf_util.protobuf_schema_to_mce_fields(
                    ProtobufSchema(
                        (
                            f"{base_name}-key.proto"
                            if is_key_schema
                            else f"{base_name}-value.proto"
                        ),
                        schema.schema_str,
                    ),
                    imported_schemas,
                    is_key_schema=is_key_schema,
                )
                # protobuf_schema_to_mce_fields returns [] (not raises) on compile
                # failures like a duplicate symbol from a shared message type, so
                # surface the empty result here rather than letting it pass silently.
                if not fields:
                    self.report.warning(
                        title="Protobuf schema produced no fields",
                        message="The protobuf schema failed to compile or contained no fields; "
                        "the topic will have no schema fields.",
                        context=topic,
                    )
            except Exception as e:
                self.report.warning(
                    title="Failed to parse protobuf schema",
                    message="The topic will have no schema fields.",
                    context=topic,
                    exc=e,
                )
                fields = []
        elif schema.schema_type == "JSON":
            base_name = topic.replace(".", "_")
            canonical_name = (
                f"{base_name}-key" if is_key_schema else f"{base_name}-value"
            )
            jsonref_schema = self._load_json_schema_with_resolved_references(
                schema=schema,
                name=canonical_name,
                subject=f"{topic}-key" if is_key_schema else f"{topic}-value",
            )
            fields = list(
                JsonSchemaTranslator.get_fields_from_schema(
                    jsonref_schema, is_key_schema=is_key_schema
                )
            )
        elif not self.source_config.ignore_warnings_on_schema_type:
            self.report.warning(
                message="Parsing kafka schema type is currently not implemented",
                context=f"schema_type={schema.schema_type}, topic={topic}",
                log=False,
            )
        return fields

    def get_schema_metadata(
        self, topic: str, platform_urn: str, is_subject: bool
    ) -> Optional[SchemaMetadata]:
        logger.debug(f"Inside get_schema_metadata {topic} {platform_urn}")

        # Process the value schema
        schema, fields = self._get_schema_and_fields(
            topic=topic,
            is_key_schema=False,
            is_subject=is_subject,
        )

        # Process the key schema
        key_schema, key_fields = self._get_schema_and_fields(
            topic=topic,
            is_key_schema=True,
            is_subject=is_subject,
        )

        # Create the schemaMetadata aspect.
        if schema is not None or key_schema is not None:
            # create a merged string for the combined schemas and compute an md5 hash across
            schema_as_string = (schema.schema_str if schema is not None else "") + (
                key_schema.schema_str if key_schema is not None else ""
            )
            md5_hash = md5(schema_as_string.encode()).hexdigest()

            return SchemaMetadata(
                schemaName=topic,
                version=0,
                hash=md5_hash,
                platform=platform_urn,
                platformSchema=KafkaSchema(
                    documentSchema=schema.schema_str if schema else "",
                    documentSchemaType=schema.schema_type if schema else None,
                    keySchema=key_schema.schema_str if key_schema else None,
                    keySchemaType=key_schema.schema_type if key_schema else None,
                ),
                fields=key_fields + fields,
            )
        return None

    def build_schema_metadata_with_key(
        self,
        topic: str,
        platform_urn: str,
        value_schema: Optional[Schema],
        value_fields: List[SchemaField],
        key_schema: Optional[Schema],
        key_fields: List[SchemaField],
    ) -> Optional[SchemaMetadata]:
        """Build SchemaMetadata from pre-fetched value and key schema data."""
        if (
            value_schema is not None
            or key_schema is not None
            or value_fields
            or key_fields
        ):
            # Create a hash from both schemas
            value_schema_str = (
                value_schema.schema_str if value_schema is not None else ""
            )
            key_schema_str = key_schema.schema_str if key_schema is not None else ""
            schema_as_string = value_schema_str + key_schema_str
            md5_hash = md5(schema_as_string.encode()).hexdigest()

            return SchemaMetadata(
                schemaName=topic,
                version=0,
                hash=md5_hash,
                platform=platform_urn,
                platformSchema=KafkaSchema(
                    documentSchema=value_schema.schema_str if value_schema else "",
                    documentSchemaType=value_schema.schema_type
                    if value_schema
                    else None,
                    keySchema=key_schema.schema_str if key_schema else None,
                    keySchemaType=key_schema.schema_type if key_schema else None,
                ),
                fields=key_fields + value_fields,  # Combine key and value fields
            )
        return None

    def get_subjects(self) -> List[str]:
        return self.known_schema_registry_subjects

    def get_schema_registry_client(self) -> SchemaRegistryClient:
        """Get the underlying SchemaRegistryClient for direct access."""
        return self.schema_registry_client
