from typing import TypedDict, Optional, Any
from dataclasses import dataclass
from datetime import datetime
import time

@dataclass
class AgentConstraints:
    max_cost_per_hour: float = 2.0  # Dollar amount that triggers emergency stop
    max_retries: int = 3  # Prevent infinite retry loops
    max_session_length: int = 3600  # Seconds before context poisoning
    required_context_keys: set = frozenset({"project_id", "user_id", "budget_limit"})

# This state type tracks what WILL break your system if you ignore it
class ProductionAgentState(TypedDict):
    goal: str
    context: dict
    current_task: str
    tools_used: list
    total_cost: float  # Running token/AI cost in dollars
    retry_count: int  # How many times we've tried
    last_check: str  # Timestamp for tracking drift
    meta_info: dict  # Business context (avoid audit hell)
    session_start: float  # When this started (for cost tracking)
    error: Optional[str]  # What went wrong (for debugging)
    status: str

def check_agent_safety(state: ProductionAgentState, constraints: AgentConstraints) -> dict:
    """
    This is your agent's guardian angel. EVERY action goes through this.

    Here's what this catches that demos ignore:

    1. Cost explosion: "We're burning $5/hour on tokens"
    2. Retry loops: Agent retrying itself into bankruptcy
    3. Session length: Context poisoning over long sessions
    4. Missing context: Agent running without required business context
    """
    current_cost = state["total_cost"]
    session_time = time.time() - state["session_start"]
    hourly_burn = (current_cost / session_time) * 3600 if session_time > 0 else 0

    # This check right here saves companies thousands in token bills
    if hourly_burn > constraints.max_cost_per_hour:
        return {
            "status": "critical",
            "message": f"🔥 COST EXPLOSION: Current rate ${hourly_burn:.2f}/hr exceeds ${constraints.max_cost_per_hour}",
            "next_action": "emergency_halt"
        }

    # This prevents the "infinite apology loop" - agent tries again, fails, apologizes, tries again...
    if state["retry_count"] > constraints.max_retries:
        return {
            "status": "critical",
            "message": f"STOP. It's tried {state['retry_count']} times. Human intervention required.",
            "next_action": "wake_engineer"
        }

    # This catches the context poison cycle before it kills production
    if session_time > constraints.max_session_length:
        return {
            "status": "warning",
            "message": f"Session at {session_time:.0f}s - risk of context degradation",
            "next_action": "checkpoint_and_refresh"
        }

    return {"status": "safe", "message": "🟢 Healthy", "next_action": "continue"}

from langchain_openai import ChatOpenAI
import logging

