import logging
import os
import redis, time

# PyFlink imports
from pyflink.table import EnvironmentSettings, StreamTableEnvironment
from pyflink.datastream import StreamExecutionEnvironment
from pyflink.common.configuration import Configuration
from pyflink.common import RowKind
from pyflink.datastream.functions import MapFunction

aws_region = "us-east-2"

# Configuration class to read runtime properties
class SimpleConfig:
    
    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"
        
        # 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
                
        except Exception as e:
            logger.error(f"Error getting application properties: {e}")
            logger.info("Using default local development configuration")

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

class CollaborativeFilteringService:
    """Service for collaborative filtering recommendations"""
    
    def __init__(self, config):
        self.config = config
        from pathlib import Path

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

            self.stream_env = StreamExecutionEnvironment.get_execution_environment()

            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.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()

        else:
            # Initialize Flink environment for dev
            current_dir = Path(__file__).resolve().parent.parent
            kinesis_jar_path = current_dir / "jar/flink-sql-connector-kinesis-5.0.0-1.20.jar"

            self.stream_env = StreamExecutionEnvironment.get_execution_environment()

            conf = Configuration()
            if kinesis_jar_path.exists():
                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_path} not found")

            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()
    
    def _configure_environment(self):
        if self.config.env == "prod" or self.config.env == "staging":
            configuration = self.table_env.get_config().get_configuration()

            # 1. Set State TTL to 24 hours (automatically removes old state)
            configuration.set_string("table.exec.state.ttl", "1h")

            # 2. Enable idle state cleanup (removes unused state)
            configuration.set_string("table.exec.state.ttl.cleanup.strategy", "delete")

            # 3. Handle data skew and idle subtasks for watermark advancement
            configuration.set_string("table.exec.source.idle-timeout", "30s")

            # optional delete optimization
            #configuration.set_string("table.exec.state.ttl.cleanup.interval", "1h")
            configuration.set_string("table.exec.mini-batch.enabled", "false")
            configuration.set_string("table.exec.mini-batch.allow-latency", "5s")
            configuration.set_string("table.exec.mini-batch.size", "1000")
        else:
            self.stream_env.set_parallelism(8)
            self.stream_env.enable_checkpointing(60000)  # 1 minute
            
            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")

            # 1. Set State TTL to 24 hours (automatically removes old state)
            configuration.set_string("table.exec.state.ttl", "1h")

            # 2. Enable idle state cleanup (removes unused state)
            configuration.set_string("table.exec.state.ttl.cleanup.strategy", "delete")

            # 3. Handle data skew and idle subtasks for watermark advancement
            configuration.set_string("table.exec.source.idle-timeout", "30s")

            # optional delete optimization
            #configuration.set_string("table.exec.state.ttl.cleanup.interval", "1h")
            configuration.set_string("table.exec.mini-batch.enabled", "false")
            configuration.set_string("table.exec.mini-batch.allow-latency", "5s")
            configuration.set_string("table.exec.mini-batch.size", "1000")
    
    def create_kinesis_source_table(self):
        """Create Kinesis source table for hook events"""
        ddl = f"""
        CREATE TABLE hook_events (
            `timestamp` STRING,
            `name` STRING,
            `user_id` STRING,
            `properties` ROW<`hook_id` STRING>,
            `proc_time` AS PROCTIME()
        ) WITH (
            'connector' = 'kinesis',
            'stream.arn' = '{self.config.stream_arn}',
            'aws.credentials.provider' = 'AUTO',
            'source.init.position' = 'LATEST',
            'aws.region' = '{aws_region}',
            'format' = 'json',
            'source.reader.type' = 'EFO',
            'source.efo.consumer.name' = 'co-existing-like-consumer-{self.config.env}',
            'source.efo.lifecycle' = 'SELF_MANAGED'
        )
        """
        self.table_env.execute_sql(ddl)
    
    def run_pipeline(self):
        try:
            self.create_kinesis_source_table()
            job = self.create_co_occuring()            
            job.wait()
        except Exception as e:
            logger.error(f"Error in pipeline execution: {e}")
            raise

    def create_co_occuring(self):
        """Create co-occurrence view and write results to sink"""

        most_recent_like_song_per_user = """
        CREATE VIEW most_recent_like_song_per_user AS
        SELECT
            user_id,
            hook_id,
            proc_time
        FROM (
            SELECT
                user_id,
                properties.hook_id AS hook_id,
                proc_time,
                name AS event_type,
                ROW_NUMBER() OVER (
                    PARTITION BY user_id, properties.hook_id
                    ORDER BY proc_time DESC
                ) AS rn
            FROM hook_events
            WHERE user_id IS NOT NULL
            AND properties.hook_id IS NOT NULL
            AND user_id <> ''
            AND properties.hook_id <> ''
            AND name IN ('HookLike','HookUndoLike')
            AND proc_time >= CURRENT_TIMESTAMP - INTERVAL '24' HOUR
        )
        WHERE rn = 1 AND event_type = 'HookLike';
        """
        self.table_env.execute_sql(most_recent_like_song_per_user)

        # 1. Create intermediate view with co-occurring hook pairs
        co_occurrence_matrix_sql = """
        CREATE VIEW co_occurrence_matrix_sql AS
        SELECT
            a.hook_id AS hook1,
            b.hook_id AS hook2,
            COUNT(*) as co_occurrence_count
        FROM most_recent_like_song_per_user a
        JOIN most_recent_like_song_per_user b
          ON a.user_id = b.user_id
        WHERE a.hook_id < b.hook_id
        GROUP BY a.hook_id, b.hook_id
        """
        self.table_env.execute_sql(co_occurrence_matrix_sql)

        top_hook_co_occurrence = """
        CREATE VIEW top_hook_recommendations AS
        WITH ranked_a AS (
        SELECT
            hook1,
            hook2,
            co_occurrence_count,
            ROW_NUMBER() OVER (
                PARTITION BY hook1
                ORDER BY co_occurrence_count DESC, hook2, hook1
            ) AS rn
        FROM co_occurrence_matrix_sql
        ),
        top_a AS (
        SELECT hook1, hook2, co_occurrence_count, rn
        FROM ranked_a
        WHERE rn <= 10
        ),
        ranked_b AS (
        SELECT
            hook2 AS hook1,
            hook1 AS hook2,
            co_occurrence_count,
            ROW_NUMBER() OVER (
                PARTITION BY hook1
                ORDER BY co_occurrence_count DESC, hook2, hook1
            ) AS rn
        FROM co_occurrence_matrix_sql
        ),
        top_b AS (
        SELECT hook1, hook2, co_occurrence_count, rn
        FROM ranked_b
        WHERE rn <= 10
        )
        SELECT * FROM top_a
        UNION ALL
        SELECT * FROM top_b;
        """
        self.table_env.execute_sql(top_hook_co_occurrence)

        recommendations_table = self.table_env.sql_query("""
            SELECT
                hook1,
                LISTAGG(CONCAT(hook2, ':', CAST(co_occurrence_count AS STRING)), ',') AS recommendations_str
            FROM top_hook_recommendations
            GROUP BY hook1
        """)
    
        recommendations_stream = self.table_env.to_changelog_stream(recommendations_table)
        
        class BufferedRedisWriter(MapFunction):
            def __init__(self, redis_host, redis_port):
                self.redis_host = redis_host
                self.redis_port = redis_port
                
            def open(self, runtime_context):
                self.redis_client = redis.Redis(host=self.redis_host, port=self.redis_port, decode_responses=True)
                self._buffer = {}
                self._last_flush = time.time()
                self._batch_size = 5
                self._flush_interval = 5

            def map(self, value):
                row_kind = value.get_row_kind()
                hook_id = str(value[0])
                recommendations = str(value[1]) if value[1] else ""

                if row_kind in [RowKind.INSERT, RowKind.UPDATE_AFTER, RowKind.DELETE]:
                    if row_kind == RowKind.DELETE:
                        recommendations = ""
                    self._buffer[hook_id] = recommendations
                
                current_time = time.time()
                if len(self._buffer) >= self._batch_size or (current_time - self._last_flush) >= self._flush_interval:
                    self._flush()
                    self._last_flush = current_time

                return f"Buffered: {hook_id}"

            def _flush(self):
                if not self._buffer:
                    return
                pipe = self.redis_client.pipeline()
                for hook_id, recommendations in self._buffer.items():
                    if recommendations.strip():
                        pipe.set(f"hook_rec_co_occurence_like:{hook_id}", recommendations, ex=86400)
                pipe.execute()
                self._buffer = {}

            def close(self):
                self._flush()

        # Process stream and print results
        processed_stream = recommendations_stream.map(BufferedRedisWriter(self.config.redis_host, self.config.redis_port))
        processed_stream.print()
        
        # Execute the streaming job
        return self.stream_env.execute("Co-occurrence Recommendations with Operation Types")
      
def main():    
    try:
        config = SimpleConfig()
        service = CollaborativeFilteringService(config)        
        service.run_pipeline()        
    except Exception as e:
        logger.error(f"Service failed: {e}")
        raise
  

if __name__ == "__main__":
    main()