"""
Exploit Queue Manager — Redis-backed priority queue for dynamic testing tasks.

Manages the flow of validated, prioritized findings to the Dynamic Agent system.
Implements priority ordering, deduplication, and status tracking.
"""

from __future__ import annotations

import json
import logging
import time
from dataclasses import dataclass, field
from datetime import datetime
from enum import Enum
from typing import Any, Optional

logger = logging.getLogger(__name__)


class QueuePriority(Enum):
    """Priority levels for queue items."""
    CRITICAL = 1
    HIGH = 2
    MEDIUM = 3
    LOW = 4
    DEFERRED = 5


class QueueStatus(Enum):
    """Status of queue items."""
    QUEUED = "queued"
    ASSIGNED = "assigned"
    RUNNING = "running"
    COMPLETED = "completed"
    FAILED = "failed"
    CANCELLED = "cancelled"


@dataclass
class QueueItem:
    """
    Item in the exploitation queue.
    
    Bridges static analysis findings to dynamic testing tasks.
    """
    queue_item_id: str
    correlation_id: str
    finding_id: str
    scan_id: str
    
    # Target information
    vuln_type: str
    target_url: str = ""
    target_method: str = "GET"
    
    # Injection context from static analysis
    injection_points: dict[str, Any] = field(default_factory=dict)
    validation_strategies: list[str] = field(default_factory=list)
    static_context: dict[str, Any] = field(default_factory=dict)
    
    # Priority
    priority: QueuePriority = QueuePriority.MEDIUM
    composite_score: float = 0.0
    
    # Execution
    status: QueueStatus = QueueStatus.QUEUED
    max_retries: int = 3
    retry_count: int = 0
    
    # Metadata
    created_at: float = field(default_factory=time.time)
    assigned_at: Optional[float] = None
    completed_at: Optional[float] = None
    assigned_agent: Optional[str] = None
    
    def to_task_payload(self) -> dict[str, Any]:
        """Convert to payload for Dynamic Agent consumption."""
        return {
            "task_id": self.queue_item_id,
            "finding_id": self.finding_id,
            "scan_id": self.scan_id,
            "vuln_type": self.vuln_type,
            "target_url": self.target_url,
            "target_method": self.target_method,
            "injection_points": [
                {
                    "parameter_name": name,
                    "parameter_location": info.get("location", "query"),
                    "parameter_type": info.get("type", "string"),
                    "current_value": info.get("value", ""),
                }
                for name, info in self.injection_points.items()
            ],
            "validation_strategies": self.validation_strategies,
            "static_context": self.static_context,
            "max_attempts": 10,
            "timeout_seconds": 300,
            "priority": self.priority.value,
        }
    
    def to_dict(self) -> dict[str, Any]:
        """Serialize to dictionary."""
        return {
            "queue_item_id": self.queue_item_id,
            "correlation_id": self.correlation_id,
            "finding_id": self.finding_id,
            "vuln_type": self.vuln_type,
            "priority": self.priority.name,
            "priority_value": self.priority.value,
            "composite_score": self.composite_score,
            "status": self.status.value,
            "target_url": self.target_url,
            "created_at": self.created_at,
        }


class ExploitQueueManager:
    """
    Redis-backed priority queue for exploitation tasks.
    
    Uses Redis Sorted Sets for priority ordering and Redis Streams
    for event-driven agent consumption.
    
    Queue semantics:
    - Items are ordered by priority (lower value = higher priority)
    - Agents claim items via atomic pop operations
    - Failed items are requeued with incremented retry count
    - Completed items are moved to a separate archive set
    """
    
    # Redis key prefixes
    KEY_QUEUE = "shannon:exploit:queue"           # Sorted set: score = priority
    KEY_ASSIGNED = "shannon:exploit:assigned"     # Hash: task_id -> agent_id
    KEY_RUNNING = "shannon:exploit:running"       # Hash: task_id -> start_time
    KEY_COMPLETED = "shannon:exploit:completed"   # Sorted set: completion time
    KEY_FAILED = "shannon:exploit:failed"         # Sorted set: failure time
    KEY_STREAM = "shannon:exploit:stream"         # Redis Stream for events
    
    def __init__(
        self,
        redis_host: str = "localhost",
        redis_port: int = 6379,
        redis_db: int = 0,
        password: Optional[str] = None,
    ):
        self.redis_host = redis_host
        self.redis_port = redis_port
        self.redis_db = redis_db
        self.password = password
        self._redis = None
    
    async def connect(self) -> None:
        """Connect to Redis."""
        import redis.asyncio as aioredis
        
        self._redis = aioredis.Redis(
            host=self.redis_host,
            port=self.redis_port,
            db=self.redis_db,
            password=self.password,
            decode_responses=True,
        )
        
        # Test connection
        await self._redis.ping()
        logger.info(f"Connected to Redis at {self.redis_host}:{self.redis_port}")
    
    async def disconnect(self) -> None:
        """Disconnect from Redis."""
        if self._redis:
            await self._redis.close()
    
    async def enqueue(
        self,
        item: QueueItem,
    ) -> str:
        """
        Add item to exploitation queue.
        
        Args:
            item: QueueItem to enqueue
            
        Returns:
            queue_item_id
        """
        # Store item data
        item_key = f"shannon:exploit:item:{item.queue_item_id}"
        await self._redis.hset(item_key, mapping={
            "data": json.dumps(item.to_dict()),
            "payload": json.dumps(item.to_task_payload()),
            "created_at": str(item.created_at),
        })
        
        # Add to priority queue (lower score = higher priority)
        score = item.priority.value * 1000 - item.composite_score
        await self._redis.zadd(
            self.KEY_QUEUE,
            {item.queue_item_id: score},
        )
        
        # Publish event
        await self._redis.xadd(
            self.KEY_STREAM,
            {
                "event": "item_queued",
                "item_id": item.queue_item_id,
                "finding_id": item.finding_id,
                "priority": item.priority.name,
                "score": item.composite_score,
            },
        )
        
        logger.info(
            f"Enqueued {item.queue_item_id} "
            f"(finding={item.finding_id}, priority={item.priority.name})"
        )
        
        return item.queue_item_id
    
    async def dequeue(
        self,
        agent_id: str,
        timeout_seconds: int = 30,
    ) -> Optional[QueueItem]:
        """
        Claim next item from queue for an agent.
        
        Uses atomic ZPOPMIN to ensure no double-claiming.
        
        Args:
            agent_id: ID of claiming agent
            timeout_seconds: How long to wait for item
            
        Returns:
            QueueItem or None if queue empty
        """
        # Atomically pop highest priority item
        result = await self._redis.zpopmin(self.KEY_QUEUE, count=1)
        
        if not result:
            return None
        
        queue_item_id, score = result[0]
        
        # Get item data
        item_key = f"shannon:exploit:item:{queue_item_id}"
        item_data = await self._redis.hget(item_key, "payload")
        
        if not item_data:
            logger.warning(f"Item data missing for {queue_item_id}")
            return None
        
        payload = json.loads(item_data)
        
        # Mark as assigned
        await self._redis.hset(
            self.KEY_ASSIGNED,
            queue_item_id,
            agent_id,
        )
        
        # Update item status
        item = QueueItem(
            queue_item_id=queue_item_id,
            correlation_id=payload.get("finding_id", ""),  # Simplified
            finding_id=payload["finding_id"],
            scan_id=payload.get("scan_id", ""),
            vuln_type=payload["vuln_type"],
            target_url=payload.get("target_url", ""),
            target_method=payload.get("target_method", "GET"),
            injection_points={
                ip["parameter_name"]: ip
                for ip in payload.get("injection_points", [])
            },
            validation_strategies=payload.get("validation_strategies", []),
            static_context=payload.get("static_context", {}),
            priority=QueuePriority(payload.get("priority", 3)),
            status=QueueStatus.ASSIGNED,
            assigned_at=time.time(),
            assigned_agent=agent_id,
        )
        
        # Publish event
        await self._redis.xadd(
            self.KEY_STREAM,
            {
                "event": "item_assigned",
                "item_id": queue_item_id,
                "agent_id": agent_id,
            },
        )
        
        logger.info(f"Dequeued {queue_item_id} for agent {agent_id}")
        
        return item
    
    async def mark_running(self, queue_item_id: str) -> None:
        """Mark item as running."""
        await self._redis.hset(
            self.KEY_RUNNING,
            queue_item_id,
            str(time.time()),
        )
        
        await self._redis.xadd(
            self.KEY_STREAM,
            {
                "event": "item_running",
                "item_id": queue_item_id,
            },
        )
    
    async def mark_completed(
        self,
        queue_item_id: str,
        result_summary: dict[str, Any],
    ) -> None:
        """Mark item as completed."""
        # Remove from running
        await self._redis.hdel(self.KEY_RUNNING, queue_item_id)
        await self._redis.hdel(self.KEY_ASSIGNED, queue_item_id)
        
        # Add to completed set
        await self._redis.zadd(
            self.KEY_COMPLETED,
            {queue_item_id: time.time()},
        )
        
        # Store result
        result_key = f"shannon:exploit:result:{queue_item_id}"
        await self._redis.hset(result_key, mapping={
            "summary": json.dumps(result_summary),
            "completed_at": str(time.time()),
        })
        
        await self._redis.xadd(
            self.KEY_STREAM,
            {
                "event": "item_completed",
                "item_id": queue_item_id,
                "validated": str(result_summary.get("is_validated", False)),
            },
        )
        
        logger.info(f"Completed {queue_item_id}")
    
    async def mark_failed(
        self,
        queue_item_id: str,
        error: str,
        retry_allowed: bool = True,
    ) -> None:
        """
        Mark item as failed.
        
        If retry is allowed and under max retries, requeue the item.
        """
        # Remove from running/assigned
        await self._redis.hdel(self.KEY_RUNNING, queue_item_id)
        await self._redis.hdel(self.KEY_ASSIGNED, queue_item_id)
        
        # Get current retry count
        item_key = f"shannon:exploit:item:{queue_item_id}"
        retry_data = await self._redis.hget(item_key, "retry_count")
        retry_count = int(retry_data or 0)
        
        if retry_allowed and retry_count < 3:
            # Requeue with incremented retry count
            await self._redis.hset(item_key, "retry_count", str(retry_count + 1))
            
            # Re-add to queue with slightly lower priority
            await self._redis.zadd(
                self.KEY_QUEUE,
                {queue_item_id: (retry_count + 1) * 1000 + 500},
            )
            
            await self._redis.xadd(
                self.KEY_STREAM,
                {
                    "event": "item_requeued",
                    "item_id": queue_item_id,
                    "retry": retry_count + 1,
                },
            )
            
            logger.info(f"Requeued {queue_item_id} (retry {retry_count + 1})")
        else:
            # Move to failed set
            await self._redis.zadd(
                self.KEY_FAILED,
                {queue_item_id: time.time()},
            )
            
            # Store error
            error_key = f"shannon:exploit:error:{queue_item_id}"
            await self._redis.hset(error_key, mapping={
                "error": error,
                "failed_at": str(time.time()),
                "total_retries": str(retry_count),
            })
            
            await self._redis.xadd(
                self.KEY_STREAM,
                {
                    "event": "item_failed",
                    "item_id": queue_item_id,
                    "error": error,
                },
            )
            
            logger.info(f"Failed {queue_item_id} after {retry_count} retries")
    
    async def get_queue_depth(self) -> dict[str, int]:
        """Get current queue statistics."""
        return {
            "queued": await self._redis.zcard(self.KEY_QUEUE),
            "assigned": await self._redis.hlen(self.KEY_ASSIGNED),
            "running": await self._redis.hlen(self.KEY_RUNNING),
            "completed": await self._redis.zcard(self.KEY_COMPLETED),
            "failed": await self._redis.zcard(self.KEY_FAILED),
        }
    
    async def get_next_items(self, count: int = 10) -> list[dict[str, Any]]:
        """Peek at next items without dequeuing."""
        items = await self._redis.zrange(
            self.KEY_QUEUE,
            0,
            count - 1,
            withscores=True,
        )
        
        result = []
        for item_id, score in items:
            item_key = f"shannon:exploit:item:{item_id}"
            data = await self._redis.hget(item_key, "data")
            if data:
                item_dict = json.loads(data)
                item_dict["queue_score"] = score
                result.append(item_dict)
        
        return result
    
    async def cancel_item(self, queue_item_id: str) -> bool:
        """Cancel a queued item."""
        # Remove from queue
        removed = await self._redis.zrem(self.KEY_QUEUE, queue_item_id)
        
        if removed:
            await self._redis.xadd(
                self.KEY_STREAM,
                {
                    "event": "item_cancelled",
                    "item_id": queue_item_id,
                },
            )
            logger.info(f"Cancelled {queue_item_id}")
            return True
        
        return False
    
    async def health_check(self) -> dict[str, Any]:
        """Check queue health."""
        try:
            await self._redis.ping()
            depths = await self.get_queue_depth()
            return {
                "status": "healthy",
                **depths,
            }
        except Exception as e:
            return {
                "status": "unhealthy",
                "error": str(e),
            }