class ProductionAgent:
    def __init__(self, model_name="gpt-4o", safety_config=None):
        self.model = ChatOpenAI(model=model_name, temperature=0)
        self.safety = safety_config or get_production_safety_config()

        # This logging setup isn't decorative - it's how you debug when things break at 3 AM
        self.logger = setup_production_logging()

        # Track current execution state for safety checks
        self.current_state: ProductionAgentState = {
            "goal": "",
            "context": {},
            "current_task": "",
            "tools_used": [],
            "total_cost": 0.0,
            "retry_count": 0,
            "last_check": "",
            "meta_info": {},
            "session_start": time.time(),
            "error": None,
            "status": "initialized"
        }

    async def execute(self, goal: str, context: dict) -> ExecutionResult:
        """
        Here's the real art: understanding that agents fail in categories, not randomness.

        Category 1: Budget failures (we're spending more than the task is worth)
        Category 2: Quality failures (the answer is wrong/unhelpful)
        Category 3: System failures (the infrastructure is dying)
        Category 4: Context failures (we're asking the wrong thing)
        """

        # Create enriched context (not just slapping it in a prompt)
        enriched_context = self._enrich_context(goal, context)

        # Start billion-dollar checklist
        safety_check = self._pre_execution_safety_check(enriched_context)
        if safety_check.status != "safe":
            self.logger.critical(f"Agent would fail safety: {safety_check.message}")
            return self._handle_safety_failure(safety_check)

        # Only now do we let it use the compute budget
        plan = await self._create_robust_plan(goal, enriched_context)

        # This prevents the most common failure: context poisoning
        self._checkpoint_context_before_execution(plan)

        execution_result = await self._execute_with_production_escalation(plan)

        return self._post_execution_cleanup(execution_result)

    def _create_robust_plan(self, goal: str, context: dict) -> dict:
        """Here's the critical difference: we plan BEFORE we panic"""

        # Production planning includes failure strategies BY DEFAULT
        planning_prompt = f"""
        Create a production plan for: {goal}

        IMPORTANT: This plan must include:
        1. What to do if each step fails
        2. Maximum time budget per step
        3. How to roll back if we need to
        4. Success criteria for each step

        Business context: {context}
        Safety constraints: {self.safety}

        Remember: This will run without human supervision.
        """

        response = self.model.invoke(planning_prompt).content

        return {
            "plan": self._parse_production_plan(response),
            "rollback_commands": self._extract_rollback(response),
            "success_criteria": self._extract_criteria(response)
        }

    def _execute_with_production_escalation(self, plan: dict) -> dict:
        """
        The key insight: execution succeeds or fails in ESCALATION patterns.

        Pattern 1: Step succeeds immediately (< 30 seconds)
        Pattern 2: Requires retry with different approach (30s-2min)
        Pattern 3: Requires human review (2-10min)
        Pattern 4: Must halt immediately (emergency)
        """

        results = []
        final_status = "completed"

        for i, step in enumerate(plan["plan"]["steps"]):

            # Check business safety before EVERY step
            step_safety = check_agent_safety(self.current_state, self.safety)
            if step_safety["status"] != "safe":
                if step_safety["next_action"] == "emergency_halt":
                    final_status = "halted_for_safety"
                    break
                elif step_safety["next_action"] == "call_human":
                    final_status = "waiting_for_human"
                    break

            # Execute step with production patterns
            step_result = self._execute_step_with_patterns(step, plan, i)
            results.append(step_result)

            # Real-time decision making based on outcome
            if step_result["requires_action"]:
                self.logger.info(f"Step {i} escalated to manual review")
                final_status = "escalated_to_manual"
                break

        return {
            "final_status": final_status,
            "step_results": results,
            "business_impact": self._calculate_impact(results),
            "rollback_ready": plan["rollback"]
        }

    def _execute_step_with_patterns(self, step: dict, plan: dict, step_index: int) -> dict:
        """Critical insight: steps don't "fail" - they escalate through known patterns"""

        start_time = time.time()
        retry_count = 0
        max_retries = 2
        original_step = step.copy()

        while retry_count <= max_retries:
            try:
                # First attempt with original approach
                if retry_count == 0:
                    result = self._attempt_step_original_way(step)
                # Second attempt with alternatives
                elif retry_count == 1:
                    result = self._attempt_step_alternative_way(step)
                # Third attempt gets human help
                else:
                    return {
                        "status": "requires_manual_intervention",
                        "message": "Automated patterns exhausted, calling human",
                        "requires_action": True,
                        "evidence": self._collect_debug_info(["current_attempts"])
                    }

                # Did we succeed by current business definition?
                if self._step_meets_success_criteria(result, plan["success_criteria"][step_index]):
                    return {
                        "status": "success",
                        "result": result,
                        "time_seconds": time.time() - start_time,
                        "retries": retry_count
                    }

                # No success - log what went wrong for patterns
                self.logger.debug(f"Step {step_index}: attempt {retry_count} failed pattern {result['failure_pattern']}")
                retry_count += 1

            except Exception as e:
                # This captures unexpected failures (network, API, etc)
                self.logger.error(f"Unexpected failure in step {step_index}: {str(e)}")
                return {
                    "status": "unexpected_failure",
                    "message": str(e),
                    "requires_action": True,
                    "evidence": self._collect_debug_info(["exception", "step_trace"])
                }

        return {
            "status": "pattern_failure",
            "message": f"Step {step_index} failed all patterns",
            "requires_action": True,
            "evidence": self._log_failure_evidence(original_step, attempts={"all_attempts": retry_count})
        }

