import time
from functools import wraps
from typing import Callable, Dict, Any
import redis

# Configuración
REDIS_CLIENT = redis.Redis(host='localhost', port=6379, db=0)
DAILY_BUDGET_USD = 100  # Presupuesto diario máximo
COST_PER_1K_INPUT_TOKENS = 0.01  # GPT-4o pricing
COST_PER_1K_OUTPUT_TOKENS = 0.03
MAX_TOKENS_PER_REQUEST = 8000  # Límite de seguridad
MAX_REQUESTS_PER_MINUTE = 10  # Rate limiting


def cost_monitor(func: Callable) -> Callable:
    """
    Decorator que monitorea costes de LLM calls y aplica circuit breakers.
    Tracking en Redis para persistencia cross-process.
    """
    @wraps(func)
    def wrapper(*args, **kwargs) -> Any:
        user_id = kwargs.get('user_id', 'default')

        # 1. CHECK RATE LIMITING
        rate_key = f"rate_limit:{user_id}:{int(time.time() / 60)}"
        current_requests = REDIS_CLIENT.incr(rate_key)
        REDIS_CLIENT.expire(rate_key, 60)  # TTL 1 minuto

        if current_requests > MAX_REQUESTS_PER_MINUTE:
            raise Exception(
                f"⚠️ Rate limit exceeded: {current_requests} requests/min"
            )

        # 2. CHECK DAILY BUDGET
        today = time.strftime("%Y-%m-%d")
        budget_key = f"daily_cost:{today}"
        current_cost = float(REDIS_CLIENT.get(budget_key) or 0)

        if current_cost >= DAILY_BUDGET_USD:
            raise Exception(
                f"🚨 DAILY BUDGET EXCEEDED: ${current_cost:.2f} / ${DAILY_BUDGET_USD}"
            )

        # 3. ESTIMATE COST BEFORE EXECUTION
        estimated_input_tokens = len(str(kwargs.get('prompt', ''))) / 4  # Aproximación rough

        if estimated_input_tokens > MAX_TOKENS_PER_REQUEST:
            raise Exception(
                f"⚠️ Request too large: ~{estimated_input_tokens} tokens"
            )

        # 4. EXECUTE FUNCTION
        start_time = time.time()
        result = func(*args, **kwargs)
        execution_time = time.time() - start_time

        # 5. TRACK ACTUAL COST
        actual_input_tokens = result.get('usage', {}).get('prompt_tokens', 0)
        actual_output_tokens = result.get('usage', {}).get('completion_tokens', 0)

        cost = (
            actual_input_tokens / 1000 * COST_PER_1K_INPUT_TOKENS +
            actual_output_tokens / 1000 * COST_PER_1K_OUTPUT_TOKENS
        )

        # Update daily cost
        REDIS_CLIENT.incrbyfloat(budget_key, cost)
        REDIS_CLIENT.expire(budget_key, 86400 * 7)  # Keep 7 days

        # Log para observability
        print(
            f"💰 Cost: ${cost:.4f} | "
            f"Tokens: {actual_input_tokens}in/{actual_output_tokens}out | "
            f"Time: {execution_time:.2f}s"
        )

        return result

    return wrapper


# USO
@cost_monitor
def call_llm(prompt: str, user_id: str, **kwargs) -> Dict[str, Any]:
    """Tu función que llama a OpenAI/Anthropic/etc"""
    # Implementación real aquí
    pass
 

from collections import defaultdict
from typing import List, Tuple


class DeadlockDetector:
    """
    Detecta loops entre agents comparando conversation history.
    Si mismos 2 agentes intercambian >N mensajes sin progreso → DEADLOCK.
    """

    def __init__(self, max_exchanges: int = 5):
        self.max_exchanges = max_exchanges
        self.conversation_history: List[Tuple[str, str]] = []  # (sender, receiver)

    def record_message(self, sender: str, receiver: str) -> bool:
        """
        Registra mensaje entre agentes.
        Retorna True si deadlock detectado.
        """
        self.conversation_history.append((sender, receiver))

        # Contar exchanges consecutivos entre mismos 2 agentes
        if len(self.conversation_history) < self.max_exchanges:
            return False

        recent = self.conversation_history[-self.max_exchanges:]

        # Check si todos los mensajes son entre mismos 2 agents
        agents = set()
        for sender, receiver in recent:
            agents.add(sender)
            agents.add(receiver)

        if len(agents) == 2:
            print(
                f"🚨 DEADLOCK DETECTED: {agents} intercambiaron "
                f"{self.max_exchanges} mensajes sin progreso"
            )
            return True

        return False

    def reset(self):
        """Limpia history después de resolver deadlock"""
        self.conversation_history = []


# USO EN LANGGRAPH
detector = DeadlockDetector(max_exchanges=5)


def agent_step(state, sender, receiver):
    """Ejecuta un paso de agente con detección de deadlock"""
    # Check deadlock ANTES de ejecutar
    if detector.record_message(sender, receiver):
        # Estrategia de escape: involucrar supervisor o human-in-the-loop
        return {
            "action": "escalate_to_supervisor",
            "reason": "deadlock_detected"
        }

    # Proceder con ejecución normal
    return execute_agent_logic(state)