import os
import json
import numpy as np
import gradio as gr
import requests
import pandas as pd
import datetime
from typing import List, Dict, Any
import hashlib
import asyncio
from enum import Enum
from dataclasses import dataclass

# Import our modules
from models import ReliabilityEvent, EventSeverity, AnomalyResult, HealingAction
from healing_policies import PolicyEngine

# === Configuration ===
HF_TOKEN = os.getenv("HF_TOKEN", "").strip()
HF_API_URL = "https://router.huggingface.co/hf-inference/v1/completions"
HEADERS = {"Authorization": f"Bearer {HF_TOKEN}"} if HF_TOKEN else {}

# === FAISS & Embeddings Setup ===
try:
    from sentence_transformers import SentenceTransformer
    import faiss
    
    VECTOR_DIM = 384
    INDEX_FILE = "incident_vectors.index"
    TEXTS_FILE = "incident_texts.json"
    
    # Try to load model with error handling
    try:
        model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
    except Exception as e:
        print(f"Model loading warning: {e}")
        from sentence_transformers import SentenceTransformer as ST
        model = ST("sentence-transformers/all-MiniLM-L6-v2")
    
    if os.path.exists(INDEX_FILE):
        index = faiss.read_index(INDEX_FILE)
        with open(TEXTS_FILE, "r") as f:
            incident_texts = json.load(f)
    else:
        index = faiss.IndexFlatL2(VECTOR_DIM)
        incident_texts = []
        
except ImportError as e:
    print(f"Warning: FAISS or SentenceTransformers not available: {e}")
    index = None
    incident_texts = []
    model = None

def save_index():
    """Save FAISS index and incident texts"""
    if index is not None:
        faiss.write_index(index, INDEX_FILE)
        with open(TEXTS_FILE, "w") as f:
            json.dump(incident_texts, f)

# === Predictive Models ===
@dataclass
class ForecastResult:
    metric: str
    predicted_value: float
    confidence: float
    trend: str  # "increasing", "decreasing", "stable"
    time_to_threshold: Any = None
    risk_level: str = "low"  # low, medium, high, critical

