import logging
import os
import time
from pathlib import Path

# PyFlink imports
from pyflink.table import EnvironmentSettings, StreamTableEnvironment, DataTypes
from pyflink.table.udf import udf
from pyflink.datastream import StreamExecutionEnvironment
from pyflink.common.configuration import Configuration

# Configuration
aws_region = "us-east-2"

ONE_HOUR_IN_SECONDS = 3600


# Configuration class to read runtime properties
class SimpleConfig:
    """Configuration for hooks positive signal service with CDK property group support"""

    def __init__(self):
        # Set default values first (for local development)
        self.stream_arn = (
            "arn:aws:kinesis:us-east-2:590183763515:stream/rec-events-stream"
        )
        self.redis_host = "localhost"
        self.redis_port = 6379
        self.env = "dev"
        self.playback_threshold = 5  # seconds

        # Try to get runtime properties (for cloud deployment)
        try:
            from pyflink.java_gateway import get_gateway

            gateway = get_gateway()

            # Get application properties using KinesisAnalyticsRuntime
            j_kinesis_analytics_runtime = gateway.jvm.com.amazonaws.services.kinesisanalytics.runtime.KinesisAnalyticsRuntime
            application_properties = (
                j_kinesis_analytics_runtime.getApplicationProperties()
            )

            # Get the 'rec.config' property group (as defined in CDK)
            rec_config_properties = application_properties.get("rec.config")

            if rec_config_properties:
                # Override defaults with runtime properties
                self.stream_arn = (
                    rec_config_properties.getProperty("stream.arn") or self.stream_arn
                )
                self.redis_host = (
                    rec_config_properties.getProperty("redis.host") or self.redis_host
                )
                self.redis_port = int(
                    rec_config_properties.getProperty("redis.port")
                    or str(self.redis_port)
                )
                self.env = rec_config_properties.getProperty("env.stage") or self.env
                self.playback_threshold = float(
                    rec_config_properties.getProperty("playback.threshold")
                    or str(self.playback_threshold)
                )

        except Exception as e:
            logging.error(f"Error getting application properties: {e}")
            logging.info("Using default local development configuration")


logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


