"""Redis patching for AutoDebug capture.

This module captures Redis commands at the execute_command level,
supporting both redis-py sync and async clients.
"""

from __future__ import annotations

import hashlib
import json
import traceback
from typing import Any, Callable, Dict, List, Optional, Tuple, TYPE_CHECKING

if TYPE_CHECKING:
    from autodebug_probe.recorder import Recorder

# Original functions
_original_execute_command: Optional[Callable] = None
_original_pipeline_execute: Optional[Callable] = None

# Global recorder reference
_recorder: Optional["Recorder"] = None

# Commands that are safe to capture (read operations)
READ_COMMANDS = {
    "GET", "MGET", "HGET", "HMGET", "HGETALL", "HKEYS", "HVALS", "HLEN",
    "LRANGE", "LINDEX", "LLEN", "SISMEMBER", "SMEMBERS", "SCARD",
    "ZRANGE", "ZRANGEBYSCORE", "ZSCORE", "ZCARD", "ZRANK",
    "EXISTS", "TYPE", "TTL", "PTTL", "KEYS", "SCAN",
    "INFO", "DBSIZE", "PING",
}

# Commands that modify state (write operations)
WRITE_COMMANDS = {
    "SET", "MSET", "SETNX", "SETEX", "PSETEX",
    "HSET", "HMSET", "HSETNX", "HDEL",
    "LPUSH", "RPUSH", "LPOP", "RPOP", "LSET", "LREM",
    "SADD", "SREM", "SPOP",
    "ZADD", "ZREM", "ZINCRBY",
    "DEL", "UNLINK", "EXPIRE", "EXPIREAT", "PERSIST",
    "INCR", "DECR", "INCRBY", "DECRBY", "INCRBYFLOAT",
    "APPEND", "GETSET",
}

# Commands that are not supported for replay
UNSUPPORTED_COMMANDS = {
    "SUBSCRIBE", "UNSUBSCRIBE", "PSUBSCRIBE", "PUNSUBSCRIBE",
    "PUBLISH", "PUBSUB",
    "WATCH", "UNWATCH", "MULTI", "EXEC", "DISCARD",
    "BLPOP", "BRPOP", "BLMOVE", "BRPOPLPUSH",  # Blocking commands
    "WAIT", "WAITAOF",
    "CLIENT", "CONFIG", "DEBUG", "SHUTDOWN", "SLAVEOF", "REPLICAOF",
}


def _get_recorder() -> "Recorder":
    """Get the global recorder instance."""
    if _recorder is None:
        from autodebug_probe.recorder import get_recorder
        return get_recorder()
    return _recorder


def _serialize_arg(arg: Any) -> Any:
    """Serialize a single argument for recording."""
    if arg is None:
        return None
    if isinstance(arg, bytes):
        try:
            return arg.decode("utf-8")
        except UnicodeDecodeError:
            return {"__bytes__": arg.hex()}
    if isinstance(arg, (str, int, float, bool)):
        return arg
    if isinstance(arg, (list, tuple)):
        return [_serialize_arg(a) for a in arg]
    if isinstance(arg, dict):
        return {str(k): _serialize_arg(v) for k, v in arg.items()}
    return str(arg)


def _serialize_args(args: Tuple) -> List[Any]:
    """Serialize command arguments for recording."""
    return [_serialize_arg(arg) for arg in args]


def _serialize_result(result: Any) -> Any:
    """Serialize command result for recording."""
    if result is None:
        return None
    if isinstance(result, bytes):
        try:
            return result.decode("utf-8")
        except UnicodeDecodeError:
            return {"__bytes__": result.hex()}
    if isinstance(result, (str, int, float, bool)):
        return result
    if isinstance(result, (list, tuple)):
        return [_serialize_result(r) for r in result]
    if isinstance(result, dict):
        return {str(k): _serialize_result(v) for k, v in result.items()}
    return str(result)


def _extract_key_pattern(command: str, args: Tuple) -> Optional[str]:
    """Extract key pattern from command for matching."""
    command = command.upper()

    # Most commands have key as first argument
    if args and isinstance(args[0], (str, bytes)):
        key = args[0]
        if isinstance(key, bytes):
            try:
                key = key.decode("utf-8")
            except UnicodeDecodeError:
                key = key.hex()
        return key

    return None


def _hash_args(args: Tuple) -> str:
    """Hash arguments for canonicalization."""
    serialized = json.dumps(_serialize_args(args), sort_keys=True)
    return hashlib.sha256(serialized.encode()).hexdigest()[:16]


