# Production-ready step limit implementation con LangGraph
# Previene infinite loops como el caso $47k
from langgraph.graph import StateGraph, END
from typing import TypedDict, Annotated
from datetime import datetime, timedelta
import operator

# Definimos el estado del agente con contadores de seguridad
class AgentState(TypedDict):
    messages: Annotated[list, operator.add]  # Historial conversación
    step_count: int  # Contador de pasos ejecutados
    max_steps: int  # Límite hard de pasos
    start_time: datetime  # Timestamp inicio para timeout
    timeout_minutes: int  # Timeout en minutos
    task_completed: bool  # Flag de completación
    error_message: str  # Mensaje de error si ocurre

# Node de validación que previene infinite loops
def check_termination_criteria(state: AgentState) -> str:
    """
    Valida si el agente debe continuar o terminar.
    Implementa 3 checks críticos:
    1. Step limit excedido
    2. Timeout excedido
    3. Task completada
    """
    # Check #1: Step limit
    if state["step_count"] >= state["max_steps"]:
        state["error_message"] = f"TERMINATED: Exceeded max steps ({state['max_steps']})"
        print(f"⚠️ {state['error_message']}")
        return "TERMINATE"

    # Check #2: Timeout
    elapsed = datetime.now() - state["start_time"]
    if elapsed > timedelta(minutes=state["timeout_minutes"]):
        state["error_message"] = f"TERMINATED: Timeout after {state['timeout_minutes']} minutes"
        print(f"⚠️ {state['error_message']}")
        return "TERMINATE"

    # Check #3: Task completed
    if state["task_completed"]:
        print(f"✅ Task completed successfully in {state['step_count']} steps")
        return "SUCCESS"

    # Continuar ejecución
    return "CONTINUE"

# Node de agente que documenta planning (Magentic pattern)
def agent_with_planning(state: AgentState):
    """
    Agente que PRIMERO planifica, LUEGO ejecuta.
    Esto permite validar el plan antes de ejecución.
    """
    # Incrementar contador ANTES de ejecutar
    state["step_count"] += 1
    print(f"\n🤖 Step {state['step_count']}/{state['max_steps']}")

    # PASO 1: Documentar plan (Magentic pattern)
    plan_prompt = f"""
    Current task: {state['messages'][-1]}
    Steps completed so far: {state['step_count'] - 1}
    
    Before executing, document your plan:
    1. What is the next action?
    2. Why is this action necessary?
    3. What is the expected outcome?
    4. How will you know if it succeeded?
    """

    # Aquí llamarías a tu LLM para generar el plan
    # plan = llm.invoke(plan_prompt)
    # print(f"📋 Plan: {plan}")

    # PASO 2: Validar plan (podrías usar otro LLM call o reglas)
    # if not validate_plan(plan):
    #     return {"error": "Invalid plan detected"}

    # PASO 3: Ejecutar acción
    # result = execute_action(plan)

    # PASO 4: Validar resultado
    # state["task_completed"] = validate_result(result)

    # Simulación para demo
    state["task_completed"] = state["step_count"] >= 5  # Completa en 5 pasos

    return state

# Construir el workflow con safety guards
workflow = StateGraph(AgentState)

# Añadir nodes
workflow.add_node("agent", agent_with_planning)

# Añadir conditional edges (routing basado en termination criteria)
workflow.add_conditional_edges(
    "agent",
    check_termination_criteria,
    {
        "CONTINUE": "agent",  # Loop back si debe continuar
        "SUCCESS": END,  # Terminar si completado
        "TERMINATE": END  # Terminar si excedió límites
    }
)

# Entry point
workflow.set_entry_point("agent")

# Compilar grafo
app = workflow.compile()

