# 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/.

"""
Location classes for Featureform resources.

This module contains all location-related classes including SQL tables, file stores,
stream channels, and catalog tables.
"""

import re
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any, Tuple

from typeguard import typechecked

from ..enums import OperationType, ResourceType, TableFormat
from ..proto import metadata_pb2 as pb
from .base import Location
from .constants import STREAM_CHANNEL_VARIANT

if TYPE_CHECKING:
    from .mappings import FeaturesSchema
else:
    # At runtime, use Any to avoid circular import issues with typeguard
    FeaturesSchema = Any


class StreamingInput(Location):
    """
    Abstract base class for streaming input sources.

    All streaming input types (e.g., KafkaTopic) should extend this class.
    This allows for type checking to determine if an input is a streaming input.
    """

    pass


@dataclass
class StreamChannelReference(StreamingInput):
    """
    Lightweight reference to a StreamChannel by name.

    Used when users pass just the name string to a streaming transformation.
    The actual StreamChannel is looked up on the server side during execution.

    Example:
        @spark.streaming_sql_transformation(inputs=["events"])
        def my_transform(events):
            return "SELECT * FROM {{ events }}"

        # The string "events" is normalized to StreamChannelReference(name="events")
    """

    name: str

    def name_variant(self):
        """Return name and variant tuple for transformation input resolution."""
        return (self.name, STREAM_CHANNEL_VARIANT)

    def resource_identifier(self):
        return self.name


@typechecked
@dataclass
class SQLTable(Location):
    name: str
    schema: str = ""
    database: str = ""

    def resource_identifier(self):
        return f"{self.database}.{self.schema}.{self.name}"

    @staticmethod
    def from_proto(source_table):
        return SQLTable(
            schema=source_table.schema,
            database=source_table.database,
            name=source_table.name,
        )


@typechecked
@dataclass
class FileStore(Location):
    """
    Contains the location of a path in a cloud file store (S3, GCS, Azure)
    """

    path_uri: str
    format: str = ""
    multiline: bool = False

    def resource_identifier(self):
        return self.path_uri

    @staticmethod
    def from_proto(source_filestore):
        return FileStore(
            path_uri=source_filestore.path,
            format=source_filestore.format,
            multiline=source_filestore.multiline,
        )


@typechecked
@dataclass
class KafkaTopic(StreamingInput):
    """
    KafkaTopic represents a Kafka topic that can be used as a streaming input.

    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

    Common timestamp_expression patterns:
    - "timestamp" - Use Kafka's ingestion timestamp
    - "CAST(get_json_object(CAST(value AS STRING), '$.event_time') AS TIMESTAMP)" - JSON field
    - "TIMESTAMP_MILLIS(CAST(get_json_object(CAST(value AS STRING), '$.ts_ms') AS BIGINT))" - Epoch ms
    """

    name: str
    topic: str
    provider: str
    timestamp_expression: str = ""
    watermark_delay: str = "5 minutes"
    description: str = ""
    owner: str = ""
    created: str = None
    last_updated: str = None

    def __post_init__(self):
        """Validate watermark_delay format matches Spark duration syntax."""
        if self.watermark_delay:
            pattern = r"^(\d+)\s+(second|seconds|minute|minutes|hour|hours|day|days|week|weeks)$"
            if not re.match(pattern, self.watermark_delay.strip().lower()):
                raise ValueError(
                    f"Invalid watermark_delay format: '{self.watermark_delay}'. "
                    f"Expected format: 'N unit' (e.g., '5 minutes', '1 hour')"
                )

    @staticmethod
    def from_proto(source_stream):
        return KafkaTopic(
            name=source_stream.name,
            topic=source_stream.channel_name,
            provider=source_stream.provider,
            timestamp_expression=source_stream.timestamp_expression,
            watermark_delay=source_stream.watermark_delay or "5 minutes",
        )

    @staticmethod
    def get_resource_type() -> ResourceType:
        return ResourceType.STREAM_CHANNEL

    @staticmethod
    def operation_type() -> OperationType:
        return OperationType.CREATE

    @staticmethod
    def get(stub, name: str) -> "KafkaTopic":
        name_req = pb.NameRequest(name=pb.Name(name=name))
        kafka_topic = next(stub.GetStreamChannels(iter([name_req])))

        if not kafka_topic:
            raise ValueError(f"Kafka topic {name} not found")

        return KafkaTopic(
            name=kafka_topic.name,
            topic=kafka_topic.channel_name,
            provider=kafka_topic.provider,
            timestamp_expression=kafka_topic.timestamp_expression,
            watermark_delay=kafka_topic.watermark_delay or "5 minutes",
            description=kafka_topic.description,
            owner=kafka_topic.owner,
            created=kafka_topic.created,
            last_updated=kafka_topic.last_updated,
        )

    def resource_identifier(self):
        return self.name

    def name_variant(self):
        """Return name and variant tuple for transformation input resolution."""
        return (self.name, STREAM_CHANNEL_VARIANT)

    def _create(self, req_id, stub) -> Tuple[None, None]:
        serialized = pb.StreamChannelRequest(
            stream_channel=pb.StreamChannel(
                name=self.name,
                channel_name=self.topic,
                provider=self.provider,
                timestamp_expression=self.timestamp_expression,
                watermark_delay=self.watermark_delay,
                description=self.description,
                owner=self.owner,
            ),
            request_id="",
        )

        stub.CreateStreamChannel(serialized)
        return None, None


