"""
Real-time monitoring and analytics dashboard
"""
from fastapi import APIRouter, WebSocket, WebSocketDisconnect
from fastapi.responses import HTMLResponse
import asyncio
import json
import time
from typing import List, Dict, Any
import logging
from datetime import datetime, timedelta
import psutil
import torch

logger = logging.getLogger(__name__)
router = APIRouter()

class ConnectionManager:
    def __init__(self):
        self.active_connections: List[WebSocket] = []
        self.stats = {
            'total_predictions': 0,
            'hallucinations_detected': 0,
            'average_confidence': 0.0,
            'average_response_time': 0.0,
            'method_usage': {},
            'hourly_stats': [],
            'error_count': 0
        }
        self.recent_predictions = []

    async def connect(self, websocket: WebSocket):
        await websocket.accept()
        self.active_connections.append(websocket)
        # Send current stats immediately
        await self.send_stats_update()

    def disconnect(self, websocket: WebSocket):
        self.active_connections.remove(websocket)

    async def send_personal_message(self, message: str, websocket: WebSocket):
        await websocket.send_text(message)

    async def broadcast(self, message: dict):
        for connection in self.active_connections:
            try:
                await connection.send_text(json.dumps(message))
            except Exception as e:
                logger.error(f"Failed to send message to websocket: {e}")

    def update_stats(self, prediction_result: Dict[str, Any], response_time: float):
        """Update statistics with new prediction result"""
        try:
            self.stats['total_predictions'] += 1
            
            if prediction_result.get('is_hallucination', False):
                self.stats['hallucinations_detected'] += 1
            
            # Update running averages
            confidence = prediction_result.get('confidence', 0.5)
            total = self.stats['total_predictions']
            
            self.stats['average_confidence'] = (
                (self.stats['average_confidence'] * (total - 1) + confidence) / total
            )
            
            self.stats['average_response_time'] = (
                (self.stats['average_response_time'] * (total - 1) + response_time) / total
            )
            
            # Track method usage
            method = prediction_result.get('method', 'unknown')
            self.stats['method_usage'][method] = self.stats['method_usage'].get(method, 0) + 1
            
            # Store recent prediction (keep last 100)
            self.recent_predictions.append({
                'timestamp': datetime.now().isoformat(),
                'is_hallucination': prediction_result.get('is_hallucination', False),
                'confidence': confidence,
                'method': method,
                'response_time': response_time
            })
            
            if len(self.recent_predictions) > 100:
                self.recent_predictions.pop(0)
            
            # Update hourly stats
            self._update_hourly_stats()
            
        except Exception as e:
            logger.error(f"Failed to update stats: {e}")
            self.stats['error_count'] += 1

    def _update_hourly_stats(self):
        """Update hourly statistics"""
        current_hour = datetime.now().replace(minute=0, second=0, microsecond=0)
        
        # Find or create current hour entry
        current_hour_entry = None
        for entry in self.stats['hourly_stats']:
            if entry['hour'] == current_hour.isoformat():
                current_hour_entry = entry
                break
        
        if not current_hour_entry:
            current_hour_entry = {
                'hour': current_hour.isoformat(),
                'predictions': 0,
                'hallucinations': 0,
                'avg_confidence': 0.0,
                'avg_response_time': 0.0
            }
            self.stats['hourly_stats'].append(current_hour_entry)
        
        # Update current hour
        current_hour_entry['predictions'] += 1
        
        # Keep only last 24 hours
        cutoff_time = current_hour - timedelta(hours=24)
        self.stats['hourly_stats'] = [
            entry for entry in self.stats['hourly_stats']
            if datetime.fromisoformat(entry['hour']) >= cutoff_time
        ]

    async def send_stats_update(self):
        """Send current statistics to all connected clients"""
        try:
            # Get system stats
            system_stats = self._get_system_stats()
            
            message = {
                'type': 'stats_update',
                'data': {
                    **self.stats,
                    'system': system_stats,
                    'recent_predictions': self.recent_predictions[-10:],  # Last 10
                    'timestamp': datetime.now().isoformat()
                }
            }
            
            await self.broadcast(message)
        except Exception as e:
            logger.error(f"Failed to send stats update: {e}")

    def _get_system_stats(self) -> Dict[str, Any]:
        """Get system performance statistics"""
        try:
            return {
                'cpu_percent': psutil.cpu_percent(),
                'memory_percent': psutil.virtual_memory().percent,
                'gpu_memory_used': torch.cuda.memory_allocated() / 1024**3 if torch.cuda.is_available() else 0,
                'gpu_memory_total': torch.cuda.get_device_properties(0).total_memory / 1024**3 if torch.cuda.is_available() else 0
            }
        except Exception as e:
            logger.error(f"Failed to get system stats: {e}")
            return {}