from typing import Dict, List, Optional
from dataclasses import dataclass
from functools import lru_cache
import time

@dataclass
class ResolutionOption:
    action: str
    impact: float
    risk: float
    reversibility: float

@dataclass
class Conflict:
    agent_a_priority: float
    agent_a_risk: float
    agent_a_reversible: float
    agent_b_priority: float
    agent_b_risk: float
    agent_b_reversible: float

class Resolution:
    pass

# Action constants for resolution options
agent_a_wins = "agent_a_wins"
agent_b_wins = "agent_b_wins"
both_lose_restart = "both_lose_restart"
escalate_to_human = "escalate_to_human"
timebox_experiment = "timebox_experiment"

class ProductionMultiAgentSystem:
    def __init__(self, max_agents: int = 5):
        self.agents: Dict[str, SafetyAwareAgent] = {}
        self.agent_timeout = 300  # 5 minutes max per agent
        self.conflict_resolver = AgentConflictResolver()

        # This is where we prevent the "agent civil war" disaster
        self.escalation_rules = {
            "agent_disagreement": "human_decision_required",
            "circular_dependency": "timeout_and_stalemate",
            "resource_conflict": "priority_order_with_fallback",
            "deadlock": "force_completion_with_logging"
        }

    def coordinate_agents(self, workflow: dict) -> dict:
        """
        The secret: coordinate by constraint satisfaction, not conversation.

        Instead of "Agent A, convince Agent B" we do:
        "Agent A, what's your constraint? Agent B, what's yours?"
        "Resolve mathematically who wins based on business rules"
        """

        agent_constraints = self._extract_all_agent_constraints(workflow)

        # If two agents want opposite things, resolve by rules, not negotiation
        conflicts = self._detect_agent_conflicts(agent_constraints)

        for conflict in conflicts:
            resolution = self.conflict_resolver.resolve_by_business_rules(conflict)

            # Document the decision for audit trail
            self._log_decision_making(conflict, resolution)

            # Update workflow based on resolution
            workflow = self._update_workflow_with_resolution(workflow, resolution)

        return workflow


class AgentConflictResolver:
    def resolve_by_business_rules(self, conflict: Conflict) -> Resolution:
        """
        Production secret: Legal mathematics > Agent intelligence.

        Don't ask agents who's right. Ask:
        - Which resolution has higher ROI?
        - Which has lower risk?
        - Which follows established policy?
        - What can we UNDO if wrong?
        """

        # This math literally saves companies from agents making bad deals
        options = self._generate_all_resolution_options(conflict)

        scored_options = []
        for option in options:
            score = self._score_resolution_business_impact(option)
            scored_options.append((score, option))

        # Sort by business score (ROI > risk > reversibility > everything else)
        scored_options.sort(key=lambda x: x[0], reverse=True)

        return scored_options[0][1]

    @lru_cache
    def _generate_all_resolution_options(self, conflict: Conflict) -> List[ResolutionOption]:
        """
        Generate RESOLUTIONS, not compromises.

        Instead of "let's split the difference" generate:
        - Agent A wins completely
        - Agent B wins completely
        - Both lose (clear restart)
        - Neither wins (get human decision)
        - Defer to external validation (get data)
        - Timebox experiment (safe test)
        """
        return [
            ResolutionOption(agent_a_wins, impact=conflict.agent_a_priority, risk=conflict.agent_a_risk, reversibility=conflict.agent_a_reversible),
            ResolutionOption(agent_b_wins, impact=conflict.agent_b_priority, risk=conflict.agent_b_risk, reversibility=conflict.agent_b_reversible),
            ResolutionOption(both_lose_restart, impact=0, risk=0, reversibility=100),
            ResolutionOption(escalate_to_human, impact=0, risk=0, reversibility=100),
            ResolutionOption(timebox_experiment, impact=50, risk=20, reversibility=80)
        ]

    def _score_resolution_business_impact(self, resolution: ResolutionOption) -> float:
        """
        Convert rules into numbers that survive agent complexity.

        Score = (Business Impact * 0.5) - (Risk Score * 0.3) + (Reversibility * 0.2)

        This math ensures we optimize for:
        1. Business value (50% weight)
        2. Risk reduction (30% weight)
        3. Ability to fix mistakes (20% weight)

        The weights are hard-coded because in production you want published rules, not AI opinions.
        """
        return (
            resolution.impact * 0.5 -
            resolution.risk * 0.3 +
            resolution.reversibility * 0.2
        )