# EJEMPLO DE USO
if __name__ == "__main__":
    # Configuración inicial con safety limits
    initial_state = {
        "messages": ["Research market trends for autonomous AI agents in 2025"],
        "step_count": 0,
        "max_steps": 10,  # CRÍTICO: Límite hard de pasos
        "start_time": datetime.now(),
        "timeout_minutes": 5,  # CRÍTICO: Timeout de 5 minutos
        "task_completed": False,
        "error_message": ""
    }

    # Ejecutar agente con protecciones
    print("🚀 Starting agent with safety protections...")
    result = app.invoke(initial_state)

    print(f"\n{'='*60}")
    print(f"Final result:")
    print(f"  Steps executed: {result['step_count']}")
    print(f"  Task completed: {result['task_completed']}")
    print(f"  Error: {result['error_message']}")
    print(f"{'='*60}")

    # RESULTADO ESPERADO:
    # - Agente ejecuta máximo 10 steps O 5 minutos (lo que ocurra primero)
    # - Si completa task antes, termina exitosamente
    # - Si excede límites, termina con error message
    # - IMPOSSIBLE tener infinite loop como el caso $47k 

# Production-grade RAG con anti-hallucination de 4 capas
# Reduce hallucinations de 15-20% a 

# ============================================
# BEFORE: Legacy LangChain (verbose, unmaintainable)
# ============================================
from langchain.agents import AgentExecutor, initialize_agent, AgentType
from langchain.chains import LLMChain
from langchain.memory import ConversationBufferMemory
from langchain.callbacks import StdOutCallbackHandler
from langchain.tools import Tool

# 1. Define tools (verboso)
tool1 = Tool(
    name="search",
    func=lambda x: f"Search results for {x}",
    description="Search the web"
)
tool2 = Tool(
    name="calculator",
    func=lambda x: eval(x),
    description="Calculate mathematical expressions"
)
tool3 = Tool(
    name="database",
    func=lambda x: f"Database query: {x}",
    description="Query internal database"
)

# 2. Setup memory (custom config)
memory = ConversationBufferMemory(
    memory_key="chat_history",
    return_messages=True,
    output_key="output"
)

# 3. Initialize agent (muchas opciones, confusing)
agent = initialize_agent(
    tools=[tool1, tool2, tool3],
    llm=llm,
    agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION,  # ¿Cuál agent type usar?
    verbose=True,
    max_iterations=10,  # Manual step limit
    early_stopping_method="generate",  # ¿Qué hace esto?
    memory=memory,
    handle_parsing_errors=True,  # Handling errors oscuro
    callbacks=[StdOutCallbackHandler()],  # Custom logging
    # ... 10+ más opciones no documentadas
)

# 4. Execute (sin state visibility)
try:
    result = agent.run("Find information about AI agents and calculate 2+2")
    print(result)
except Exception as e:
    # Error handling manual, no structured logging
    print(f"Agent failed: {e}")
    # ¿Cómo hacer retry? ¿Cómo guardar state parcial?

# PROBLEMAS:
# - 50+ líneas solo setup
# - Opciones confusas (AgentType, early_stopping_method)
# - No state persistence built-in
# - Error handling manual
# - Debugging imposible (no execution replay)
# - Cyclic workflows = hacks con callbacks


# ============================================
# AFTER: Modern LangGraph (concise, maintainable)
# ============================================
from langgraph.prebuilt import create_react_agent
from langgraph.checkpoint.memory import MemorySaver

# 1. Define tools (mismo, más conciso)
tools = [
    Tool(name="search", func=lambda x: f"Search: {x}", description="Search"),
    Tool(name="calculator", func=lambda x: eval(x), description="Calculate"),
    Tool(name="database", func=lambda x: f"DB: {x}", description="Query DB")
]

# 2. Create agent (1 línea, sane defaults)
agent = create_react_agent(
    model=llm,
    tools=tools,
    checkpointer=MemorySaver()  # Built-in state persistence!
)

# 3. Execute con thread_id para state tracking
config = {"configurable": {"thread_id": "user-123"}}