class SimplePredictiveEngine:
    """Lightweight forecasting engine optimized for Hugging Face Spaces"""
    
    def __init__(self, history_window: int = 50):
        self.history_window = history_window
        self.service_history: Dict[str, List] = {}
        self.prediction_cache: Dict[str, ForecastResult] = {}
        
    def add_telemetry(self, service: str, event_data: Dict):
        """Add telemetry data to service history"""
        if service not in self.service_history:
            self.service_history[service] = []
        
        telemetry_point = {
            'timestamp': datetime.datetime.now(),
            'latency': event_data.get('latency_p99', 0),
            'error_rate': event_data.get('error_rate', 0),
            'throughput': event_data.get('throughput', 0),
            'cpu_util': event_data.get('cpu_util'),
            'memory_util': event_data.get('memory_util')
        }
        
        self.service_history[service].append(telemetry_point)
        
        # Keep only recent history
        if len(self.service_history[service]) > self.history_window:
            self.service_history[service].pop(0)
    
    def forecast_service_health(self, service: str, lookahead_minutes: int = 15) -> List[ForecastResult]:
        """Forecast service health metrics"""
        if service not in self.service_history or len(self.service_history[service]) < 10:
            return []
        
        history = self.service_history[service]
        forecasts = []
        
        # Forecast latency
        latency_forecast = self._forecast_latency(history, lookahead_minutes)
        if latency_forecast:
            forecasts.append(latency_forecast)
        
        # Forecast error rate
        error_forecast = self._forecast_error_rate(history, lookahead_minutes)
        if error_forecast:
            forecasts.append(error_forecast)
        
        # Forecast resource utilization
        resource_forecasts = self._forecast_resources(history, lookahead_minutes)
        forecasts.extend(resource_forecasts)
        
        # Cache results
        for forecast in forecasts:
            cache_key = f"{service}_{forecast.metric}"
            self.prediction_cache[cache_key] = forecast
        
        return forecasts
    
    def _forecast_latency(self, history: List, lookahead_minutes: int) -> Any:
        """Forecast latency using linear regression and trend analysis"""
        try:
            latencies = [point['latency'] for point in history[-20:]]
            
            if len(latencies) < 5:
                return None
            
            # Simple linear trend
            x = np.arange(len(latencies))
            slope, intercept = np.polyfit(x, latencies, 1)
            
            # Predict next value
            next_x = len(latencies)
            predicted_latency = slope * next_x + intercept
            
            # Calculate confidence based on data quality
            residuals = latencies - (slope * x + intercept)
            confidence = max(0, 1 - (np.std(residuals) / max(1, np.mean(latencies))))
            
            # Determine trend
            if slope > 5:
                trend = "increasing"
                risk = "high" if predicted_latency > 300 else "medium"
            elif slope < -2:
                trend = "decreasing" 
                risk = "low"
            else:
                trend = "stable"
                risk = "low"
            
            # Calculate time to reach critical threshold (500ms)
            time_to_critical = None
            if slope > 0 and predicted_latency < 500:
                time_to_critical = datetime.timedelta(
                    minutes=lookahead_minutes * (500 - predicted_latency) / max(0.1, (predicted_latency - latencies[-1]))
                )
            
            return ForecastResult(
                metric="latency",
                predicted_value=predicted_latency,
                confidence=confidence,
                trend=trend,
                time_to_threshold=time_to_critical,
                risk_level=risk
            )
            
        except Exception as e:
            print(f"Latency forecast error: {e}")
            return None
    
    def _forecast_error_rate(self, history: List, lookahead_minutes: int) -> Any:
        """Forecast error rate using exponential smoothing"""
        try:
            error_rates = [point['error_rate'] for point in history[-15:]]
            
            if len(error_rates) < 5:
                return None
            
            # Exponential smoothing
            alpha = 0.3
            forecast = error_rates[0]
            for rate in error_rates[1:]:
                forecast = alpha * rate + (1 - alpha) * forecast
            
            predicted_rate = forecast
            
            # Trend analysis
            recent_trend = np.mean(error_rates[-3:]) - np.mean(error_rates[-6:-3])
            
            if recent_trend > 0.02:
                trend = "increasing"
                risk = "high" if predicted_rate > 0.1 else "medium"
            elif recent_trend < -0.01:
                trend = "decreasing"
                risk = "low"
            else:
                trend = "stable"
                risk = "low"
            
            # Confidence based on volatility
            confidence = max(0, 1 - (np.std(error_rates) / max(0.01, np.mean(error_rates))))
            
            return ForecastResult(
                metric="error_rate",
                predicted_value=predicted_rate,
                confidence=confidence,
                trend=trend,
                risk_level=risk
            )
            
        except Exception as e:
            print(f"Error rate forecast error: {e}")
            return None
    
    def _forecast_resources(self, history: List, lookahead_minutes: int) -> List[ForecastResult]:
        """Forecast CPU and memory utilization"""
        forecasts = []
        
        # CPU forecast
        cpu_values = [point['cpu_util'] for point in history if point.get('cpu_util') is not None]
        if len(cpu_values) >= 5:
            try:
                predicted_cpu = np.mean(cpu_values[-5:])
                trend = "increasing" if cpu_values[-1] > np.mean(cpu_values[-10:-5]) else "stable"
                
                risk = "low"
                if predicted_cpu > 0.8:
                    risk = "critical" if predicted_cpu > 0.9 else "high"
                elif predicted_cpu > 0.7:
                    risk = "medium"
                
                forecasts.append(ForecastResult(
                    metric="cpu_util",
                    predicted_value=predicted_cpu,
                    confidence=0.7,
                    trend=trend,
                    risk_level=risk
                ))
            except Exception as e:
                print(f"CPU forecast error: {e}")
        
        # Memory forecast
        memory_values = [point['memory_util'] for point in history if point.get('memory_util') is not None]
        if len(memory_values) >= 5:
            try:
                predicted_memory = np.mean(memory_values[-5:])
                trend = "increasing" if memory_values[-1] > np.mean(memory_values[-10:-5]) else "stable"
                
                risk = "low"
                if predicted_memory > 0.8:
                    risk = "critical" if predicted_memory > 0.9 else "high"
                elif predicted_memory > 0.7:
                    risk = "medium"
                
                forecasts.append(ForecastResult(
                    metric="memory_util",
                    predicted_value=predicted_memory,
                    confidence=0.7,
                    trend=trend,
                    risk_level=risk
                ))
            except Exception as e:
                print(f"Memory forecast error: {e}")
        
        return forecasts
    
    def get_predictive_insights(self, service: str) -> Dict[str, Any]:
        """Generate actionable insights from forecasts"""
        forecasts = self.forecast_service_health(service)
        
        critical_risks = [f for f in forecasts if f.risk_level in ["high", "critical"]]
        warnings = []
        recommendations = []
        
        for forecast in critical_risks:
            if forecast.metric == "latency" and forecast.risk_level in ["high", "critical"]:
                warnings.append(f"📈 Latency expected to reach {forecast.predicted_value:.0f}ms")
                if forecast.time_to_threshold:
                    minutes = int(forecast.time_to_threshold.total_seconds() / 60)
                    recommendations.append(f"⏰ Critical latency (~500ms) in ~{minutes} minutes")
                recommendations.append("🔧 Consider scaling or optimizing dependencies")
            
            elif forecast.metric == "error_rate" and forecast.risk_level in ["high", "critical"]:
                warnings.append(f"🚨 Errors expected to reach {forecast.predicted_value*100:.1f}%")
                recommendations.append("🐛 Investigate recent deployments or dependency issues")
            
            elif forecast.metric == "cpu_util" and forecast.risk_level in ["high", "critical"]:
                warnings.append(f"🔥 CPU expected at {forecast.predicted_value*100:.1f}%")
                recommendations.append("⚡ Consider scaling compute resources")
            
            elif forecast.metric == "memory_util" and forecast.risk_level in ["high", "critical"]:
                warnings.append(f"💾 Memory expected at {forecast.predicted_value*100:.1f}%")
                recommendations.append("🧹 Check for memory leaks or optimize usage")
        
        return {
            'service': service,
            'forecasts': [f.__dict__ for f in forecasts],
            'warnings': warnings[:3],
            'recommendations': list(dict.fromkeys(recommendations))[:3],
            'critical_risk_count': len(critical_risks),
            'forecast_timestamp': datetime.datetime.now().isoformat()
        }