from dataclasses import dataclass
from datetime import datetime, timedelta
from typing import List, Dict
import statistics

@dataclass
class AgentMetric:
    timestamp: datetime
    cost_usd: float
    latency_ms: float
    success: bool
    error_type: str | None
    agent_id: str
    task_type: str


class ProductionMonitor:
    """
    This is what you check at 3 AM when PagerDuty goes off.

    Key insight: You don't monitor "agent performance."
    You monitor "is this costing more than it's worth?"
    """

    def __init__(self):
        self.cost_alert_threshold = 10.0  # $10/hour triggers investigation
        self.error_rate_threshold = 0.15  # 15% error rate triggers alert
        self.latency_threshold_ms = 5000  # 5 seconds triggers investigation

    def check_agent_health(self, metrics: List[AgentMetric]) -> Dict:
        """Run this every minute. It catches disasters before they become headlines."""

        if not metrics:
            return {"status": "no_data", "action": "check_agent_connectivity"}

        recent_metrics = [m for m in metrics if m.timestamp > datetime.now() - timedelta(hours=1)]

        # The metrics that actually matter to your career
        hourly_cost = sum(m.cost_usd for m in recent_metrics)
        error_rate = sum(1 for m in recent_metrics if not m.success) / len(recent_metrics)
        avg_latency = statistics.mean([m.latency_ms for m in recent_metrics])

        alerts = []

        # Cost explosion detection
        if hourly_cost > self.cost_alert_threshold:
            alerts.append({
                "severity": "high",
                "message": f"💰 Cost spike: ${hourly_cost:.2f}/hour",
                "action": "review_recent_tasks_and_maybe_kill_agent"
            })

        # Error rate spike
        if error_rate > self.error_rate_threshold:
            alerts.append({
                "severity": "medium",
                "message": f"⚠️ Error rate: {error_rate:.1%}",
                "action": "investigate_common_failure_patterns"
            })

        # Latency degradation
        if avg_latency > self.latency_threshold_ms:
            alerts.append({
                "severity": "low",
                "message": f"🐌 Slow responses: {avg_latency:.0f}ms avg",
                "action": "check_model_availability_and_context_size"
            })

        return {
            "status": "alerting" if alerts else "healthy",
            "hourly_cost": hourly_cost,
            "error_rate": error_rate,
            "avg_latency_ms": avg_latency,
            "alerts": alerts
        }

    def generate_daily_report(self, metrics: List[AgentMetric]) -> str:
        """
        This is what you send to your boss every morning.

        Format: "Here's what our agents did yesterday and whether it was worth it"
        """
        yesterday = [m for m in metrics if m.timestamp > datetime.now() - timedelta(days=1)]

        total_cost = sum(m.cost_usd for m in yesterday)
        total_tasks = len(yesterday)
        successful_tasks = sum(1 for m in yesterday if m.success)
        cost_per_success = total_cost / successful_tasks if successful_tasks > 0 else float('inf')

        # This calculation answers: "Would it have been cheaper to hire a human?"
        human_cost_equivalent = total_tasks * 0.50  # Assume $0.50 per task for human
        roi = (human_cost_equivalent - total_cost) / human_cost_equivalent * 100 if human_cost_equivalent > 0 else 0

        return f"""
        🤖 Daily Agent Report

        Tasks completed: {total_tasks} ({successful_tasks} successful)
        Total cost: ${total_cost:.2f}
        Cost per successful task: ${cost_per_success:.4f}
        ROI vs human labor: {roi:.1f}%

        {'✅ Agents paying for themselves' if roi > 0 else '🔴 Agents costing more than humans - investigate'}
        """

