# This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0. If a copy of the MPL was not distributed with this
# file, You can obtain one at https://mozilla.org/MPL/2.0/.

"""
Streaming provider wrapper classes for Featureform.
"""

from typing import TYPE_CHECKING, List, Optional, Union

from ..resources import KafkaTopic
from ..resources.constants import STREAM_CHANNEL_VARIANT
from ..utils.helpers import set_tags_properties

if TYPE_CHECKING:
    from ..registrar import UserRegistrar

from .offline import OfflineProvider

__all__ = ["KafkaProvider"]


class KafkaProvider(OfflineProvider):
    def __init__(self, registrar, provider):
        super().__init__(registrar, provider)
        self.__registrar = registrar
        self.__provider = provider

    def register_kafka_topic(
        self,
        name: str,
        topic: str,
        timestamp_expression: str = "",
        watermark_delay: str = "5 minutes",
        description: str = "",
        tags: Optional[List[str]] = None,
        properties: Optional[dict] = None,
        owner: Union[str, "UserRegistrar"] = "",
    ):
        """Register a Kafka topic as a primary data source.

        The `timestamp_expression` is a Spark SQL expression that extracts the event
        timestamp from the Kafka message. The expression has access to all Kafka columns:
        - `key` (binary) - Message key
        - `value` (binary) - Message payload (typically JSON)
        - `topic` (string) - Kafka topic name
        - `partition` (int) - Partition number
        - `offset` (long) - Message offset
        - `timestamp` (timestamp) - Kafka ingestion timestamp
        - `timestampType` (int) - Timestamp type

        **Examples**

        ```python
        spark = client.get_provider("my_spark")

        # Use Kafka's ingestion timestamp
        transactions_stream = spark.register_kafka_topic(
            name="transactions_stream",
            topic="transactions",
            timestamp_expression="timestamp",
            watermark_delay="5 minutes",
        )

        # Extract timestamp from JSON payload
        transactions_stream = spark.register_kafka_topic(
            name="transactions_stream",
            topic="transactions",
            timestamp_expression="CAST(get_json_object(CAST(value AS STRING), '$.event_time') AS TIMESTAMP)",
            watermark_delay="5 minutes",
        )

        # Convert epoch milliseconds from JSON payload
        transactions_stream = spark.register_kafka_topic(
            name="transactions_stream",
            topic="transactions",
            timestamp_expression="TIMESTAMP_MILLIS(CAST(get_json_object(CAST(value AS STRING), '$.ts_epoch_ms') AS BIGINT))",
            watermark_delay="10 minutes",
        )
        ```

        Args:
            name (str): Name to register the Kafka topic under.
            topic (str): The name of the Kafka topic.
            timestamp_expression (str): A Spark SQL expression that extracts the event timestamp.
                Must evaluate to a TIMESTAMP type. This is used for watermarking in streaming
                JOINs and late data handling. If empty, watermarking will not be applied.
            watermark_delay (str): How late data can arrive and still be processed
                (e.g., "5 minutes", "1 hour"). Defaults to "5 minutes".
            description (str, optional): Description of the Kafka topic.
            tags (List[str], optional): Tags associated with the data source.
            properties (dict, optional): Additional properties.
            owner (Union[str, "UserRegistrar"], optional): Owner of the data source.

        Returns:
            KafkaTopic: The registered Kafka topic.
        """

        tags, properties = set_tags_properties(tags, properties)

        kafka_topic = KafkaTopic(
            name=name,
            topic=topic,
            provider=self.name(),
            timestamp_expression=timestamp_expression,
            watermark_delay=watermark_delay,
            owner=owner,
            description=description,
        )
        # Register the StreamChannel resource
        self.__registrar.add_resource(kafka_topic)

        # Also register a SourceVariant with PrimaryData so that transformations
        # can reference this Kafka topic as an input. The variant uses the
        # STREAM_CHANNEL_VARIANT constant to match KafkaTopic.name_variant().
        self.__registrar.register_primary_data(
            name=name,
            location=kafka_topic,
            provider=self.name(),
            variant=STREAM_CHANNEL_VARIANT,
            owner=owner,
            description=description,
            tags=tags or [],
            properties=properties or {},
        )
        return kafka_topic
