#!/usr/bin/env python3
"""
Genesis Browser Swarm Coordinator
==================================
Orchestrates 20+ specialized browser agents for revenue pipelines,
customer support, and onboarding automation.

Architecture:
- Agent Registry: Track all active agents
- Task Queue: Distribute work to specialized agents
- Session Pool: Manage browser sessions
- Error Recovery: Handle failures gracefully
"""

import os
import json
import asyncio
from datetime import datetime
from typing import Dict, List, Optional, Any
from dataclasses import dataclass, field
from enum import Enum
import uuid


class AgentType(Enum):
    """Agent specialization categories."""
    REVENUE = "revenue"
    SUPPORT = "support"
    ONBOARDING = "onboarding"
    YOUTUBE = "youtube"
    UTILITY = "utility"
    CAMEL = "camel"


class AgentStatus(Enum):
    """Agent operational status."""
    IDLE = "idle"
    BUSY = "busy"
    ERROR = "error"
    OFFLINE = "offline"


class TaskStatus(Enum):
    """Task execution status."""
    PENDING = "pending"
    IN_PROGRESS = "in_progress"
    COMPLETED = "completed"
    FAILED = "failed"
    CANCELLED = "cancelled"


@dataclass
class Agent:
    """Browser swarm agent definition."""
    agent_id: str
    name: str
    agent_type: AgentType
    description: str
    tools: List[str]
    status: AgentStatus = AgentStatus.IDLE
    current_task: Optional[str] = None
    tasks_completed: int = 0
    tasks_failed: int = 0
    last_active: Optional[str] = None

    def to_dict(self) -> Dict:
        return {
            'agent_id': self.agent_id,
            'name': self.name,
            'type': self.agent_type.value,
            'description': self.description,
            'tools': self.tools,
            'status': self.status.value,
            'current_task': self.current_task,
            'tasks_completed': self.tasks_completed,
            'tasks_failed': self.tasks_failed,
            'last_active': self.last_active
        }


@dataclass
class Task:
    """Task definition for agent execution."""
    task_id: str
    title: str
    description: str
    agent_type: AgentType
    assigned_agent: Optional[str] = None
    status: TaskStatus = TaskStatus.PENDING
    priority: int = 5  # 1-10, higher = more urgent
    created_at: str = field(default_factory=lambda: datetime.now().isoformat())
    started_at: Optional[str] = None
    completed_at: Optional[str] = None
    result: Optional[Dict] = None
    error: Optional[str] = None
    retries: int = 0
    max_retries: int = 3


# ============================================================================
# AGENT DEFINITIONS - 20 Specialized Browser Agents
# ============================================================================

AGENT_DEFINITIONS = [
    # Revenue Pipeline Agents (7)
    Agent("REV-001", "Lead Scraper", AgentType.REVENUE,
          "Extract leads from LinkedIn, directories, and web sources",
          ["browser_navigate", "browser_snapshot", "browser_evaluate"]),

    Agent("REV-002", "CRM Automator", AgentType.REVENUE,
          "Sync leads to HubSpot/Salesforce, update contact records",
          ["browser_fill_form", "browser_click", "browser_type"]),

    Agent("REV-003", "Outreach Sender", AgentType.REVENUE,
          "Send automated email sequences via web interfaces",
          ["browser_type", "browser_click", "browser_wait_for"]),

    Agent("REV-004", "Meeting Scheduler", AgentType.REVENUE,
          "Book calls via Calendly/Cal.com, manage calendar",
          ["browser_navigate", "browser_select_option", "browser_click"]),

    Agent("REV-005", "Invoice Generator", AgentType.REVENUE,
          "Create and send invoices via Stripe/PayPal dashboards",
          ["browser_fill_form", "browser_take_screenshot", "browser_click"]),

    Agent("REV-006", "Payment Tracker", AgentType.REVENUE,
          "Monitor payment status, send payment reminders",
          ["browser_snapshot", "browser_evaluate", "browser_navigate"]),

    Agent("REV-007", "Revenue Reporter", AgentType.REVENUE,
          "Generate revenue dashboards and reports",
          ["browser_take_screenshot", "browser_evaluate", "browser_snapshot"]),

    # Customer Support Agents (6)
    Agent("SUP-001", "Ticket Triager", AgentType.SUPPORT,
          "Classify and route support tickets by priority/category",
          ["browser_snapshot", "browser_evaluate", "browser_click"]),

    Agent("SUP-002", "FAQ Responder", AgentType.SUPPORT,
          "Auto-respond to common questions with KB articles",
          ["browser_type", "browser_click", "browser_navigate"]),

    Agent("SUP-003", "Escalation Handler", AgentType.SUPPORT,
          "Route complex issues to human agents via Slack",
          ["browser_fill_form", "browser_click", "slack_post_message"]),

    Agent("SUP-004", "Status Updater", AgentType.SUPPORT,
          "Update ticket status and notify customers",
          ["browser_click", "browser_type", "browser_select_option"]),

    Agent("SUP-005", "Feedback Collector", AgentType.SUPPORT,
          "Gather CSAT/NPS scores after resolution",
          ["browser_navigate", "browser_fill_form", "browser_click"]),

    Agent("SUP-006", "Knowledge Builder", AgentType.SUPPORT,
          "Update knowledge base from resolved tickets",
          ["browser_type", "browser_click", "memory_create_entities"]),

    # Customer Onboarding Agents (5)
    Agent("ONB-001", "Welcome Sequencer", AgentType.ONBOARDING,
          "Send welcome emails and setup guides",
          ["browser_navigate", "browser_type", "browser_click"]),

    Agent("ONB-002", "Account Provisioner", AgentType.ONBOARDING,
          "Create accounts and set permissions",
          ["browser_fill_form", "browser_click", "browser_select_option"]),

    Agent("ONB-003", "Tutorial Guide", AgentType.ONBOARDING,
          "Walk customers through product features",
          ["browser_navigate", "browser_take_screenshot", "browser_click"]),

    Agent("ONB-004", "Progress Tracker", AgentType.ONBOARDING,
          "Monitor onboarding completion milestones",
          ["browser_snapshot", "browser_evaluate", "browser_navigate"]),

    Agent("ONB-005", "Success Checker", AgentType.ONBOARDING,
          "Verify customer can perform key actions",
          ["browser_click", "browser_evaluate", "browser_wait_for"]),

    # Utility Agents (2)
    Agent("UTL-001", "Session Manager", AgentType.UTILITY,
          "Maintain auth sessions, handle cookies and tokens",
          ["browser_evaluate", "browser_file_upload", "browser_navigate"]),

    Agent("UTL-002", "Error Recovery", AgentType.UTILITY,
          "Handle browser crashes, retry failed tasks",
          ["browser_navigate", "browser_wait_for", "browser_take_screenshot"]),
]


class BrowserSwarmCoordinator:
    """
    Central coordinator for the Genesis Browser Swarm.

    Responsibilities:
    - Agent lifecycle management
    - Task distribution and load balancing
    - Error handling and recovery
    - Performance monitoring
    """

    def __init__(self):
        self.agents: Dict[str, Agent] = {}
        self.tasks: Dict[str, Task] = {}
        self.task_queue: List[str] = []
        self._initialize_agents()

    def _initialize_agents(self):
        """Register all agent definitions."""
        for agent in AGENT_DEFINITIONS:
            self.agents[agent.agent_id] = agent

    def get_agent(self, agent_id: str) -> Optional[Agent]:
        """Get agent by ID."""
        return self.agents.get(agent_id)

    def get_agents_by_type(self, agent_type: AgentType) -> List[Agent]:
        """Get all agents of a specific type."""
        return [a for a in self.agents.values() if a.agent_type == agent_type]

    def get_available_agent(self, agent_type: AgentType) -> Optional[Agent]:
        """Get an idle agent of the specified type."""
        for agent in self.get_agents_by_type(agent_type):
            if agent.status == AgentStatus.IDLE:
                return agent
        return None

    def create_task(
        self,
        title: str,
        description: str,
        agent_type: AgentType,
        priority: int = 5
    ) -> Task:
        """Create a new task."""
        task = Task(
            task_id=f"task-{uuid.uuid4().hex[:8]}",
            title=title,
            description=description,
            agent_type=agent_type,
            priority=priority
        )
        self.tasks[task.task_id] = task
        self.task_queue.append(task.task_id)
        self._sort_queue()
        return task

    def _sort_queue(self):
        """Sort task queue by priority (highest first)."""
        self.task_queue.sort(
            key=lambda tid: self.tasks[tid].priority,
            reverse=True
        )

    def assign_task(self, task_id: str, agent_id: str) -> bool:
        """Assign a task to an agent."""
        task = self.tasks.get(task_id)
        agent = self.agents.get(agent_id)

        if not task or not agent:
            return False

        if agent.status != AgentStatus.IDLE:
            return False

        task.assigned_agent = agent_id
        task.status = TaskStatus.IN_PROGRESS
        task.started_at = datetime.now().isoformat()

        agent.status = AgentStatus.BUSY
        agent.current_task = task_id
        agent.last_active = datetime.now().isoformat()

        if task_id in self.task_queue:
            self.task_queue.remove(task_id)

        return True

    def complete_task(self, task_id: str, result: Dict) -> bool:
        """Mark a task as completed."""
        task = self.tasks.get(task_id)
        if not task:
            return False

        task.status = TaskStatus.COMPLETED
        task.completed_at = datetime.now().isoformat()
        task.result = result

        if task.assigned_agent:
            agent = self.agents.get(task.assigned_agent)
            if agent:
                agent.status = AgentStatus.IDLE
                agent.current_task = None
                agent.tasks_completed += 1

        return True

    def fail_task(self, task_id: str, error: str) -> bool:
        """Mark a task as failed, potentially retry."""
        task = self.tasks.get(task_id)
        if not task:
            return False

        task.retries += 1
        task.error = error

        if task.assigned_agent:
            agent = self.agents.get(task.assigned_agent)
            if agent:
                agent.status = AgentStatus.IDLE
                agent.current_task = None

        if task.retries < task.max_retries:
            # Re-queue for retry
            task.status = TaskStatus.PENDING
            task.assigned_agent = None
            self.task_queue.append(task_id)
            self._sort_queue()
        else:
            task.status = TaskStatus.FAILED
            if task.assigned_agent:
                agent = self.agents.get(task.assigned_agent)
                if agent:
                    agent.tasks_failed += 1

        return True

    def get_status(self) -> Dict[str, Any]:
        """Get overall swarm status."""
        agents_by_status = {}
        for status in AgentStatus:
            agents_by_status[status.value] = len([
                a for a in self.agents.values() if a.status == status
            ])

        tasks_by_status = {}
        for status in TaskStatus:
            tasks_by_status[status.value] = len([
                t for t in self.tasks.values() if t.status == status
            ])

        return {
            'total_agents': len(self.agents),
            'agents_by_status': agents_by_status,
            'total_tasks': len(self.tasks),
            'tasks_by_status': tasks_by_status,
            'queue_length': len(self.task_queue),
            'timestamp': datetime.now().isoformat()
        }

    def get_agent_prompt(self, agent_id: str, task_description: str, context: str = "") -> str:
        """Generate the master prompt for an agent."""
        agent = self.agents.get(agent_id)
        if not agent:
            return ""

        return MASTER_PROMPT_TEMPLATE.format(
            AGENT_ID=agent.agent_id,
            AGENT_ROLE=agent.name,
            SPECIALIZATION=agent.description,
            TASK_DESCRIPTION=task_description,
            MEMORY_CONTEXT=context
        )