try:
    # State management automático
    result = agent.invoke(
        {"messages": [("user", "Find info about AI agents and calculate 2+2")]},
        config=config
    )
    print(result["messages"][-1].content)

    # Continuar conversación (state persisted)
    result2 = agent.invoke(
        {"messages": [("user", "What was the calculation result?")]},
        config=config  # MISMO thread_id = recuerda contexto
    )
    print(result2["messages"][-1].content)

except Exception as e:
    # Error handling structured con state snapshot
    print(f"Agent failed at step: {agent.get_state(config)}")

# BENEFICIOS:
# - 90% menos código (50+ → 5 líneas)
# - State persistence built-in (checkpointer)
# - Execution replay posible (get_state)
# - Cyclic workflows nativos (add_edge)
# - Debugging trivial (LangSmith tracing automático)
# - Maintainability 10x mejor (team onboarding 

# Production Orchestrator-Worker pattern con LangGraph
# 1 orchestrator coordina 3 workers especializados
from langgraph.graph import StateGraph, END
from typing import TypedDict, Literal, Annotated
import operator


# State schema compartido
class OrchestratorState(TypedDict):
    messages: Annotated[list, operator.add]
    task_description: str
    assigned_worker: str  # "search" | "calculator" | "database" | None
    worker_results: dict  # {worker_name: result}
    final_response: str
    error: str


# === ORCHESTRATOR NODE ===
def orchestrator_node(state: OrchestratorState) -> OrchestratorState:
    """
    Central coordinator. Analiza task y decide qué worker asignar.
    NUNCA ejecuta tasks directamente—solo routing.
    """
    task = state["task_description"]

    # Routing logic (podría ser LLM-based para más inteligencia)
    if "search" in task.lower() or "find" in task.lower():
        state["assigned_worker"] = "search"
    elif "calculate" in task.lower() or any(op in task for op in ["+", "-", "*", "/"]):
        state["assigned_worker"] = "calculator"
    elif "database" in task.lower() or "query" in task.lower():
        state["assigned_worker"] = "database"
    else:
        # Fallback: usa LLM para decidir
        state["assigned_worker"] = "search"  # Default

    print(f"🎯 Orchestrator assigned task to: {state['assigned_worker']}")
    return state


# === WORKER NODES ===
def search_worker(state: OrchestratorState) -> OrchestratorState:
    """Worker especializado en búsquedas web."""
    print(f"🔍 Search worker executing...")

    # Aquí llamarías a Tavily API, Google Search, etc
    result = f"Search results for: {state['task_description']}"
    state["worker_results"]["search"] = result

    print(f"✅ Search completed")
    return state


def calculator_worker(state: OrchestratorState) -> OrchestratorState:
    """Worker especializado en cálculos."""
    print(f"🔢 Calculator worker executing...")

    # Aquí harías cálculos reales (eval con safety checks)
    try:
        # NOTA: eval() es inseguro—usa ast.literal_eval o math parser
        result = f"Calculation result (simulated)"
        state["worker_results"]["calculator"] = result
        print(f"✅ Calculation completed")
    except Exception as e:
        state["error"] = f"Calculator error: {e}"

    return state


def database_worker(state: OrchestratorState) -> OrchestratorState:
    """Worker especializado en database queries."""
    print(f"💾 Database worker executing...")

    # Aquí ejecutarías SQL queries, MongoDB queries, etc
    result = f"Database query results (simulated)"
    state["worker_results"]["database"] = result

    print(f"✅ Database query completed")
    return state


# === AGGREGATOR NODE ===
def aggregator_node(state: OrchestratorState) -> OrchestratorState:
    """
    Recibe results de workers y genera final response.
    Podría usar LLM para sintetizar results.
    """
    print(f"📊 Aggregating results from workers...")

    # Combina worker results
    all_results = "\n".join([
        f"{worker}: {result}"
        for worker, result in state["worker_results"].items()
    ])

    # Genera final response (aquí usarías LLM en producción)
    state["final_response"] = f"Task completed. Results:\n{all_results}"

    print(f"✅ Aggregation completed")
    return state