def _patched_execute_command(self, *args, **kwargs):
    """Patched Redis.execute_command that captures commands."""
    command = args[0].upper() if args else "UNKNOWN"
    cmd_args = args[1:] if len(args) > 1 else ()

    recorder = _get_recorder()

    # Check for unsupported commands
    if command in UNSUPPORTED_COMMANDS:
        recorder.record_event({
            "type": "REDIS_UNSUPPORTED",
            "command": command,
            "message": f"Redis command {command} is not supported for replay",
        })

    # Execute the command
    try:
        result = _original_execute_command(self, *args, **kwargs)
    except Exception as exc:
        # Record the command with exception details
        key_pattern = _extract_key_pattern(command, cmd_args)
        recorder.record_redis({
            "command": command,
            "args": _serialize_args(cmd_args),
            "args_hash": _hash_args(cmd_args),
            "result": None,
            "key": key_pattern,
            "is_read": command in READ_COMMANDS,
            "is_write": command in WRITE_COMMANDS,
            "exception": {
                "type": exc.__class__.__name__,
                "message": str(exc),
            },
        })

        # Record an EXCEPTION event with stack trace
        recorder.record_event({
            "type": "EXCEPTION",
            "command": command,
            "key": key_pattern,
            "exception_type": exc.__class__.__name__,
            "exception_message": str(exc),
            "traceback": traceback.format_exc(),
        })

        # Re-raise the original exception unchanged
        raise

    # Record the command
    key_pattern = _extract_key_pattern(command, cmd_args)

    payload = {
        "command": command,
        "args": _serialize_args(cmd_args),
        "args_hash": _hash_args(cmd_args),
        "result": _serialize_result(result),
        "key": key_pattern,
        "is_read": command in READ_COMMANDS,
        "is_write": command in WRITE_COMMANDS,
    }

    recorder.record_redis(payload)

    return result


def _patched_pipeline_execute(self, raise_on_error=True):
    """Patched Pipeline.execute that captures pipeline commands."""
    recorder = _get_recorder()

    # Get commands from pipeline
    commands = []
    if hasattr(self, "command_stack"):
        for cmd in self.command_stack:
            if hasattr(cmd, "args"):
                command = cmd.args[0].upper() if cmd.args else "UNKNOWN"
                cmd_args = cmd.args[1:] if len(cmd.args) > 1 else ()
                commands.append({
                    "command": command,
                    "args": _serialize_args(cmd_args),
                    "key": _extract_key_pattern(command, cmd_args),
                })

    # Execute the pipeline
    results = _original_pipeline_execute(self, raise_on_error)

    # Record the pipeline execution
    recorder.record_redis({
        "command": "PIPELINE",
        "commands": commands,
        "results": [_serialize_result(r) for r in results] if results else [],
        "is_pipeline": True,
    })

    return results


def install() -> None:
    """Install Redis capture patches."""
    global _original_execute_command, _original_pipeline_execute

    try:
        import redis
    except ImportError:
        return

    if _original_execute_command is not None:
        return  # Already patched

    # Patch Redis.execute_command
    _original_execute_command = redis.Redis.execute_command
    redis.Redis.execute_command = _patched_execute_command

    # Patch Pipeline.execute
    if hasattr(redis, "client") and hasattr(redis.client, "Pipeline"):
        _original_pipeline_execute = redis.client.Pipeline.execute
        redis.client.Pipeline.execute = _patched_pipeline_execute


def uninstall() -> None:
    """Restore original Redis functions."""
    global _original_execute_command, _original_pipeline_execute

    try:
        import redis
    except ImportError:
        return

    if _original_execute_command is not None:
        redis.Redis.execute_command = _original_execute_command
        _original_execute_command = None

    if _original_pipeline_execute is not None:
        if hasattr(redis, "client") and hasattr(redis.client, "Pipeline"):
            redis.client.Pipeline.execute = _original_pipeline_execute
        _original_pipeline_execute = None


def get_redis_info() -> Dict[str, Any]:
    """Get Redis detection info for manifest."""
    try:
        import redis
        return {
            "detected": True,
            "version": redis.__version__,
            "supported_commands": list(READ_COMMANDS | WRITE_COMMANDS),
            "unsupported_commands": list(UNSUPPORTED_COMMANDS),
        }
    except ImportError:
        return {"detected": False}