# ============================================================================
# MASTER PROMPT TEMPLATE
# ============================================================================

MASTER_PROMPT_TEMPLATE = """# GENESIS BROWSER SWARM - MASTER SYSTEM PROMPT

You are a specialized browser automation agent in the Genesis Browser Swarm.

## Your Identity
- Agent ID: {AGENT_ID}
- Role: {AGENT_ROLE}
- Specialization: {SPECIALIZATION}

## Core Capabilities
You have access to Playwright browser automation tools:
- Navigation: browser_navigate, browser_navigate_back, browser_tabs
- Interaction: browser_click, browser_type, browser_fill_form, browser_select_option
- State: browser_snapshot, browser_take_screenshot, browser_evaluate
- Waiting: browser_wait_for, browser_handle_dialog
- Advanced: browser_drag, browser_hover, browser_file_upload

## Operating Principles

### 1. Reliability First
- Always verify page load before interaction
- Use browser_wait_for before clicking dynamic elements
- Take screenshots on errors for debugging
- Retry failed operations up to 3 times with exponential backoff

### 2. Memory Integration
- Log all significant actions to Genesis Memory via PostToolUse hook
- Query memory for context before starting tasks
- Store learned patterns (selectors, workflows) for reuse

### 3. Coordination Protocol
- Check task queue before starting new work
- Report completion status to coordinator
- Escalate failures after 3 retries
- Share session state with related agents

### 4. Security Compliance
- Never store credentials in memory
- Use session manager for authentication
- Respect rate limits and robots.txt
- Log all data access for audit

## Task Execution Flow

1. **RECEIVE** task from coordinator
2. **VERIFY** browser session is active
3. **NAVIGATE** to target URL
4. **WAIT** for page ready state
5. **EXECUTE** task-specific actions
6. **CAPTURE** result (screenshot/data)
7. **STORE** in Genesis Memory
8. **REPORT** completion to coordinator

## Error Handling

If error occurs:
1. Capture screenshot: browser_take_screenshot
2. Log to Genesis Memory with error details
3. Attempt recovery (refresh, re-auth)
4. If recovery fails, escalate to UTL-002

## Current Task
{TASK_DESCRIPTION}

## Available Context
{MEMORY_CONTEXT}

Begin execution.
"""


# Singleton coordinator instance
_coordinator: Optional[BrowserSwarmCoordinator] = None

def get_coordinator() -> BrowserSwarmCoordinator:
    """Get the singleton coordinator instance."""
    global _coordinator
    if _coordinator is None:
        _coordinator = BrowserSwarmCoordinator()
    return _coordinator


# CLI interface
if __name__ == "__main__":
    coordinator = get_coordinator()
    status = coordinator.get_status()

    print("=" * 60)
    print("GENESIS BROWSER SWARM STATUS")
    print("=" * 60)
    print(f"\nTotal Agents: {status['total_agents']}")
    print("\nAgents by Type:")
    for agent_type in AgentType:
        agents = coordinator.get_agents_by_type(agent_type)
        print(f"  {agent_type.value}: {len(agents)} agents")
        for agent in agents:
            print(f"    - {agent.agent_id}: {agent.name}")

    print(f"\nAgents by Status: {status['agents_by_status']}")
    print(f"Tasks by Status: {status['tasks_by_status']}")
    print(f"Queue Length: {status['queue_length']}")