# === Core Engine Components ===
policy_engine = PolicyEngine()
events_history: List[ReliabilityEvent] = []

class BusinessImpactCalculator:
    """Calculate business impact of anomalies"""
    
    def __init__(self, revenue_per_request: float = 0.01):
        self.revenue_per_request = revenue_per_request
    
    def calculate_impact(self, event: ReliabilityEvent, duration_minutes: int = 5) -> Dict[str, Any]:
        base_revenue_per_minute = 100
        
        impact_multiplier = 1.0
        
        if event.latency_p99 > 300:
            impact_multiplier += 0.5
        if event.error_rate > 0.1:
            impact_multiplier += 0.8
        if event.cpu_util and event.cpu_util > 0.9:
            impact_multiplier += 0.3
        
        revenue_loss = base_revenue_per_minute * impact_multiplier * (duration_minutes / 60)
        
        base_users_affected = 1000
        user_impact_multiplier = (event.error_rate * 10) + (max(0, event.latency_p99 - 100) / 500)
        affected_users = int(base_users_affected * user_impact_multiplier)
        
        if revenue_loss > 500 or affected_users > 5000:
            severity = "CRITICAL"
        elif revenue_loss > 100 or affected_users > 1000:
            severity = "HIGH"
        elif revenue_loss > 50 or affected_users > 500:
            severity = "MEDIUM"
        else:
            severity = "LOW"
        
        return {
            'revenue_loss_estimate': round(revenue_loss, 2),
            'affected_users_estimate': affected_users,
            'severity_level': severity,
            'throughput_reduction_pct': round(min(100, user_impact_multiplier * 100), 1)
        }

business_calculator = BusinessImpactCalculator()

class AdvancedAnomalyDetector:
    """Enhanced anomaly detection with adaptive thresholds"""
    
    def __init__(self):
        self.historical_data = []
        self.adaptive_thresholds = {
            'latency_p99': 150,
            'error_rate': 0.05
        }
    
    def detect_anomaly(self, event: ReliabilityEvent) -> bool:
        latency_anomaly = event.latency_p99 > self.adaptive_thresholds['latency_p99']
        error_anomaly = event.error_rate > self.adaptive_thresholds['error_rate']
        
        resource_anomaly = False
        if event.cpu_util and event.cpu_util > 0.9:
            resource_anomaly = True
        if event.memory_util and event.memory_util > 0.9:
            resource_anomaly = True
        
        self._update_thresholds(event)
        
        return latency_anomaly or error_anomaly or resource_anomaly
    
    def _update_thresholds(self, event: ReliabilityEvent):
        self.historical_data.append(event)
        
        if len(self.historical_data) > 100:
            self.historical_data.pop(0)
        
        if len(self.historical_data) > 10:
            recent_latencies = [e.latency_p99 for e in self.historical_data[-20:]]
            self.adaptive_thresholds['latency_p99'] = np.percentile(recent_latencies, 90)

anomaly_detector = AdvancedAnomalyDetector()

# === Multi-Agent Foundation ===
class AgentSpecialization(Enum):
    DETECTIVE = "anomaly_detection"
    DIAGNOSTICIAN = "root_cause_analysis"
    PREDICTIVE = "predictive_analytics"

class BaseAgent:
    def __init__(self, specialization: AgentSpecialization):
        self.specialization = specialization
    
    async def analyze(self, event: ReliabilityEvent) -> Dict[str, Any]:
        raise NotImplementedError