# Global connection manager
manager = ConnectionManager()

@router.websocket("/ws/monitor")
async def websocket_endpoint(websocket: WebSocket):
    await manager.connect(websocket)
    try:
        while True:
            # Send periodic updates every 5 seconds
            await asyncio.sleep(5)
            await manager.send_stats_update()
    except WebSocketDisconnect:
        manager.disconnect(websocket)

@router.get("/monitor", response_class=HTMLResponse)
async def get_monitor_dashboard():
    """Serve the monitoring dashboard"""
    html_content = """
    <!DOCTYPE html>
    <html>
    <head>
        <title>Hallucination Detection Monitor</title>
        <script src="https://cdn.jsdelivr.net/npm/chart.js"></script>
        <style>
            body { font-family: Arial, sans-serif; margin: 20px; background: #f5f5f5; }
            .dashboard { display: grid; grid-template-columns: 1fr 1fr; gap: 20px; }
            .widget { background: white; padding: 20px; border-radius: 8px; box-shadow: 0 2px 4px rgba(0,0,0,0.1); }
            .metric { text-align: center; margin: 10px 0; }
            .metric-value { font-size: 2em; font-weight: bold; color: #007bff; }
            .metric-label { color: #666; text-transform: uppercase; font-size: 0.8em; }
            .status-good { color: #28a745; }
            .status-warning { color: #ffc107; }
            .status-error { color: #dc3545; }
            .recent-predictions { max-height: 300px; overflow-y: auto; }
            .prediction-item { padding: 5px; border-bottom: 1px solid #eee; font-size: 0.9em; }
            .hallucination { background-color: #ffe6e6; }
            .normal { background-color: #e6ffe6; }
            h1, h2 { color: #333; }
            .system-stats { display: flex; justify-content: space-between; }
            .system-stat { text-align: center; }
        </style>
    </head>
    <body>
        <h1>🔍 Hallucination Detection System - Real-time Monitor</h1>
        
        <div class="dashboard">
            <div class="widget">
                <h2>📊 Key Metrics</h2>
                <div class="metric">
                    <div class="metric-value" id="total-predictions">0</div>
                    <div class="metric-label">Total Predictions</div>
                </div>
                <div class="metric">
                    <div class="metric-value" id="hallucination-rate">0%</div>
                    <div class="metric-label">Hallucination Rate</div>
                </div>
                <div class="metric">
                    <div class="metric-value" id="avg-confidence">0%</div>
                    <div class="metric-label">Avg Confidence</div>
                </div>
                <div class="metric">
                    <div class="metric-value" id="avg-response-time">0ms</div>
                    <div class="metric-label">Avg Response Time</div>
                </div>
            </div>

            <div class="widget">
                <h2>🖥️ System Health</h2>
                <div class="system-stats">
                    <div class="system-stat">
                        <div class="metric-value" id="cpu-usage">0%</div>
                        <div class="metric-label">CPU Usage</div>
                    </div>
                    <div class="system-stat">
                        <div class="metric-value" id="memory-usage">0%</div>
                        <div class="metric-label">Memory Usage</div>
                    </div>
                    <div class="system-stat">
                        <div class="metric-value" id="gpu-usage">0%</div>
                        <div class="metric-label">GPU Memory</div>
                    </div>
                </div>
            </div>

            <div class="widget">
                <h2>📈 Prediction Methods</h2>
                <canvas id="methods-chart" width="400" height="200"></canvas>
            </div>

            <div class="widget">
                <h2>🕒 Recent Predictions</h2>
                <div class="recent-predictions" id="recent-predictions">
                    <!-- Recent predictions will be populated here -->
                </div>
            </div>
        </div>

        <script>
            const ws = new WebSocket('ws://localhost:8000/ws/monitor');
            let methodsChart;

            ws.onmessage = function(event) {
                const data = JSON.parse(event.data);
                if (data.type === 'stats_update') {
                    updateDashboard(data.data);
                }
            };

            function updateDashboard(stats) {
                // Update key metrics
                document.getElementById('total-predictions').textContent = stats.total_predictions || 0;
                
                const hallucinationRate = stats.total_predictions > 0 ? 
                    ((stats.hallucinations_detected / stats.total_predictions) * 100).toFixed(1) : 0;
                document.getElementById('hallucination-rate').textContent = hallucinationRate + '%';
                
                document.getElementById('avg-confidence').textContent = 
                    ((stats.average_confidence || 0) * 100).toFixed(1) + '%';
                
                document.getElementById('avg-response-time').textContent = 
                    ((stats.average_response_time || 0) * 1000).toFixed(0) + 'ms';

                // Update system health
                if (stats.system) {
                    document.getElementById('cpu-usage').textContent = 
                        (stats.system.cpu_percent || 0).toFixed(1) + '%';
                    document.getElementById('memory-usage').textContent = 
                        (stats.system.memory_percent || 0).toFixed(1) + '%';
                    
                    const gpuUsage = stats.system.gpu_memory_total > 0 ? 
                        ((stats.system.gpu_memory_used / stats.system.gpu_memory_total) * 100).toFixed(1) : 0;
                    document.getElementById('gpu-usage').textContent = gpuUsage + '%';
                }

                // Update methods chart
                updateMethodsChart(stats.method_usage || {});

                // Update recent predictions
                updateRecentPredictions(stats.recent_predictions || []);
            }

            function updateMethodsChart(methodUsage) {
                const ctx = document.getElementById('methods-chart').getContext('2d');
                
                if (methodsChart) {
                    methodsChart.destroy();
                }

                const labels = Object.keys(methodUsage);
                const data = Object.values(methodUsage);

                methodsChart = new Chart(ctx, {
                    type: 'doughnut',
                    data: {
                        labels: labels,
                        datasets: [{
                            data: data,
                            backgroundColor: [
                                '#007bff', '#28a745', '#ffc107', '#dc3545', '#6f42c1', '#fd7e14'
                            ]
                        }]
                    },
                    options: {
                        responsive: true,
                        plugins: {
                            legend: {
                                position: 'bottom'
                            }
                        }
                    }
                });
            }

            function updateRecentPredictions(predictions) {
                const container = document.getElementById('recent-predictions');
                container.innerHTML = '';

                predictions.reverse().forEach(pred => {
                    const item = document.createElement('div');
                    item.className = `prediction-item ${pred.is_hallucination ? 'hallucination' : 'normal'}`;
                    
                    const time = new Date(pred.timestamp).toLocaleTimeString();
                    const status = pred.is_hallucination ? '❌ Hallucination' : '✅ Valid';
                    const confidence = (pred.confidence * 100).toFixed(1);
                    const responseTime = (pred.response_time * 1000).toFixed(0);
                    
                    item.innerHTML = `
                        <strong>${time}</strong> - ${status} 
                        (${confidence}% confidence, ${responseTime}ms, ${pred.method})
                    `;
                    
                    container.appendChild(item);
                });
            }

            // Connect to WebSocket
            ws.onopen = function() {
                console.log('Connected to monitoring WebSocket');
            };

            ws.onerror = function(error) {
                console.error('WebSocket error:', error);
            };
        </script>
    </body>
    </html>
    """
    return HTMLResponse(content=html_content)

@router.get("/api/monitor/stats")
async def get_monitor_stats():
    """Get current monitoring statistics"""
    return {
        "stats": manager.stats,
        "system": manager._get_system_stats(),
        "recent_predictions": manager.recent_predictions[-20:]
    }

@router.get("/monitor/metrics")
async def get_metrics():
    """Get system metrics and performance data"""
    system_stats = manager._get_system_stats()
    return {
        "system": system_stats,
        "predictions": {
            "total": manager.stats['total_predictions'],
            "hallucinations": manager.stats['hallucinations_detected'],
            "accuracy_rate": (manager.stats['total_predictions'] - manager.stats['hallucinations_detected']) / max(manager.stats['total_predictions'], 1),
            "average_confidence": manager.stats['average_confidence'],
            "average_response_time": manager.stats['average_response_time']
        },
        "status": "healthy",
        "timestamp": datetime.now().isoformat()
    }