#!/usr/bin/env python3
"""
Redis Validation Script for Hooks Recommender (runs on Modal)

Usage:
    uv run modal run validate_redis_hooks.py
    uv run modal run validate_redis_hooks.py --user-id 12345 --user-uuid "uuid-here"
    uv run modal run validate_redis_hooks.py --json-output
"""

import json
import os
from typing import Any, Dict, List
import modal
from redis import Redis

from suno_recs.config import get_modal_secrets
from suno_recs.worker.modal_recs_app_config import app, config

# Default test values
DEFAULT_USER_ID = 62406164
DEFAULT_USER_UUID = "3788780e-8909-405e-a05b-c334f8387d69"


def format_table(rows: List[List[str]], headers: List[str]) -> str:
    """Format data as a properly aligned ASCII table"""
    # Calculate column widths
    col_widths = []
    for i, header in enumerate(headers):
        max_width = len(header)
        for row in rows:
            if i < len(row):
                max_width = max(max_width, len(str(row[i])))
        col_widths.append(max_width)
    
    # Build table
    lines = []
    
    # Header separator
    separator = "+" + "+".join(["-" * (w + 2) for w in col_widths]) + "+"
    lines.append(separator)
    
    # Headers
    header_row = "|"
    for i, header in enumerate(headers):
        header_row += f" {header:<{col_widths[i]}} |"
    lines.append(header_row)
    lines.append(separator)
    
    # Data rows
    for row in rows:
        data_row = "|"
        for i in range(len(headers)):
            if i < len(row):
                cell = str(row[i])
            else:
                cell = ""
            data_row += f" {cell:<{col_widths[i]}} |"
        lines.append(data_row)
    
    # Bottom border
    lines.append(separator)
    
    return "\n".join(lines)