class AnomalyDetectionAgent(BaseAgent):
    def __init__(self):
        super().__init__(AgentSpecialization.DETECTIVE)
    
    async def analyze(self, event: ReliabilityEvent) -> Dict[str, Any]:
        anomaly_score = self._calculate_anomaly_score(event)
        
        return {
            'specialization': self.specialization.value,
            'confidence': anomaly_score,
            'findings': {
                'anomaly_score': anomaly_score,
                'severity_tier': self._classify_severity(anomaly_score),
                'primary_metrics_affected': self._identify_affected_metrics(event)
            },
            'recommendations': self._generate_detection_recommendations(event, anomaly_score)
        }
    
    def _calculate_anomaly_score(self, event: ReliabilityEvent) -> float:
        scores = []
        
        if event.latency_p99 > 150:
            latency_score = min(1.0, (event.latency_p99 - 150) / 500)
            scores.append(0.4 * latency_score)
        
        if event.error_rate > 0.05:
            error_score = min(1.0, event.error_rate / 0.3)
            scores.append(0.3 * error_score)
        
        resource_score = 0
        if event.cpu_util and event.cpu_util > 0.8:
            resource_score += 0.15 * min(1.0, (event.cpu_util - 0.8) / 0.2)
        if event.memory_util and event.memory_util > 0.8:
            resource_score += 0.15 * min(1.0, (event.memory_util - 0.8) / 0.2)
        scores.append(resource_score)
        
        return min(1.0, sum(scores))
    
    def _classify_severity(self, anomaly_score: float) -> str:
        if anomaly_score > 0.8:
            return "CRITICAL"
        elif anomaly_score > 0.6:
            return "HIGH"
        elif anomaly_score > 0.4:
            return "MEDIUM"
        else:
            return "LOW"
    
    def _identify_affected_metrics(self, event: ReliabilityEvent) -> List[Dict[str, Any]]:
        affected = []
        
        if event.latency_p99 > 500:
            affected.append({"metric": "latency", "value": event.latency_p99, "severity": "CRITICAL", "threshold": 150})
        elif event.latency_p99 > 300:
            affected.append({"metric": "latency", "value": event.latency_p99, "severity": "HIGH", "threshold": 150})
        elif event.latency_p99 > 150:
            affected.append({"metric": "latency", "value": event.latency_p99, "severity": "MEDIUM", "threshold": 150})
        
        if event.error_rate > 0.3:
            affected.append({"metric": "error_rate", "value": event.error_rate, "severity": "CRITICAL", "threshold": 0.05})
        elif event.error_rate > 0.15:
            affected.append({"metric": "error_rate", "value": event.error_rate, "severity": "HIGH", "threshold": 0.05})
        elif event.error_rate > 0.05:
            affected.append({"metric": "error_rate", "value": event.error_rate, "severity": "MEDIUM", "threshold": 0.05})
        
        if event.cpu_util and event.cpu_util > 0.9:
            affected.append({"metric": "cpu", "value": event.cpu_util, "severity": "CRITICAL", "threshold": 0.8})
        elif event.cpu_util and event.cpu_util > 0.8:
            affected.append({"metric": "cpu", "value": event.cpu_util, "severity": "HIGH", "threshold": 0.8})
            
        if event.memory_util and event.memory_util > 0.9:
            affected.append({"metric": "memory", "value": event.memory_util, "severity": "CRITICAL", "threshold": 0.8})
        elif event.memory_util and event.memory_util > 0.8:
            affected.append({"metric": "memory", "value": event.memory_util, "severity": "HIGH", "threshold": 0.8})
        
        return affected
    
    def _generate_detection_recommendations(self, event: ReliabilityEvent, anomaly_score: float) -> List[str]:
        recommendations = []
        affected_metrics = self._identify_affected_metrics(event)
        
        for metric in affected_metrics:
            metric_name = metric["metric"]
            severity = metric["severity"]
            value = metric["value"]
            threshold = metric["threshold"]
            
            if metric_name == "latency":
                if severity == "CRITICAL":
                    recommendations.append(f"🚨 CRITICAL: Latency {value}ms (>{threshold}ms) - Check database & external dependencies")
                elif severity == "HIGH":
                    recommendations.append(f"⚠️ HIGH: Latency {value}ms (>{threshold}ms) - Investigate service performance")
                else:
                    recommendations.append(f"📈 Latency elevated: {value}ms (>{threshold}ms) - Monitor trend")
                    
            elif metric_name == "error_rate":
                if severity == "CRITICAL":
                    recommendations.append(f"🚨 CRITICAL: Error rate {value*100:.1f}% (>{threshold*100}%) - Check recent deployments")
                elif severity == "HIGH":
                    recommendations.append(f"⚠️ HIGH: Error rate {value*100:.1f}% (>{threshold*100}%) - Review application logs")
                else:
                    recommendations.append(f"📈 Errors increasing: {value*100:.1f}% (>{threshold*100}%)")
                    
            elif metric_name == "cpu":
                recommendations.append(f"🔥 CPU {severity}: {value*100:.1f}% utilization - Consider scaling")
                
            elif metric_name == "memory":
                recommendations.append(f"💾 Memory {severity}: {value*100:.1f}% utilization - Check for memory leaks")
        
        if anomaly_score > 0.8:
            recommendations.append("🎯 IMMEDIATE ACTION REQUIRED: Multiple critical metrics affected")
        elif anomaly_score > 0.6:
            recommendations.append("🎯 INVESTIGATE: Significant performance degradation detected")
        elif anomaly_score > 0.4:
            recommendations.append("📊 MONITOR: Early warning signs detected")
        
        return recommendations[:4]