class HooksPositiveSignalService:
    """Service for processing hooks positive signals (likes, unlikes, playback)"""

    def __init__(self, config):
        self.config = config

        if self.config.env == "prod" or self.config.env == "staging":
            # Initialize Flink environment for production/staging
            current_dir = Path(__file__).resolve().parent
            pyflink_jar_path = current_dir / "lib/pyflink-dependencies.jar"

            conf = Configuration()
            if pyflink_jar_path.exists():
                conf.set_string("pipeline.jars", f"file://{pyflink_jar_path}")
                print(f"Setting pipeline.jars to: file://{pyflink_jar_path}")
            else:
                print(f"⚠️ Warning: {pyflink_jar_path} not found")

            self.stream_env = StreamExecutionEnvironment.get_execution_environment()

            self.env_settings = (
                EnvironmentSettings.new_instance()
                .with_configuration(conf)
                .in_streaming_mode()
                .build()
            )

            self.table_env = StreamTableEnvironment.create(
                self.stream_env, self.env_settings
            )

            # Add JARs to the table environment
            if os.path.exists(pyflink_jar_path):
                jar_url = f"file://{pyflink_jar_path}"
                print(f"Setting pipeline.jars to: {jar_url}")
                self.table_env.get_config().get_configuration().set_string(
                    "pipeline.jars", jar_url
                )

            self._configure_environment()
            self._register_udfs()
        else:
            # Initialize Flink environment for local development
            self.stream_env = StreamExecutionEnvironment.get_execution_environment()

            # Add connector JAR for Kinesis
            current_dir = os.path.dirname(os.path.abspath(__file__))
            parent_dir = os.path.dirname(current_dir)
            kinesis_jar_path = os.path.join(
                parent_dir, "jar/flink-sql-connector-kinesis-5.0.0-1.20.jar"
            )

            conf = Configuration()
            if os.path.exists(kinesis_jar_path):
                conf.set_string("pipeline.jars", f"file://{kinesis_jar_path}")
                print(f"Loaded Kinesis connector from: {kinesis_jar_path}")
            else:
                print(f"⚠️ Warning: Kinesis JAR not found at {kinesis_jar_path}")

            self.env_settings = (
                EnvironmentSettings.new_instance()
                .with_configuration(conf)
                .in_streaming_mode()
                .build()
            )

            self.table_env = StreamTableEnvironment.create(
                self.stream_env, self.env_settings
            )

            self._configure_environment()
            self._register_udfs()

    def _configure_environment(self):
        configuration = self.table_env.get_config().get_configuration()

        # Exactly-once processing
        configuration.set_string("execution.checkpointing.mode", "EXACTLY_ONCE")
        configuration.set_string(
            "execution.checkpointing.externalized-checkpoint-retention",
            "RETAIN_ON_CANCELLATION",
        )

        # State TTL configuration
        configuration.set_string("table.exec.state.ttl", "36h")
        configuration.set_string("table.exec.state.ttl.cleanup.strategy", "delete")

        # Mini-batch processing for better performance
        configuration.set_string("table.exec.mini-batch.enabled", "true")
        configuration.set_string("table.exec.mini-batch.allow-latency", "5s")
        configuration.set_string("table.exec.mini-batch.size", "5000")

    def _register_udfs(self):
        # Store config values in local variables to avoid serialization issues
        redis_host = self.config.redis_host
        redis_port = self.config.redis_port
        playback_threshold = self.config.playback_threshold

        @udf(result_type=DataTypes.BOOLEAN())
        def process_playback_event(
            user_id: str, hook_id: str, play_duration: float
        ) -> bool:
            try:
                import redis

                redis_client = redis.Redis(
                    host=redis_host, port=redis_port, decode_responses=True
                )

                # Key for tracking playback duration per user-hook combination
                playback_key = f"hooks_playback_duration:{user_id}:{hook_id}"
                positive_signals_long_watch_key = (
                    f"hooks_positive_signal_long_watch:{user_id}"
                )

                # Get current playback duration for this user-hook combination
                current_duration_str = redis_client.get(playback_key)
                current_duration = (
                    float(current_duration_str) if current_duration_str else 0.0
                )

                # Add new playback duration
                new_duration = current_duration + play_duration
                redis_client.set(playback_key, str(new_duration))
                redis_client.expire(playback_key, ONE_HOUR_IN_SECONDS)

                print(
                    f"Playback: User {user_id}, Hook {hook_id}, Duration: {current_duration:.1f}s + {play_duration:.1f}s = {new_duration:.1f}s"
                )

                # Check if threshold is exceeded
                if new_duration >= playback_threshold:
                    # Get current timestamp for scoring
                    current_timestamp = time.time()

                    # Add hook to positive signals sorted set with timestamp as score
                    # This automatically handles duplicates and maintains chronological order
                    added_count = redis_client.zadd(
                        positive_signals_long_watch_key, {hook_id: current_timestamp}
                    )

                    # Keep only the most recent 200 entries in the sorted set
                    # Remove all elements except the top 200 (highest scores/most recent)
                    redis_client.zremrangebyrank(positive_signals_long_watch_key, 0, -201)

                    if added_count > 0:
                        print(
                            f"Added hook {hook_id} to positive signals for user {user_id} (playback threshold exceeded) at timestamp {current_timestamp}"
                        )
                    else:
                        print(
                            f"Updated timestamp for existing hook {hook_id} in positive signals for user {user_id}"
                        )
                    return True

                return False  # No change to positive signals

            except Exception as e:
                print(
                    f"Error processing playback event for user {user_id}, hook {hook_id}: {e}"
                )
                return False

        self.table_env.create_temporary_system_function(
            "process_playback_event", process_playback_event
        )

        logger.info("UDFs registered successfully")

    def create_kinesis_source_table(self):
        """Create Kinesis source table for hook events"""
        ddl = f"""
            CREATE TABLE hook_events (
                name STRING,
                `timestamp` STRING,
                source STRING,
                user_id STRING,
                session_id STRING,
                properties ROW<
                    hook_id STRING,
                    hook_status STRING,
                    recommendation_item_id STRING,
                    hook_play_duration DOUBLE
                >,
                request_id STRING
            ) WITH (
            'connector' = 'kinesis',
            'stream.arn' = '{self.config.stream_arn}',
            'aws.credentials.provider' = 'AUTO',
            'aws.region' = '{aws_region}',
            'source.init.position' = 'LATEST',
            'format' = 'json',
            'source.reader.type' = 'EFO',
            'source.efo.consumer.name' = 'hooks-positive-signal-consumer-{self.config.env}',
            'source.efo.lifecycle' = 'SELF_MANAGED'
        )
        """
        self.table_env.execute_sql(ddl)
        print("Executing DDL:")
        print(ddl)

    def run_pipeline(self):
        try:
            self.create_kinesis_source_table()
            job = self.process_hook_events()
            job.wait()
        except Exception as e:
            logger.error(f"Error in pipeline execution: {e}")
            raise

    def process_hook_events(self):
        # Create sink table for processing results
        create_sink_sql = """
        CREATE TABLE hook_events_sink (
            user_id STRING,
            hook_id STRING,
            success BOOLEAN
        ) WITH (
            'connector' = 'print',
            'print-identifier' = 'HOOK-EVENTS'
        )
        """
        self.table_env.execute_sql(create_sink_sql)

        playback_sql = """
        INSERT INTO hook_events_sink
        SELECT
            user_id,
            properties.hook_id,
            process_playback_event(user_id, properties.hook_id, properties.hook_play_duration) as success
        FROM hook_events
        WHERE name = 'HooksPlayDuration'
          AND user_id IS NOT NULL
          AND properties.hook_id IS NOT NULL
          AND properties.hook_play_duration IS NOT NULL
        """

        print("Processing hook events...")

        playback_job = self.table_env.execute_sql(playback_sql)

        return playback_job


def main():
    try:
        config = SimpleConfig()
        service = HooksPositiveSignalService(config)
        service.run_pipeline()
    except Exception as e:
        logger.error(f"Service failed: {e}")
        raise


if __name__ == "__main__":
    main()