@typechecked
@dataclass
class FileStoreStreamingSource(StreamingInput):
    """
    FileStoreStreamingSource represents a file-based streaming source for testing.

    Reads parquet/CSV/JSON files from a filestore (S3/GCS/Azure) as a bounded stream.
    The stream terminates after processing all available files using Spark's
    `availableNow` trigger. Schema is inferred automatically from the files.

    This source type is useful for:
    - End-to-end testing of streaming feature views without Kafka
    - Processing bounded datasets through the streaming code path
    - Development and debugging of streaming transformations

    Note: Internally, this uses the StreamChannel resource with trigger_mode=AVAILABLE_NOW.

    Example:
        spark = client.get_provider("my_spark")
        streaming_source = spark._register_filestore_streaming_source(
            name="test_transactions",
            path="s3://bucket/test-data/transactions.parquet",
            file_format="parquet",
            description="Test transaction data for streaming E2E tests",
        )
    """

    name: str
    path: str  # S3/GCS/Azure path to files
    provider: str
    file_format: str = "parquet"  # parquet, csv, json
    description: str = ""
    owner: str = ""
    created: str = None
    last_updated: str = None

    def __post_init__(self):
        """Validate file_format is a supported format."""
        supported_formats = ("parquet", "csv", "json")
        if self.file_format.lower() not in supported_formats:
            raise ValueError(
                f"Invalid file_format: '{self.file_format}'. "
                f"Supported formats: {supported_formats}"
            )

    @staticmethod
    def from_proto(source_stream):
        """Create FileStoreStreamingSource from a StreamChannel proto with trigger_mode=AVAILABLE_NOW."""
        return FileStoreStreamingSource(
            name=source_stream.name,
            path=source_stream.channel_name,  # path is stored in channel_name
            provider=source_stream.provider,
            file_format=source_stream.file_format or "parquet",
            description=source_stream.description,
            owner=source_stream.owner,
        )

    @staticmethod
    def get_resource_type() -> ResourceType:
        # Uses STREAM_CHANNEL with trigger_mode=AVAILABLE_NOW internally
        return ResourceType.STREAM_CHANNEL

    @staticmethod
    def operation_type() -> OperationType:
        return OperationType.CREATE

    @staticmethod
    def get(stub, name: str) -> "FileStoreStreamingSource":
        """Get a FileStoreStreamingSource by name (retrieves a bounded StreamChannel)."""
        name_req = pb.NameRequest(name=pb.Name(name=name))
        source = next(stub.GetStreamChannels(iter([name_req])))

        if not source:
            raise ValueError(f"FileStoreStreamingSource {name} not found")

        # Verify it's a bounded stream (file-based) using AVAILABLE_NOW trigger mode
        if source.trigger_mode != pb.TRIGGER_MODE_AVAILABLE_NOW:
            raise ValueError(
                f"StreamChannel {name} is not a bounded stream (FileStoreStreamingSource)"
            )

        return FileStoreStreamingSource(
            name=source.name,
            path=source.channel_name,  # path is stored in channel_name
            provider=source.provider,
            file_format=source.file_format or "parquet",
            description=source.description,
            owner=source.owner,
            created=source.created,
            last_updated=source.last_updated,
        )

    def resource_identifier(self):
        return self.name

    def name_variant(self):
        """Return name and variant tuple for transformation input resolution."""
        return (self.name, STREAM_CHANNEL_VARIANT)

    def _create(self, req_id, stub) -> Tuple[None, None]:
        """Create a StreamChannel with trigger_mode=AVAILABLE_NOW for file-based streaming."""
        serialized = pb.StreamChannelRequest(
            stream_channel=pb.StreamChannel(
                name=self.name,
                channel_name=self.path,  # store path in channel_name
                provider=self.provider,
                file_format=self.file_format,
                trigger_mode=pb.TRIGGER_MODE_AVAILABLE_NOW,  # bounded file-based stream
                description=self.description,
                owner=self.owner,
            ),
            request_id="",
        )

        stub.CreateStreamChannel(serialized)
        return None, None


@typechecked
@dataclass
class GlueCatalogTable(Location):
    database: str
    table: str
    table_format: TableFormat = field(default_factory=lambda: TableFormat.ICEBERG)

    def resource_identifier(self):
        return f"{self.database}.{self.table}"

    @staticmethod
    def from_proto(source_catalog_table):
        return GlueCatalogTable(
            database=source_catalog_table.database,
            table=source_catalog_table.table,
            table_format=TableFormat.get_format(source_catalog_table.table_format),
        )


@typechecked
@dataclass
class UnityCatalogTable(Location):
    database: str
    schema: str
    table: str
    table_format: TableFormat = field(default_factory=lambda: TableFormat.ICEBERG)

    def resource_identifier(self):
        return f"{self.database}.{self.schema}.{self.table}"

    @staticmethod
    def from_proto(source_catalog_table):
        return UnityCatalogTable(
            database=source_catalog_table.database,
            schema=source_catalog_table.schema,
            table=source_catalog_table.table,
            table_format=TableFormat.get_format(source_catalog_table.table_format),
        )


@typechecked
@dataclass
class Directory(Location):
    path: str

    def path(self):
        return self.path


@typechecked
@dataclass
class OnlineLocation(Location):
    """Online location for feature serving."""

    table: str
    version: int
    schema: "FeaturesSchema"  # Forward reference to avoid circular import

    def resource_identifier(self):
        return f"{self.table}__{self.version}"

    @staticmethod
    def from_proto(proto: pb.Location):
        # Import here to avoid circular dependency
        from .mappings import FeaturesSchema

        online_proto = proto.online
        return (
            OnlineLocation(
                table=online_proto.table,
                version=online_proto.version,
                schema=FeaturesSchema.from_proto(online_proto.schema),
            )
            if online_proto
            else None
        )