class RootCauseAgent(BaseAgent):
    def __init__(self):
        super().__init__(AgentSpecialization.DIAGNOSTICIAN)
    
    async def analyze(self, event: ReliabilityEvent) -> Dict[str, Any]:
        causes = self._analyze_potential_causes(event)
        
        return {
            'specialization': self.specialization.value,
            'confidence': 0.7,
            'findings': {
                'likely_root_causes': causes,
                'evidence_patterns': self._identify_evidence(event),
                'investigation_priority': self._prioritize_investigation(causes)
            },
            'recommendations': [
                f"Check {cause['cause']} for issues" for cause in causes[:2]
            ]
        }
    
    def _analyze_potential_causes(self, event: ReliabilityEvent) -> List[Dict[str, Any]]:
        causes = []
        
        if event.latency_p99 > 500 and event.error_rate > 0.2:
            causes.append({
                "cause": "Database/External Dependency Failure",
                "confidence": 0.85,
                "evidence": f"Extreme latency ({event.latency_p99}ms) with high errors ({event.error_rate*100:.1f}%)",
                "investigation": "Check database connection pool, external API health"
            })
        
        if event.cpu_util and event.cpu_util > 0.9 and event.memory_util and event.memory_util > 0.9:
            causes.append({
                "cause": "Resource Exhaustion",
                "confidence": 0.90,
                "evidence": f"CPU ({event.cpu_util*100:.1f}%) and Memory ({event.memory_util*100:.1f}%) critically high",
                "investigation": "Check for memory leaks, infinite loops, insufficient resources"
            })
        
        if event.error_rate > 0.3 and event.latency_p99 < 200:
            causes.append({
                "cause": "Application Bug / Configuration Issue",
                "confidence": 0.75,
                "evidence": f"High error rate ({event.error_rate*100:.1f}%) without latency impact",
                "investigation": "Review recent deployments, configuration changes, application logs"
            })
        
        if 200 <= event.latency_p99 <= 400 and 0.05 <= event.error_rate <= 0.15:
            causes.append({
                "cause": "Gradual Performance Degradation",
                "confidence": 0.65,
                "evidence": f"Moderate latency ({event.latency_p99}ms) and errors ({event.error_rate*100:.1f}%)",
                "investigation": "Check resource trends, dependency performance, capacity planning"
            })
        
        if not causes:
            causes.append({
                "cause": "Unknown - Requires Investigation",
                "confidence": 0.3,
                "evidence": "Pattern does not match known failure modes",
                "investigation": "Complete system review needed"
            })
        
        return causes
    
    def _identify_evidence(self, event: ReliabilityEvent) -> List[str]:
        evidence = []
        if event.latency_p99 > event.error_rate * 1000:
            evidence.append("latency_disproportionate_to_errors")
        if event.cpu_util and event.cpu_util > 0.8 and event.memory_util and event.memory_util > 0.8:
            evidence.append("correlated_resource_exhaustion")
        return evidence
    
    def _prioritize_investigation(self, causes: List[Dict[str, Any]]) -> str:
        for cause in causes:
            if "Database" in cause["cause"] or "Resource Exhaustion" in cause["cause"]:
                return "HIGH"
        return "MEDIUM"

class PredictiveAgent(BaseAgent):
    def __init__(self):
        super().__init__(AgentSpecialization.PREDICTIVE)
        self.engine = SimplePredictiveEngine()
    
    async def analyze(self, event: ReliabilityEvent) -> Dict[str, Any]:
        """Predictive analysis for future risks"""
        event_data = {
            'latency_p99': event.latency_p99,
            'error_rate': event.error_rate,
            'throughput': event.throughput,
            'cpu_util': event.cpu_util,
            'memory_util': event.memory_util
        }
        self.engine.add_telemetry(event.component, event_data)
        
        insights = self.engine.get_predictive_insights(event.component)
        
        return {
            'specialization': self.specialization.value,
            'confidence': 0.8 if insights['critical_risk_count'] > 0 else 0.5,
            'findings': insights,
            'recommendations': insights['recommendations']
        }