# Production Agent Configuration
# Copy this file and customize for your environment

agent:
  name: "production-agent"
  model: "gpt-4o"
  temperature: 0  # Deterministic outputs for production

safety:
  # Cost controls - adjust based on your budget
  max_cost_per_hour: 5.0        # Kill agent if burning > $5/hr
  max_cost_per_day: 50.0        # Daily budget cap
  max_retries: 3                # Prevent infinite loops
  max_session_length: 3600      # 1 hour max before context refresh

  # Required context - agent won't start without these
  required_context_keys:
    - project_id
    - user_id
    - budget_limit
    - task_priority

monitoring:
  # Alert thresholds
  cost_alert_threshold: 10.0    # Alert at $10/hr
  error_rate_threshold: 0.15    # Alert at 15% errors
  latency_threshold_ms: 5000    # Alert at 5s latency

  # Where to send alerts
  alert_channels:
    - type: slack
      webhook_url: "${SLACK_WEBHOOK_URL}"
    - type: email
      recipients:
        - oncall@yourcompany.com

  # Metrics storage
  metrics_backend: "prometheus"
  metrics_port: 9090

escalation:
  # What happens when agent needs help
  on_failure:
    - notify_oncall: true
    - create_incident: true
    - auto_rollback: true

  # PagerDuty integration
  pagerduty:
    service_key: "${PAGERDUTY_SERVICE_KEY}"
    severity_map:
      critical: "P1"
      warning: "P2"

logging:
  level: "INFO"
  format: "json"  # Structured logs for production
  output: "/var/log/agent/production.log"
  retention_days: 30

FROM python:3.11-slim

WORKDIR /app

# Install dependencies
COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt

# Copy agent code
COPY . .

# Create non-root user for security
RUN useradd -m -u 1000 agent && \
    chown -R agent:agent /app
USER agent

# Health check
HEALTHCHECK --interval=30s --timeout=10s --start-period=5s --retries=3 \
    CMD curl -f http://localhost:8080/health || exit 1

# Run the agent service
CMD ["python", "-m", "agent.server"]

langchain-openai>=0.1.0
langchain-core>=0.1.0
pydantic>=2.0.0
prometheus-client>=0.19.0
pyyaml>=6.0
structlog>=23.0.0
httpx>=0.25.0

version: "3.8"

services:
  agent:
    build: .
    container_name: production-agent
    restart: unless-stopped
    ports:
      - "8080:8080"      # API endpoint
      - "9090:9090"      # Metrics endpoint
    environment:
      - OPENAI_API_KEY=${OPENAI_API_KEY}
      - ANTHROPIC_API_KEY=${ANTHROPIC_API_KEY}
      - CONFIG_PATH=/app/config.yaml
      - LOG_LEVEL=INFO
    volumes:
      - ./config.yaml:/app/config.yaml:ro
      - agent-logs:/var/log/agent
    depends_on:
      - redis
      - prometheus
    deploy:
      resources:
        limits:
          memory: 2G
          cpus: "2.0"
    healthcheck:
      test: ["CMD", "curl", "-f", "http://localhost:8080/health"]
      interval: 30s
      timeout: 10s
      retries: 3

  redis:
    image: redis:7-alpine
    container_name: agent-redis
    restart: unless-stopped
    volumes:
      - redis-data:/data
    command: redis-server --appendonly yes

  prometheus:
    image: prom/prometheus:latest
    container_name: agent-prometheus
    restart: unless-stopped
    ports:
      - "9091:9090"
    volumes:
      - ./prometheus.yml:/etc/prometheus/prometheus.yml:ro
      - prometheus-data:/prometheus
    command:
      - '--config.file=/etc/prometheus/prometheus.yml'
      - '--storage.tsdb.retention.time=30d'

  alertmanager:
    image: prom/alertmanager:latest
    container_name: agent-alertmanager
    restart: unless-stopped
    ports:
      - "9093:9093"
    volumes:
      - ./alertmanager.yml:/etc/alertmanager/alertmanager.yml:ro

volumes:
  agent-logs:
  redis-data:
  prometheus-data: