"""
Module de sécurité pour le SDK STRIGIVAL.

Contient les validations, détecteurs de divergence et protections
contre les attaques identifiées dans l'audit Full-Stack Web3.
"""

import re
import hashlib
import json
import time
import logging
from dataclasses import dataclass, field
from typing import Dict, Any, List, Optional, Set
from enum import Enum

from .exceptions import ValidationError

logger = logging.getLogger(__name__)


# =============================================================================
# CONSTANTS - ALLOWED PROTOCOLS & ACTIONS
# =============================================================================

class AllowedProtocol(Enum):
    """Protocoles DeFi autorisés."""
    UNISWAP_V3 = "uniswap_v3"
    UNISWAP_V2 = "uniswap_v2"
    AAVE_V3 = "aave_v3"
    AAVE_V2 = "aave_v2"
    COMPOUND_V3 = "compound_v3"
    ONE_INCH = "1inch"
    GMX = "gmx"
    CURVE = "curve"
    BALANCER = "balancer"


class AllowedAction(Enum):
    """Actions de trading autorisées."""
    SWAP = "swap"
    SUPPLY = "supply"
    BORROW = "borrow"
    REPAY = "repay"
    WITHDRAW = "withdraw"
    ADD_LIQUIDITY = "add_liquidity"
    REMOVE_LIQUIDITY = "remove_liquidity"


# Sélecteurs de fonctions connus et autorisés (4 premiers bytes du keccak256)
KNOWN_FUNCTION_SELECTORS: Dict[str, str] = {
    # Uniswap V3
    "0x414bf389": "exactInputSingle",
    "0xc04b8d59": "exactInput",
    "0x5023b4df": "exactOutputSingle",
    "0xf28c0498": "exactOutput",
    # Uniswap V2
    "0x38ed1739": "swapExactTokensForTokens",
    "0x8803dbee": "swapTokensForExactTokens",
    # Aave V3
    "0x617ba037": "supply",
    "0x69328dec": "withdraw",
    "0xa415bcad": "borrow",
    "0x573ade81": "repay",
    # Compound V3
    "0xf2b9fdb8": "supply",
    "0xf3fef3a3": "withdraw",
    # 1inch
    "0x12aa3caf": "swap",
    "0xe449022e": "uniswapV3Swap",
    # Curve
    "0x3df02124": "exchange",
    "0xa6417ed6": "exchange_underlying",
    # Balancer
    "0x52bbbe29": "swap",
    "0x945bcec9": "batchSwap",
}

# Adresses de contrats de protocoles connus (mainnet)
KNOWN_PROTOCOL_ADDRESSES: Dict[str, Set[str]] = {
    "uniswap_v3": {
        "0xe592427a0aece92de3edee1f18e0157c05861564",  # SwapRouter
        "0x68b3465833fb72a70ecdf485e0e4c7bd8665fc45",  # SwapRouter02
    },
    "aave_v3": {
        "0x87870bca3f3fd6335c3f4ce8392d69350b4fa4e2",  # Pool
    },
    "1inch": {
        "0x1111111254eeb25477b68fb85ed929f73a960582",  # AggregationRouterV5
    },
}


# =============================================================================
# CALLDATA VALIDATOR
# =============================================================================

