import json
import logging
import redis
import os

# 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


# Example with default fallback

aws_region = "us-east-2"

# Configuration class to read runtime properties
class SimpleConfig:
    """Configuration for collaborative filtering 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"
        
        # 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"

            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 - only if file exists
            if pyflink_jar_path.exists():
                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
            from pathlib import Path

            self.stream_env = StreamExecutionEnvironment.get_execution_environment()

            # Add connector JAR for Kinesis
            current_dir = Path(__file__).resolve().parent.parent
            kinesis_jar_path = current_dir / "jar/flink-sql-connector-kinesis-5.0.0-1.20.jar"

            conf = Configuration()
            conf.set_string("pipeline.jars", f"file://{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)

            print(f"Loaded Kinesis connector from: {kinesis_jar_path}")

            self._configure_environment()
            self._register_udfs()
    
    def _configure_environment(self):
        if self.config.env == "prod" or self.config.env == "staging":
            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", "36h")
            
            # 2. Enable idle state cleanup (removes unused state)
            configuration.set_string("table.exec.state.ttl.cleanup.strategy", "delete")

            # optional delete optimization
            #configuration.set_string("table.exec.state.ttl.cleanup.interval", "1h")
            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")
        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", "72h")
            
            # 2. Enable idle state cleanup (removes unused state)
            configuration.set_string("table.exec.state.ttl.cleanup.strategy", "delete")

            # optional delete optimization
            #configuration.set_string("table.exec.state.ttl.cleanup.interval", "1h")
            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
        
        @udf(result_type=DataTypes.BOOLEAN())
        def update_user_like_history(user_id: str, hook_id: str, event_type: str) -> bool:
            """Process each like/unlike event individually"""
            try:
                import redis
                import json
                # Create Redis client inside the UDF to avoid serialization issues
                redis_client = redis.Redis(host=redis_host, port=redis_port, decode_responses=True)
                rec_key = f"hook_rec_user_like_history:{user_id}"
                
                # Get current user's liked hooks
                current_hooks_str = redis_client.get(rec_key)
                print (current_hooks_str)
                if current_hooks_str:
                    try:
                        current_hooks = json.loads(current_hooks_str) if current_hooks_str.startswith('[') else []
                    except:
                        current_hooks = []
                else:
                    current_hooks = []
                
                print(f"Event: {event_type}, User: {user_id}, Hook: {hook_id}")
                
                # Process the specific event
                if event_type == 'HookLike':
                    if hook_id not in current_hooks:
                        current_hooks.append(hook_id)
                        # Keep only the most recent 20
                        current_hooks = current_hooks[-20:]
                        redis_client.set(rec_key, json.dumps(current_hooks))
                        redis_client.expire(rec_key, 86400)
                        print(f"Added hook {hook_id} to user {user_id}")
                        return True
                
                elif event_type == 'HookUndoLike':
                    if hook_id in current_hooks:
                        current_hooks.remove(hook_id)
                        if current_hooks:
                            redis_client.set(rec_key, json.dumps(current_hooks))
                            redis_client.expire(rec_key, 86400)
                        else:
                            redis_client.delete(rec_key)  # Remove empty list
                        print(f"Removed hook {hook_id} from user {user_id}")
                        return True
                
                return False  # No change needed
                
            except Exception as e:
                print(f"Error processing event for user {user_id}: {e}")
                return False

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

        logger.info("UDFs registered successfully (timestamp-based recommendations)")
    
    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',
            'aws.region' = '{aws_region}',
            'source.init.position' = 'LATEST',
            'format' = 'json',
            'source.reader.type' = 'EFO',
            'source.efo.consumer.name' = 'user-like-history-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_user_like_history()            
            job.wait()
        except Exception as e:
            logger.error(f"Error in pipeline execution: {e}")
            raise

    def process_user_like_history(self):
        # Process each event individually - no aggregation
        create_sink_sql = """
        CREATE TABLE individual_event_sink (
            user_id STRING,
            hook_id STRING,
            name STRING,
            success BOOLEAN
        ) WITH (
            'connector' = 'print',
            'print-identifier' = 'INDIVIDUAL-EVENTS'
        )
        """
        self.table_env.execute_sql(create_sink_sql)
        
        # Process each event as it comes in
        insert_sql = """
        INSERT INTO individual_event_sink
        SELECT
            user_id,
            hook_id,
            name,
            update_user_like_history(user_id, hook_id, name) as success
        FROM hook_events
        WHERE name IN ('HookLike', 'HookUndoLike')
        """
        
        print("Processing individual events as they arrive...")
        job = self.table_env.execute_sql(insert_sql)
        return job
    
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()