class OrchestrationManager:
    def __init__(self):
        self.agents = {
            AgentSpecialization.DETECTIVE: AnomalyDetectionAgent(),
            AgentSpecialization.DIAGNOSTICIAN: RootCauseAgent(),
            AgentSpecialization.PREDICTIVE: PredictiveAgent(),
        }
    
    async def orchestrate_analysis(self, event: ReliabilityEvent) -> Dict[str, Any]:
        agent_tasks = {
            spec: agent.analyze(event)
            for spec, agent in self.agents.items()
        }
        
        agent_results = {}
        for specialization, task in agent_tasks.items():
            try:
                result = await asyncio.wait_for(task, timeout=5.0)
                agent_results[specialization.value] = result
            except asyncio.TimeoutError:
                continue
        
        return self._synthesize_agent_findings(event, agent_results)
    
    def _synthesize_agent_findings(self, event: ReliabilityEvent, agent_results: Dict) -> Dict[str, Any]:
        detective_result = agent_results.get(AgentSpecialization.DETECTIVE.value)
        diagnostician_result = agent_results.get(AgentSpecialization.DIAGNOSTICIAN.value)
        predictive_result = agent_results.get(AgentSpecialization.PREDICTIVE.value)
        
        if not detective_result:
            return {'error': 'No agent results available'}
        
        synthesis = {
            'incident_summary': {
                'severity': detective_result['findings'].get('severity_tier', 'UNKNOWN'),
                'anomaly_confidence': detective_result['confidence'],
                'primary_metrics_affected': [metric["metric"] for metric in detective_result['findings'].get('primary_metrics_affected', [])]
            },
            'root_cause_insights': diagnostician_result['findings'] if diagnostician_result else {},
            'predictive_insights': predictive_result['findings'] if predictive_result else {},
            'recommended_actions': self._prioritize_actions(
                detective_result.get('recommendations', []),
                diagnostician_result.get('recommendations', []) if diagnostician_result else [],
                predictive_result.get('recommendations', []) if predictive_result else []
            ),
            'agent_metadata': {
                'participating_agents': list(agent_results.keys()),
                'analysis_timestamp': datetime.datetime.now().isoformat()
            }
        }
        
        return synthesis
    
    def _prioritize_actions(self, detection_actions: List[str], diagnosis_actions: List[str], predictive_actions: List[str]) -> List[str]:
        all_actions = detection_actions + diagnosis_actions + predictive_actions
        seen = set()
        unique_actions = []
        for action in all_actions:
            if action not in seen:
                seen.add(action)
                unique_actions.append(action)
        return unique_actions[:5]

# Initialize enhanced components
orchestration_manager = OrchestrationManager()

class EnhancedReliabilityEngine:
    def __init__(self):
        self.performance_metrics = {
            'total_incidents_processed': 0,
            'multi_agent_analyses': 0
        }
    
    async def process_event_enhanced(self, component: str, latency: float, error_rate: float,
                                   throughput: float = 1000, cpu_util: float = None,
                                   memory_util: float = None) -> Dict[str, Any]:
        
        event = ReliabilityEvent(
            component=component,
            latency_p99=latency,
            error_rate=error_rate,
            throughput=throughput,
            cpu_util=cpu_util,
            memory_util=memory_util,
            upstream_deps=["auth-service", "database"] if component == "api-service" else []
        )
        
        agent_analysis = await orchestration_manager.orchestrate_analysis(event)
        
        is_anomaly = anomaly_detector.detect_anomaly(event)

        agent_confidence = 0.0
        if agent_analysis and 'incident_summary' in agent_analysis:
            agent_confidence = agent_analysis.get('incident_summary', {}).get('anomaly_confidence', 0)
        else:
            agent_confidence = 0.8 if is_anomaly else 0.1

        if agent_confidence > 0.8:
            event.severity = EventSeverity.CRITICAL
        elif agent_confidence > 0.6:
            event.severity = EventSeverity.HIGH  
        elif agent_confidence > 0.4:
            event.severity = EventSeverity.MEDIUM
        else:
            event.severity = EventSeverity.LOW
        
        healing_actions = policy_engine.evaluate_policies(event)
        
        business_impact = business_calculator.calculate_impact(event) if is_anomaly else None
        
        if index is not None and is_anomaly:
            analysis_text = agent_analysis.get('recommended_actions', ['No analysis'])[0]
            vector_text = f"{component} {latency} {error_rate} {analysis_text}"
            vec = model.encode([vector_text])
            index.add(np.array(vec, dtype=np.float32))
            incident_texts.append(vector_text)
            save_index()
        
        result = {
            "timestamp": event.timestamp,
            "component": component,
            "latency_p99": latency,
            "error_rate": error_rate,
            "throughput": throughput,
            "status": "ANOMALY" if is_anomaly else "NORMAL",
            "multi_agent_analysis": agent_analysis,
            "healing_actions": [action.value for action in healing_actions],
            "business_impact": business_impact,
            "severity": event.severity.value,
            "similar_incidents_count": len(incident_texts) if is_anomaly else 0,
            "processing_metadata": {
                "agents_used": agent_analysis.get('agent_metadata', {}).get('participating_agents', []),
                "analysis_confidence": agent_analysis.get('incident_summary', {}).get('anomaly_confidence', 0)
            }
        }
        
        events_history.append(event)
        self.performance_metrics['total_incidents_processed'] += 1
        self.performance_metrics['multi_agent_analyses'] += 1
        
        return result

# Initialize enhanced engine
enhanced_engine = EnhancedReliabilityEngine()

def call_huggingface_analysis(prompt: str) -> str:
    if not HF_TOKEN:
        fallback_insights = [
            "High latency detected - possible resource contention or network issues",
            "Error rate increase suggests recent deployment instability", 
            "Latency spike correlates with increased user traffic patterns",
            "Intermittent failures indicate potential dependency service degradation",
            "Performance degradation detected - consider scaling compute resources"
        ]
        import random
        return random.choice(fallback_insights)

    try:
        enhanced_prompt = f"""
        As a senior reliability engineer, analyze this telemetry event and provide a concise root cause analysis:
        
        {prompt}
        
        Focus on:
        - Potential infrastructure or application issues
        - Correlation between metrics
        - Business impact assessment
        - Recommended investigation areas
        
        Provide 1-2 sentences maximum with actionable insights.
        """
        
        payload = {
            "model": "mistralai/Mixtral-8x7B-Instruct-v0.1",
            "prompt": enhanced_prompt,
            "max_tokens": 150,
            "temperature": 0.4,
        }
        response = requests.post(HF_API_URL, headers=HEADERS, json=payload, timeout=15)
        if response.status_code == 200:
            result = response.json()
            analysis_text = result.get("choices", [{}])[0].get("text", "").strip()
            if analysis_text and len(analysis_text) > 10:
                return analysis_text.split('\n')[0]
            return analysis_text
        else:
            return f"API Error {response.status_code}: Service temporarily unavailable"
    except Exception as e:
        return f"Analysis service error: {str(e)}"