class CalldataValidator:
    """
    Validateur de calldata pour prévenir les injections malveillantes.
    
    SECURITY: Empêche un bot compromis d'envoyer des calldata arbitraires
    qui pourraient drainer le vault.
    """
    
    def __init__(self, strict_mode: bool = True):
        """
        Initialise le validateur.
        
        Args:
            strict_mode: Si True, rejette tout calldata non reconnu
        """
        self.strict_mode = strict_mode
        self.allowed_protocols = {p.value for p in AllowedProtocol}
        self.allowed_actions = {a.value for a in AllowedAction}
        self._custom_selectors: Dict[str, str] = {}
    
    def add_custom_selector(self, selector: str, function_name: str) -> None:
        """Ajoute un sélecteur de fonction personnalisé."""
        if not selector.startswith("0x") or len(selector) != 10:
            raise ValidationError(f"Invalid selector format: {selector}")
        self._custom_selectors[selector.lower()] = function_name
    
    def validate_protocol(self, protocol: str) -> None:
        """Valide que le protocole est autorisé."""
        if protocol.lower() not in self.allowed_protocols:
            raise ValidationError(
                f"Protocol not allowed: {protocol}. "
                f"Allowed: {', '.join(self.allowed_protocols)}",
                field="protocol"
            )
    
    def validate_action(self, action: str) -> None:
        """Valide que l'action est autorisée."""
        if action.lower() not in self.allowed_actions:
            raise ValidationError(
                f"Action not allowed: {action}. "
                f"Allowed: {', '.join(self.allowed_actions)}",
                field="action"
            )
    
    def validate_target(self, target: str, protocol: str) -> None:
        """
        Valide que l'adresse cible est un contrat de protocole connu.
        
        SECURITY: Empêche l'interaction avec des contrats arbitraires.
        """
        if not target or target == "0x":
            return  # Sera résolu par l'API
        
        target_lower = target.lower()
        
        # Vérifier le format d'adresse
        if not re.match(r"^0x[a-fA-F0-9]{40}$", target):
            raise ValidationError(
                f"Invalid target address format: {target}",
                field="target"
            )
        
        # En mode strict, vérifier que l'adresse est connue
        if self.strict_mode:
            protocol_lower = protocol.lower()
            if protocol_lower in KNOWN_PROTOCOL_ADDRESSES:
                known_addresses = KNOWN_PROTOCOL_ADDRESSES[protocol_lower]
                if target_lower not in known_addresses:
                    logger.warning(
                        f"[SECURITY] Unknown target address for {protocol}: {target}"
                    )
                    # Ne pas bloquer mais logger l'avertissement
    
    def validate_calldata(self, calldata: str) -> str:
        """
        Valide le calldata et retourne le nom de la fonction.
        
        SECURITY: Vérifie que le calldata commence par un sélecteur connu.
        """
        if not calldata:
            return "empty"
        
        if not calldata.startswith("0x"):
            raise ValidationError(
                "Calldata must start with 0x",
                field="calldata"
            )
        
        if len(calldata) < 10:
            raise ValidationError(
                "Calldata too short (minimum 10 chars for selector)",
                field="calldata"
            )
        
        selector = calldata[:10].lower()
        
        # Vérifier dans les sélecteurs connus
        if selector in KNOWN_FUNCTION_SELECTORS:
            return KNOWN_FUNCTION_SELECTORS[selector]
        
        # Vérifier dans les sélecteurs personnalisés
        if selector in self._custom_selectors:
            return self._custom_selectors[selector]
        
        # En mode strict, rejeter les sélecteurs inconnus
        if self.strict_mode:
            raise ValidationError(
                f"Unknown function selector: {selector}. "
                f"Add it via add_custom_selector() if legitimate.",
                field="calldata"
            )
        
        logger.warning(f"[SECURITY] Unknown function selector: {selector}")
        return "unknown"
    
    def validate_params(
        self,
        protocol: str,
        action: str,
        params: Dict[str, Any]
    ) -> None:
        """
        Valide les paramètres d'un trade.
        
        SECURITY: Vérifie que les paramètres ne contiennent pas
        de valeurs malveillantes.
        """
        # 1. Valider protocole et action
        self.validate_protocol(protocol)
        self.validate_action(action)
        
        # 2. Vérifier qu'il n'y a pas de "target" arbitraire non autorisé
        if "target" in params:
            self.validate_target(params["target"], protocol)
        
        # 3. Vérifier le calldata si présent
        if "calldata" in params:
            self.validate_calldata(params["calldata"])
        
        # 4. Vérifier les montants (pas de valeurs négatives ou overflow)
        for key in ["amount", "amountIn", "amountOut", "value"]:
            if key in params:
                amount = params[key]
                if isinstance(amount, (int, float)):
                    if amount < 0:
                        raise ValidationError(
                            f"Negative amount not allowed: {key}={amount}",
                            field=key
                        )
                    if amount > 2**256 - 1:
                        raise ValidationError(
                            f"Amount overflow: {key}={amount}",
                            field=key
                        )
        
        # 5. Vérifier les adresses de tokens
        for key in ["tokenIn", "tokenOut", "token", "asset"]:
            if key in params:
                token = params[key]
                if isinstance(token, str) and token.startswith("0x"):
                    if not re.match(r"^0x[a-fA-F0-9]{40}$", token):
                        raise ValidationError(
                            f"Invalid token address: {key}={token}",
                            field=key
                        )
        
        logger.debug(f"[SECURITY] Params validated: {protocol}/{action}")


# =============================================================================
# PROTOCOL DIVERGENCE DETECTOR
# =============================================================================

@dataclass
class ProtocolState:
    """État d'un protocole DeFi à un instant T."""
    protocol_id: str
    version: str
    parameters_hash: str
    last_checked: float
    is_valid: bool = True
    check_count: int = 0


class ProtocolDivergenceDetector:
    """
    Détecte les divergences entre l'état attendu et réel des protocoles.
    
    SECURITY: Empêche les trades basés sur des paramètres de protocole
    obsolètes (ex: changement de Interest Rate Model sur Aave).
    """
    
    def __init__(self, divergence_threshold_hours: float = 1.0):
        """
        Initialise le détecteur.
        
        Args:
            divergence_threshold_hours: Durée max avant re-vérification
        """
        self.protocol_states: Dict[str, ProtocolState] = {}
        self.divergence_threshold_hours = divergence_threshold_hours
        self._divergence_callbacks: List[callable] = []
    
    def register_protocol(
        self,
        protocol_id: str,
        version: str,
        parameters: Dict[str, Any]
    ) -> None:
        """Enregistre l'état initial d'un protocole."""
        params_hash = self._hash_parameters(parameters)
        self.protocol_states[protocol_id] = ProtocolState(
            protocol_id=protocol_id,
            version=version,
            parameters_hash=params_hash,
            last_checked=time.time(),
            is_valid=True,
            check_count=0
        )
        logger.info(f"[PROTOCOL] Registered {protocol_id} v{version}")
    
    def check_divergence(
        self,
        protocol_id: str,
        current_parameters: Dict[str, Any]
    ) -> bool:
        """
        Vérifie si le protocole a divergé de son état enregistré.
        
        Returns:
            True si divergence détectée, False sinon
        """
        if protocol_id not in self.protocol_states:
            logger.warning(f"[PROTOCOL] Unknown protocol: {protocol_id}")
            return False
        
        state = self.protocol_states[protocol_id]
        current_hash = self._hash_parameters(current_parameters)
        state.check_count += 1
        
        # Vérifier si les paramètres ont changé
        if current_hash != state.parameters_hash:
            logger.error(f"[PROTOCOL] Divergence detected for {protocol_id}")
            logger.error(f"  Expected hash: {state.parameters_hash[:16]}...")
            logger.error(f"  Current hash:  {current_hash[:16]}...")
            
            state.is_valid = False
            
            # Appeler les callbacks
            for callback in self._divergence_callbacks:
                try:
                    callback(protocol_id, state.parameters_hash, current_hash)
                except Exception as e:
                    logger.error(f"[PROTOCOL] Callback error: {e}")
            
            return True
        
        # Vérifier si le check est trop vieux
        age_hours = (time.time() - state.last_checked) / 3600
        if age_hours > self.divergence_threshold_hours:
            logger.warning(
                f"[PROTOCOL] State check too old for {protocol_id}: "
                f"{age_hours:.1f}h > {self.divergence_threshold_hours}h"
            )
            return True
        
        state.last_checked = time.time()
        return False
    
    def is_protocol_valid(self, protocol_id: str) -> bool:
        """Vérifie si un protocole est dans un état valide."""
        if protocol_id not in self.protocol_states:
            return True  # Inconnu = pas de vérification
        return self.protocol_states[protocol_id].is_valid
    
    def on_divergence(self, callback: callable) -> None:
        """Enregistre un callback appelé lors d'une divergence."""
        self._divergence_callbacks.append(callback)
    
    def update_protocol_state(
        self,
        protocol_id: str,
        parameters: Dict[str, Any]
    ) -> None:
        """Met à jour l'état d'un protocole après vérification manuelle."""
        if protocol_id not in self.protocol_states:
            return
        
        state = self.protocol_states[protocol_id]
        state.parameters_hash = self._hash_parameters(parameters)
        state.last_checked = time.time()
        state.is_valid = True
        
        logger.info(f"[PROTOCOL] Updated state for {protocol_id}")
    
    def _hash_parameters(self, parameters: Dict[str, Any]) -> str:
        """Génère un hash déterministe des paramètres."""
        sorted_params = dict(sorted(parameters.items()))
        params_str = json.dumps(sorted_params, sort_keys=True, default=str)
        return hashlib.sha256(params_str.encode()).hexdigest()


# =============================================================================
# RATE CONDITION DETECTOR
# =============================================================================

@dataclass
class OrderRecord:
    """Enregistrement d'un ordre pour détection de race conditions."""
    order_id: str
    timestamp: float
    action: str
    params_hash: str
    status: str = "pending"


class RaceConditionDetector:
    """
    Détecte les race conditions dans les ordres.
    
    SECURITY: Empêche un bot d'envoyer deux ordres contradictoires
    qui pourraient être traités dans le mauvais ordre.
    """
    
    def __init__(self, window_seconds: float = 5.0):
        """
        Initialise le détecteur.
        
        Args:
            window_seconds: Fenêtre de temps pour détecter les conflits
        """
        self.window_seconds = window_seconds
        self.pending_orders: Dict[str, OrderRecord] = {}
        self._order_history: List[OrderRecord] = []
    
    def check_conflict(
        self,
        vault_address: str,
        action: str,
        params: Dict[str, Any]
    ) -> Optional[str]:
        """
        Vérifie s'il y a un conflit avec un ordre en cours.
        
        Returns:
            ID de l'ordre en conflit, ou None
        """
        now = time.time()
        params_hash = hashlib.sha256(
            json.dumps(params, sort_keys=True, default=str).encode()
        ).hexdigest()[:16]
        
        # Nettoyer les vieux ordres
        self._cleanup_old_orders(now)
        
        # Vérifier les conflits
        for order_id, order in self.pending_orders.items():
            if order.status != "pending":
                continue
            
            # Même action dans la fenêtre de temps
            if order.action == action:
                age = now - order.timestamp
                if age < self.window_seconds:
                    logger.warning(
                        f"[RACE] Potential race condition detected: "
                        f"order {order_id} ({order.action}) is still pending"
                    )
                    return order_id
            
            # Actions contradictoires (ex: buy puis sell immédiat)
            if self._are_contradictory(order.action, action):
                age = now - order.timestamp
                if age < self.window_seconds:
                    logger.warning(
                        f"[RACE] Contradictory orders detected: "
                        f"{order.action} vs {action}"
                    )
                    return order_id
        
        return None
    
    def register_order(
        self,
        order_id: str,
        action: str,
        params: Dict[str, Any]
    ) -> None:
        """Enregistre un nouvel ordre."""
        params_hash = hashlib.sha256(
            json.dumps(params, sort_keys=True, default=str).encode()
        ).hexdigest()[:16]
        
        record = OrderRecord(
            order_id=order_id,
            timestamp=time.time(),
            action=action,
            params_hash=params_hash,
            status="pending"
        )
        
        self.pending_orders[order_id] = record
        self._order_history.append(record)
    
    def complete_order(self, order_id: str, status: str = "completed") -> None:
        """Marque un ordre comme terminé."""
        if order_id in self.pending_orders:
            self.pending_orders[order_id].status = status
            del self.pending_orders[order_id]
    
    def _cleanup_old_orders(self, now: float) -> None:
        """Nettoie les ordres trop vieux."""
        cutoff = now - self.window_seconds * 10
        to_remove = [
            oid for oid, order in self.pending_orders.items()
            if order.timestamp < cutoff
        ]
        for oid in to_remove:
            del self.pending_orders[oid]
    
    def _are_contradictory(self, action1: str, action2: str) -> bool:
        """Vérifie si deux actions sont contradictoires."""
        contradictions = {
            ("supply", "withdraw"),
            ("borrow", "repay"),
            ("add_liquidity", "remove_liquidity"),
        }
        pair = tuple(sorted([action1.lower(), action2.lower()]))
        return pair in contradictions


# =============================================================================
# GLOBAL VALIDATOR INSTANCE
# =============================================================================

# Instance globale pour usage dans le client
_global_validator: Optional[CalldataValidator] = None
_global_divergence_detector: Optional[ProtocolDivergenceDetector] = None
_global_race_detector: Optional[RaceConditionDetector] = None


def get_validator(strict_mode: bool = True) -> CalldataValidator:
    """Retourne le validateur global."""
    global _global_validator
    if _global_validator is None:
        _global_validator = CalldataValidator(strict_mode=strict_mode)
    return _global_validator


def get_divergence_detector() -> ProtocolDivergenceDetector:
    """Retourne le détecteur de divergence global."""
    global _global_divergence_detector
    if _global_divergence_detector is None:
        _global_divergence_detector = ProtocolDivergenceDetector()
    return _global_divergence_detector


def get_race_detector() -> RaceConditionDetector:
    """Retourne le détecteur de race conditions global."""
    global _global_race_detector
    if _global_race_detector is None:
        _global_race_detector = RaceConditionDetector()
    return _global_race_detector