# === ROUTING FUNCTION ===
def route_to_worker(state: OrchestratorState) -> Literal["search", "calculator", "database"]:
    """
    Conditional edge routing basado en orchestrator decision.
    """
    return state["assigned_worker"]


# === BUILD GRAPH ===
workflow = StateGraph(OrchestratorState)

# Add nodes
workflow.add_node("orchestrator", orchestrator_node)
workflow.add_node("search", search_worker)
workflow.add_node("calculator", calculator_worker)
workflow.add_node("database", database_worker)
workflow.add_node("aggregator", aggregator_node)

# Define edges
workflow.set_entry_point("orchestrator")

# Conditional routing a workers
workflow.add_conditional_edges(
    "orchestrator",
    route_to_worker,
    {
        "search": "search",
        "calculator": "calculator",
        "database": "database"
    }
)

# All workers converge to aggregator
workflow.add_edge("search", "aggregator")
workflow.add_edge("calculator", "aggregator")
workflow.add_edge("database", "aggregator")

# Aggregator termina
workflow.add_edge("aggregator", END)

# Compile
app = workflow.compile()


# === EJEMPLO DE USO ===
if __name__ == "__main__":
    # Task 1: Search task
    result1 = app.invoke({
        "messages": [],
        "task_description": "Find information about LangGraph framework",
        "assigned_worker": "",
        "worker_results": {},
        "final_response": "",
        "error": ""
    })

    print(f"\n{'='*60}")
    print(f"Task 1 Result:\n{result1['final_response']}")
    print(f"{'='*60}\n")

    # Task 2: Calculation task
    result2 = app.invoke({
        "messages": [],
        "task_description": "Calculate the sum of 123 + 456",
        "assigned_worker": "",
        "worker_results": {},
        "final_response": "",
        "error": ""
    })

    print(f"\n{'='*60}")
    print(f"Task 2 Result:\n{result2['final_response']}")
    print(f"{'='*60}\n")


# BENEFICIOS DE ESTE PATTERN:
# ✅ Clear ownership: cada worker tiene responsibility específica
# ✅ Easy debugging: logs muestran exactamente qué worker ejecutó
# ✅ Parallel-ready: workers NO se bloquean entre sí
# ✅ Error isolation: si Calculator falla, Search/Database siguen OK
# ✅ Maintainable: agregar nuevo worker = add_node + routing logic
# ✅ Testable: cada worker se puede unit test independientemente

# COMPARISON VS AD-HOC:
# - Ad-hoc: 3 agents llamándose directamente = spaghetti
# - Orchestrator: 1 coordinator + 3 workers = clean architecture

# Production error handling: 3 layers (Retry + Circuit Breaker + Fallback)
# Reduce failure rate de 10-15% a 

# Production cost optimization framework
# Reduce costes 40-70% sin sacrificar quality
import tiktoken
from typing import Literal
from datetime import datetime, timedelta