# === Enhanced UI with Multi-Agent Insights ===
def create_enhanced_ui():
    with gr.Blocks(title="🧠 Enterprise Agentic Reliability Framework", theme="soft") as demo:
        gr.Markdown("""
        # 🧠 Enterprise Agentic Reliability Framework
        **Multi-Agent AI System for Production Reliability**
        
        *Specialized AI agents working together to detect, diagnose, predict, and heal system issues*
        """)
        
        with gr.Row():
            with gr.Column(scale=1):
                gr.Markdown("### 📊 Telemetry Input")
                component = gr.Dropdown(
                    choices=["api-service", "auth-service", "payment-service", "database", "cache-service"],
                    value="api-service",
                    label="Component",
                    info="Select the service being monitored"
                )
                latency = gr.Slider(
                    minimum=10, maximum=1000, value=100, step=1,
                    label="Latency P99 (ms)",
                    info="Alert threshold: >150ms (adaptive)"
                )
                error_rate = gr.Slider(
                    minimum=0, maximum=0.5, value=0.02, step=0.001,
                    label="Error Rate",
                    info="Alert threshold: >0.05"
                )
                throughput = gr.Number(
                    value=1000,
                    label="Throughput (req/sec)",
                    info="Current request rate"
                )
                cpu_util = gr.Slider(
                    minimum=0, maximum=1, value=0.4, step=0.01,
                    label="CPU Utilization",
                    info="0.0 - 1.0 scale"
                )
                memory_util = gr.Slider(
                    minimum=0, maximum=1, value=0.3, step=0.01,
                    label="Memory Utilization", 
                    info="0.0 - 1.0 scale"
                )
                submit_btn = gr.Button("🚀 Submit Telemetry Event", variant="primary", size="lg")
                
            with gr.Column(scale=2):
                gr.Markdown("### 🔍 Multi-Agent Analysis")
                output_text = gr.Textbox(
                    label="Agent Synthesis",
                    placeholder="AI agents are analyzing...",
                    lines=6
                )
                
                with gr.Accordion("🤖 Agent Specialists Analysis", open=False):
                    gr.Markdown("""
                    **Specialized AI Agents:**
                    - 🕵️ **Detective**: Anomaly detection & pattern recognition
                    - 🔍 **Diagnostician**: Root cause analysis & investigation  
                    - 🔮 **Predictive**: Future risk forecasting & trend analysis
                    """)
                    
                    agent_insights = gr.JSON(
                        label="Detailed Agent Findings",
                        value={}
                    )
                
                with gr.Accordion("🔮 Predictive Analytics & Forecasting", open=False):
                    gr.Markdown("""
                    **Future Risk Forecasting:**
                    - 📈 Latency trends and thresholds
                    - 🚨 Error rate predictions  
                    - 🔥 Resource utilization forecasts
                    - ⏰ Time-to-failure estimates
                    """)
                    
                    predictive_insights = gr.JSON(
                        label="Predictive Forecasts",
                        value={}
                    )
                
                gr.Markdown("### 📈 Recent Events (Last 15)")
                events_table = gr.Dataframe(
                    headers=["Timestamp", "Component", "Latency", "Error Rate", "Throughput", "Severity", "Analysis"],
                    label="Event History",
                    wrap=True,
                )
        
        with gr.Accordion("ℹ️ Framework Capabilities", open=False):
            gr.Markdown("""
            - **🤖 Multi-Agent AI**: Specialized agents for detection, diagnosis, prediction, and healing
            - **🔮 Predictive Analytics**: Forecast future risks and performance degradation
            - **🔧 Policy-Based Healing**: Automated recovery actions based on severity and context
            - **💰 Business Impact**: Revenue and user impact quantification
            - **🎯 Adaptive Detection**: ML-powered thresholds that learn from your environment
            - **📚 Vector Memory**: FAISS-based incident memory for similarity detection
            - **⚡ Production Ready**: Circuit breakers, cooldowns, and enterprise features
            """)
            
        with gr.Accordion("🔧 Healing Policies", open=False):
            policy_info = []
            for policy in policy_engine.policies:
                if policy.enabled:
                    actions = ", ".join([action.value for action in policy.actions])
                    policy_info.append(f"**{policy.name}**: {actions} (Priority: {policy.priority})")
            
            gr.Markdown("\n\n".join(policy_info))
        
        async def submit_event_enhanced(component, latency, error_rate, throughput, cpu_util, memory_util):
            try:
                latency = float(latency)
                error_rate = float(error_rate)
                throughput = float(throughput) if throughput else 1000
                cpu_util = float(cpu_util) if cpu_util else None
                memory_util = float(memory_util) if memory_util else None
                
                result = await enhanced_engine.process_event_enhanced(
                    component, latency, error_rate, throughput, cpu_util, memory_util
                )
                
                table_data = []
                for event in events_history[-15:]:
                    table_data.append([
                        event.timestamp[:19],
                        event.component,
                        event.latency_p99,
                        f"{event.error_rate:.3f}",
                        event.throughput,
                        event.severity.value.upper(),
                        "Multi-agent analysis" if 'multi_agent_analysis' in result else 'N/A'
                    ])
                
                status_emoji = "🚨" if result["status"] == "ANOMALY" else "✅"
                output_msg = f"{status_emoji} {result['status']}"
                
                if "multi_agent_analysis" in result:
                    analysis = result["multi_agent_analysis"]
                    confidence = analysis.get('incident_summary', {}).get('anomaly_confidence', 0)
                    output_msg += f"\n🎯 Confidence: {confidence*100:.1f}%"
                    
                    predictive_data = analysis.get('predictive_insights', {})
                    if predictive_data.get('critical_risk_count', 0) > 0:
                        output_msg += f"\n🔮 PREDICTIVE: {predictive_data['critical_risk_count']} critical risks forecast"
                    
                    if analysis.get('recommended_actions'):
                        output_msg += f"\n💡 Insights: {', '.join(analysis['recommended_actions'][:2])}"
                
                if result["business_impact"]:
                    impact = result["business_impact"]
                    output_msg += f"\n💰 Business Impact: ${impact['revenue_loss_estimate']} | 👥 {impact['affected_users_estimate']} users | 🚨 {impact['severity_level']}"
                
                if result["healing_actions"] and result["healing_actions"] != ["no_action"]:
                    actions = ", ".join(result["healing_actions"])
                    output_msg += f"\n🔧 Auto-Actions: {actions}"
                
                agent_insights_data = result.get("multi_agent_analysis", {})
                predictive_insights_data = agent_insights_data.get('predictive_insights', {})
                
                return (
                    output_msg,
                    agent_insights_data,
                    predictive_insights_data,
                    gr.Dataframe(
                        headers=["Timestamp", "Component", "Latency", "Error Rate", "Throughput", "Severity", "Analysis"],
                        value=table_data,
                        wrap=True
                    )
                )
                
            except Exception as e:
                return f"❌ Error processing event: {str(e)}", {}, {}, gr.Dataframe(value=[])
        
        submit_btn.click(
            fn=submit_event_enhanced,
            inputs=[component, latency, error_rate, throughput, cpu_util, memory_util],
            outputs=[output_text, agent_insights, predictive_insights, events_table]
        )
    
    return demo

if __name__ == "__main__":
    demo = create_enhanced_ui()
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False
    )