@app.function(
    secrets=get_modal_secrets(config.deployment_type),
    proxy=modal.Proxy.from_name("data-layer-proxy"),
    region="us-east-2",
    timeout=30,
)
def validate_redis(user_id: int = DEFAULT_USER_ID, user_uuid: str = DEFAULT_USER_UUID) -> Dict[str, Any]:
    """Run Redis validation on Modal with proper proxy access"""
    
    # Get Redis URLs from environment
    redis_endpoints = {
        "REDIS_RECS": os.getenv("REDIS_RECS_URL"),
        "REDIS_HOOK_HISTORY": os.getenv("REDIS_HOOK_LISTEN_HISTORY_URL"),
        "VALKEY": os.getenv("VALKEY_URL"),
    }
    
    print(f"\n{'='*80}")
    print("REDIS VALIDATION FOR HOOKS RECOMMENDER")
    print(f"{'='*80}")
    print(f"User ID: {user_id}")
    print(f"User UUID: {user_uuid}")
    print(f"{'='*80}\n")
    
    # Initialize Redis clients
    print("Connecting to Redis clusters...")
    clients = {}
    for name, host in redis_endpoints.items():
        try:
            client = Redis(
                host=host, 
                port=6379, 
                decode_responses=True,
                socket_connect_timeout=5,
                socket_timeout=5
            )
            client.ping()
            clients[name] = client
            print(f"  [OK] {name}: Connected to {host}")
        except Exception as e:
            print(f"  [ERROR] {name}: Failed - {e}")
            clients[name] = None
    
    print()
    
    # Define all keys to check
    checks = [
        # Watch history
        ("REDIS_HOOK_HISTORY", f"{user_id}:hook_listening_history_storage", "LRANGE", "watch_history", "Hook watch history"),
        
        # User profile (VALKEY)
        ("VALKEY", f"user:{user_id}:tag_freq", "ZREVRANGE", "user_profile", "User creation tags"),
        ("VALKEY", f"feature_top_listening_clip:{user_id}", "GET", "user_profile", "Top listening clips"),
        ("VALKEY", f"REACT:LIKE:USER:{user_id}:", "ZREVRANGE", "user_profile", "Liked clips (songs)"),
        ("VALKEY", f"REACT:DISLIKE:USER:{user_id}:", "ZREVRANGE", "user_profile", "Disliked clips"),
        
        # Hook signals (REDIS_RECS)
        ("REDIS_RECS", f"hook_rec_user_like_history:{user_id}", "GET", "hook_signals", "Liked hooks (Flink)"),
        ("REDIS_RECS", f"hooks_positive_signal_long_watch:{user_uuid}", "ZREVRANGE", "hook_signals", "Long watch signal"),
        ("REDIS_RECS", f"hooks_positive_signal_omniplay:{user_uuid}", "ZREVRANGE", "hook_signals", "Omniplay signal"),
        ("REDIS_RECS", f"hooks_positive_signal_comment:{user_id}", "ZREVRANGE", "hook_signals", "Comment signal"),
        ("REDIS_RECS", f"hooks_positive_signal_comment_view:{user_id}", "ZREVRANGE", "hook_signals", "Comment view signal"),
        ("REDIS_RECS", f"hooks_positive_signal_profile_view:{user_id}", "ZREVRANGE", "hook_signals", "Profile view signal"),
        ("REDIS_RECS", f"hooks_positive_signal_remix:{user_uuid}", "ZREVRANGE", "hook_signals", "Remix signal"),
        ("REDIS_RECS", f"hooks_positive_signal_rewatch:{user_uuid}", "ZREVRANGE", "hook_signals", "Rewatch signal"),
        
        # Collaborative filtering
        ("REDIS_RECS", f"hook_recommendations:{user_id}", "GET", "collaborative", "CF recommendations"),
        
        # Champion content
        ("REDIS_RECS", "champion_hook_creator", "SMEMBERS", "champion", "Champion creators"),
        
        # Bucket caches
        ("REDIS_RECS", "hooks_cache:fresh:v4", "TYPE", "caches", "Fresh cache"),
        ("REDIS_RECS", "hooks_cache:popular:v4", "TYPE", "caches", "Popular cache"),
        ("REDIS_RECS", "hooks_cache:creator_champion:v4", "TYPE", "caches", "Champion creator cache"),
        ("REDIS_RECS", "hooks_cache:manual_champion:v4", "TYPE", "caches", "Manual champion cache"),
    ]
    
    results = {}
    table_rows = []
    
    print(f"Checking {len(checks)} Redis keys...\n")
    
    for client_name, key, operation, category, description in checks:
        if category not in results:
            results[category] = []
        
        client = clients.get(client_name)
        if not client:
            result = {
                "key": key,
                "description": description,
                "host": redis_endpoints.get(client_name),
                "exists": False,
                "error": f"No connection to {client_name}",
                "count": 0,
                "type": "N/A"
            }
        else:
            result = {
                "key": key,
                "description": description,
                "host": redis_endpoints.get(client_name),
                "exists": False,
                "error": None,
                "count": 0,
                "type": "N/A"
            }
            
            try:
                # Get the Redis data type
                redis_type = client.type(key)
                result["type"] = redis_type if redis_type != "none" else "N/A"
                
                if operation == "GET":
                    data = client.get(key)
                    result["exists"] = data is not None
                    if data:
                        try:
                            parsed = json.loads(data)
                            result["count"] = len(parsed) if isinstance(parsed, list) else 1
                        except:
                            result["count"] = 1
                            
                elif operation == "LRANGE":
                    count = client.llen(key)
                    result["exists"] = count > 0
                    result["count"] = count
                    
                elif operation == "ZREVRANGE":
                    count = client.zcard(key)
                    result["exists"] = count > 0
                    result["count"] = count
                    
                elif operation == "SMEMBERS":
                    count = client.scard(key)
                    result["exists"] = count > 0
                    result["count"] = count
                    
                elif operation == "TYPE":
                    # For TYPE operation, just check if key exists
                    result["exists"] = redis_type != "none"
                    result["count"] = 1 if result["exists"] else 0
                
                # Get TTL for any existing key
                if result["exists"]:
                    ttl = client.ttl(key)
                    if ttl > 0:
                        result["ttl_minutes"] = round(ttl / 60, 1)
                    elif ttl == -1:
                        result["ttl_minutes"] = "no expiry"
                    else:
                        result["ttl_minutes"] = "expired"
                    
            except Exception as e:
                result["error"] = str(e)
                result["type"] = "ERROR"
        
        results[category].append(result)
        
        # Prepare row for table
        status = "EXISTS" if result["exists"] else "ERROR" if result["error"] else "EMPTY"
        count_str = str(result["count"]) if result["exists"] else "0"
        
        # Get TTL info
        ttl_str = "-"
        if "ttl_minutes" in result:
            if isinstance(result['ttl_minutes'], (int, float)):
                ttl_str = f"{result['ttl_minutes']}m"
            else:
                ttl_str = str(result['ttl_minutes'])
        
        # Truncate key for display
        display_key = key if len(key) <= 45 else key[:42] + "..."
        
        table_rows.append([
            status,
            description,
            display_key,
            result["type"],
            count_str,
            ttl_str,
            client_name
        ])
    
    # Print table
    headers = ["Status", "Description", "Key", "Type", "Count", "TTL", "Cluster"]
    print(format_table(table_rows, headers))
    
    # Simple summary
    print(f"\n{'='*80}")
    total_keys = sum(len(v) for v in results.values())
    existing_keys = sum(1 for v in results.values() for r in v if r["exists"])
    empty_keys = total_keys - existing_keys
    
    print(f"Total keys checked: {total_keys}")
    print(f"Keys with data: {existing_keys}")
    print(f"Empty keys: {empty_keys}")
    print(f"{'='*80}")
    
    return results


@app.local_entrypoint()
def main(user_id: int = DEFAULT_USER_ID, user_uuid: str = DEFAULT_USER_UUID, json_output: bool = False):
    """Local entrypoint - Modal automatically maps CLI args to these parameters"""
    
    # Run validation on Modal
    results = validate_redis.remote(user_id, user_uuid)
    
    # Output JSON if requested
    if json_output:
        print("\n" + "="*80)
        print("JSON OUTPUT")
        print("="*80)
        print(json.dumps(results, indent=2, default=str))