"""
Monitoring and logging system for workflow execution.
"""

import logging
from typing import Any, Dict, Optional, Callable
from datetime import datetime
from enum import Enum


class LogLevel(str, Enum):
    """Log level enumeration."""
    DEBUG = "debug"
    INFO = "info"
    WARNING = "warning"
    ERROR = "error"
    CRITICAL = "critical"


class WorkflowLogger:
    """
    Structured logging for workflow execution.
    """
    
    def __init__(self, workflow_name: str = "workflow", log_file: Optional[str] = None):
        """
        Initialize workflow logger.
        
        Args:
            workflow_name: Name of the workflow
            log_file: Optional file path for logging
        """
        self.workflow_name = workflow_name
        self.logger = logging.getLogger(f"ybagent.{workflow_name}")
        self.logger.setLevel(logging.DEBUG)
        
        # Console handler
        console_handler = logging.StreamHandler()
        console_handler.setLevel(logging.INFO)
        formatter = logging.Formatter(
            '%(asctime)s - %(name)s - %(levelname)s - %(message)s'
        )
        console_handler.setFormatter(formatter)
        self.logger.addHandler(console_handler)
        
        # File handler if specified
        if log_file:
            file_handler = logging.FileHandler(log_file)
            file_handler.setLevel(logging.DEBUG)
            file_handler.setFormatter(formatter)
            self.logger.addHandler(file_handler)
    
    def log(self, level: LogLevel, message: str, node_id: Optional[str] = None, **kwargs):
        """
        Log a message.
        
        Args:
            level: Log level
            message: Log message
            node_id: Optional node ID
            **kwargs: Additional context
        """
        extra_info = f" [Node: {node_id}]" if node_id else ""
        full_message = f"{message}{extra_info}"
        
        if kwargs:
            full_message += f" | Context: {kwargs}"
        
        log_method = getattr(self.logger, level.value)
        log_method(full_message)
    
    def debug(self, message: str, **kwargs):
        """Log debug message."""
        self.log(LogLevel.DEBUG, message, **kwargs)
    
    def info(self, message: str, **kwargs):
        """Log info message."""
        self.log(LogLevel.INFO, message, **kwargs)
    
    def warning(self, message: str, **kwargs):
        """Log warning message."""
        self.log(LogLevel.WARNING, message, **kwargs)
    
    def error(self, message: str, **kwargs):
        """Log error message."""
        self.log(LogLevel.ERROR, message, **kwargs)
    
    def critical(self, message: str, **kwargs):
        """Log critical message."""
        self.log(LogLevel.CRITICAL, message, **kwargs)


class ExecutionTracker:
    """
    Track execution metrics and performance.
    """
    
    def __init__(self):
        self.metrics: Dict[str, Any] = {
            "total_executions": 0,
            "successful_executions": 0,
            "failed_executions": 0,
            "node_metrics": {},
            "execution_times": []
        }
        self.event_handlers: Dict[str, list] = {
            "workflow_start": [],
            "workflow_end": [],
            "node_start": [],
            "node_end": [],
            "node_error": []
        }
    
    def on(self, event: str, handler: Callable):
        """
        Register an event handler.
        
        Args:
            event: Event name
            handler: Callback function
        """
        if event not in self.event_handlers:
            self.event_handlers[event] = []
        self.event_handlers[event].append(handler)
    
    def emit(self, event: str, data: Dict[str, Any]):
        """
        Emit an event.
        
        Args:
            event: Event name
            data: Event data
        """
        handlers = self.event_handlers.get(event, [])
        for handler in handlers:
            try:
                handler(data)
            except Exception as e:
                print(f"Error in event handler for '{event}': {e}")
    
    def track_workflow_start(self, workflow_name: str):
        """Track workflow start."""
        self.metrics["total_executions"] += 1
        self.emit("workflow_start", {
            "workflow_name": workflow_name,
            "timestamp": datetime.now().isoformat()
        })
    
    def track_workflow_end(self, workflow_name: str, success: bool, duration: float):
        """Track workflow end."""
        if success:
            self.metrics["successful_executions"] += 1
        else:
            self.metrics["failed_executions"] += 1
        
        self.metrics["execution_times"].append(duration)
        
        self.emit("workflow_end", {
            "workflow_name": workflow_name,
            "success": success,
            "duration": duration,
            "timestamp": datetime.now().isoformat()
        })
    
    def track_node_start(self, node_id: str):
        """Track node execution start."""
        if node_id not in self.metrics["node_metrics"]:
            self.metrics["node_metrics"][node_id] = {
                "executions": 0,
                "successes": 0,
                "failures": 0,
                "total_time": 0.0
            }
        
        self.metrics["node_metrics"][node_id]["executions"] += 1
        
        self.emit("node_start", {
            "node_id": node_id,
            "timestamp": datetime.now().isoformat()
        })
    
    def track_node_end(self, node_id: str, success: bool, duration: float):
        """Track node execution end."""
        if node_id in self.metrics["node_metrics"]:
            if success:
                self.metrics["node_metrics"][node_id]["successes"] += 1
            else:
                self.metrics["node_metrics"][node_id]["failures"] += 1
            
            self.metrics["node_metrics"][node_id]["total_time"] += duration
        
        self.emit("node_end", {
            "node_id": node_id,
            "success": success,
            "duration": duration,
            "timestamp": datetime.now().isoformat()
        })
    
    def track_node_error(self, node_id: str, error: str):
        """Track node error."""
        self.emit("node_error", {
            "node_id": node_id,
            "error": error,
            "timestamp": datetime.now().isoformat()
        })
    
    def get_metrics(self) -> Dict[str, Any]:
        """Get all metrics."""
        return self.metrics.copy()
    
    def get_summary(self) -> str:
        """Get a summary of metrics."""
        total = self.metrics["total_executions"]
        success = self.metrics["successful_executions"]
        failed = self.metrics["failed_executions"]
        
        avg_time = 0.0
        if self.metrics["execution_times"]:
            avg_time = sum(self.metrics["execution_times"]) / len(self.metrics["execution_times"])
        
        summary = f"""
Workflow Execution Summary:
---------------------------
Total Executions: {total}
Successful: {success}
Failed: {failed}
Success Rate: {(success/total*100) if total > 0 else 0:.1f}%
Average Execution Time: {avg_time:.2f}s

Node Metrics:
"""
        for node_id, metrics in self.metrics["node_metrics"].items():
            summary += f"\n  {node_id}:"
            summary += f"\n    Executions: {metrics['executions']}"
            summary += f"\n    Success Rate: {(metrics['successes']/metrics['executions']*100) if metrics['executions'] > 0 else 0:.1f}%"
            summary += f"\n    Avg Time: {(metrics['total_time']/metrics['executions']) if metrics['executions'] > 0 else 0:.2f}s"
        
        return summary