# ===== CAPA #1: HARD BUDGET LIMITS =====
class BudgetGuard:
    """
    Previene cost explosions con hard limits.
    """

    def __init__(
        self,
        max_tokens_per_request: int = 10000,
        max_cost_per_day: float = 500.0,  # En tu moneda
        max_requests_per_user_per_hour: int = 100
    ):
        self.max_tokens_per_request = max_tokens_per_request
        self.max_cost_per_day = max_cost_per_day
        self.max_requests_per_user_per_hour = max_requests_per_user_per_hour

        # Tracking state
        self.daily_cost = 0.0
        self.last_reset = datetime.now()
        self.user_request_counts = {}  # {user_id: [(timestamp, count), ...]}

    def check_token_budget(self, prompt: str, model: str = "gpt-4") -> bool:
        """
        Verifica si request excede token budget.
        """
        # Count tokens usando tiktoken
        encoding = tiktoken.encoding_for_model(model)
        token_count = len(encoding.encode(prompt))

        if token_count > self.max_tokens_per_request:
            raise Exception(
                f"Request rejected: {token_count} tokens exceeds limit of {self.max_tokens_per_request}"
            )

        return True

    def check_daily_budget(self, estimated_cost: float) -> bool:
        """
        Verifica si request excedería daily budget.
        """
        # Reset counter si es nuevo día
        if datetime.now().date() > self.last_reset.date():
            self.daily_cost = 0.0
            self.last_reset = datetime.now()

        if self.daily_cost + estimated_cost > self.max_cost_per_day:
            raise Exception(
                f"Daily budget exceeded: ${self.daily_cost:.2f} + ${estimated_cost:.2f} > ${self.max_cost_per_day}"
            )

        return True

    def check_rate_limit(self, user_id: str) -> bool:
        """
        Verifica rate limit per-user.
        """
        now = datetime.now()
        one_hour_ago = now - timedelta(hours=1)

        # Clean old entries
        if user_id in self.user_request_counts:
            self.user_request_counts[user_id] = [
                (ts, count) for ts, count in self.user_request_counts[user_id]
                if ts > one_hour_ago
            ]
        else:
            self.user_request_counts[user_id] = []

        # Count requests in last hour
        recent_count = sum(count for _, count in self.user_request_counts[user_id])

        if recent_count >= self.max_requests_per_user_per_hour:
            raise Exception(
                f"Rate limit exceeded: {recent_count} requests in last hour (max {self.max_requests_per_user_per_hour})"
            )

        # Increment counter
        self.user_request_counts[user_id].append((now, 1))
        return True


# ===== CAPA #2: SMART MODEL ROUTING =====
class SmartModelRouter:
    """
    Rutas requests al modelo más cost-effective.
    """

    # Pricing aproximado (verificar pricing real)
    MODEL_COSTS = {
        "gpt-4": {"input": 0.00003, "output": 0.00006},  # $/token
        "gpt-4-turbo": {"input": 0.00001, "output": 0.00003},
        "gpt-3.5-turbo": {"input": 0.000001, "output": 0.000002},
        "claude-opus": {"input": 0.000015, "output": 0.000075},
        "claude-sonnet": {"input": 0.000003, "output": 0.000015}
    }

    def classify_complexity(self, prompt: str) -> Literal["simple", "medium", "complex"]:
        """
        Clasifica complexity de task.
        En producción: usa ML classifier o LLM-based classifier.
        """
        # Heurísticas simples (mejora con ML)
        if len(prompt) < 100:
            return "simple"
        elif "analyze" in prompt.lower() or "complex" in prompt.lower():
            return "complex"
        else:
            return "medium"

    def select_model(self, prompt: str, max_cost: float = None) -> str:
        """
        Selecciona modelo óptimo basado en complexity + cost constraints.
        """
        complexity = self.classify_complexity(prompt)

        # Routing basado en complexity
        if complexity == "simple":
            return "gpt-3.5-turbo"  # 5x más barato que GPT-4
        elif complexity == "medium":
            return "gpt-4-turbo"  # Balance cost/quality
        else:
            # Complex: usa best model
            return "gpt-4"

    def estimate_cost(self, prompt: str, model: str, output_tokens: int = 500) -> float:
        """
        Estima cost de request.
        """
        encoding = tiktoken.encoding_for_model(model)
        input_tokens = len(encoding.encode(prompt))

        costs = self.MODEL_COSTS.get(model, self.MODEL_COSTS["gpt-4"])  # Default expensive

        estimated_cost = (
            input_tokens * costs["input"] +
            output_tokens * costs["output"]
        )

        return estimated_cost


# ===== CAPA #3: SEMANTIC CACHING =====
class SemanticCache:
    """
    Cache responses basado en semantic similarity.
    """

    def __init__(self, similarity_threshold: float = 0.95):
        self.cache = {}  # {query_hash: (response, timestamp, embedding)}
        self.similarity_threshold = similarity_threshold

    def get(self, query: str) -> str | None:
        """
        Busca en cache si hay query similar.
        """
        # Simplified: exact match. En producción usa embeddings + cosine similarity.
        query_normalized = query.lower().strip()

        if query_normalized in self.cache:
            cached_response, timestamp, _ = self.cache[query_normalized]
            print(f"✅ CACHE HIT (saved API call)")
            return cached_response

        # TODO: Implement semantic similarity search con embeddings
        return None

    def set(self, query: str, response: str):
        """
        Guarda en cache.
        """
        query_normalized = query.lower().strip()
        self.cache[query_normalized] = (response, datetime.now(), None)


# ===== AGENT CON COST OPTIMIZATION COMPLETO =====
class CostOptimizedAgent:
    """
    Agent con todas las capas de cost optimization.
    """

    def __init__(self, openai_api_key: str):
        import openai
        self.client = openai.OpenAI(api_key=openai_api_key)

        # Initialize cost optimization layers
        self.budget_guard = BudgetGuard(
            max_tokens_per_request=10000,
            max_cost_per_day=500.0,
            max_requests_per_user_per_hour=100
        )
        self.router = SmartModelRouter()
        self.cache = SemanticCache()

    def invoke(self, prompt: str, user_id: str = "default") -> dict:
        """
        Invoke con FULL cost optimization.
        """
        # CAPA #3: Check cache primero (fastest + cheapest)
        cached = self.cache.get(prompt)
        if cached:
            return {
                "response": cached,
                "cached": True,
                "cost": 0.0,
                "model": "cache"
            }

        # CAPA #2: Smart model routing
        selected_model = self.router.select_model(prompt)
        print(f"📊 Selected model: {selected_model}")

        # Estimate cost
        estimated_cost = self.router.estimate_cost(prompt, selected_model)
        print(f"💰 Estimated cost: ${estimated_cost:.4f}")

        # CAPA #1: Budget guards
        try:
            self.budget_guard.check_token_budget(prompt, selected_model)
            self.budget_guard.check_daily_budget(estimated_cost)
            self.budget_guard.check_rate_limit(user_id)
        except Exception as e:
            print(f"🚫 Budget guard REJECTED request: {e}")
            raise

        # Execute request
        response = self.client.chat.completions.create(
            model=selected_model,
            messages=[{"role": "user", "content": prompt}]
        )

        result_text = response.choices[0].message.content

        # Update tracking
        actual_cost = self.router.estimate_cost(
            prompt,
            selected_model,
            output_tokens=response.usage.completion_tokens
        )
        self.budget_guard.daily_cost += actual_cost

        # Cache result
        self.cache.set(prompt, result_text)

        return {
            "response": result_text,
            "cached": False,
            "cost": actual_cost,
            "model": selected_model,
            "tokens_used": response.usage.total_tokens
        }


# ===== EJEMPLO DE USO =====
if __name__ == "__main__":
    agent = CostOptimizedAgent(openai_api_key="tu-api-key")

    # Test 1: Simple query (should use GPT-3.5)
    result1 = agent.invoke("What is 2+2?", user_id="user123")
    print(f"\nResult 1: {result1}\n")

    # Test 2: Same query (should hit cache)
    result2 = agent.invoke("What is 2+2?", user_id="user123")
    print(f"\nResult 2 (cached): {result2}\n")

    # Test 3: Complex query (should use GPT-4)
    result3 = agent.invoke(
        "Analyze the geopolitical implications of...",
        user_id="user123"
    )
    print(f"\nResult 3: {result3}\n")


# RESULTADOS ESPERADOS:
# - 40-60% queries hit cache (costo $0)
# - 80% queries restantes usan GPT-3.5 (5x más barato)
# - Solo 10-20% queries usan GPT-4 (cuando realmente necesario)
# - TOTAL COST REDUCTION: 40-70% vs usar GPT-4 para todo
# - NO más $47k surprises (daily budget ceiling previene)