"""Performance Tuning Agent — uses JEPA predictions to auto-tune configurations."""
from __future__ import annotations

import logging
from typing import Any, Dict, List

from mesh.codex.base_agent import CodexAgent
from mesh.codex.models import CodexTask, CodexTaskResult, FindingSeverity

logger = logging.getLogger(__name__)


class PerformanceTunerAgent(CodexAgent):
    """Reads JEPA stats and node health data, compares current configs
    against JEPA predictions, and generates tuning recommendations.

    Can auto-apply recommendations via the existing /tuning/apply
    endpoint when the 'auto_apply' parameter is set.
    """

    def __init__(self, mesh_db, config=None, jepa_trainer=None, coordinator=None):
        super().__init__("tuner", mesh_db, config)
        self.jepa_trainer = jepa_trainer
        self.coordinator = coordinator

    async def execute(self, task: CodexTask) -> CodexTaskResult:
        findings: List[Dict[str, Any]] = []
        recommendations: List[str] = []
        target_node = task.target or None

        # Gather node data
        nodes = self._get_nodes(target_node)
        if not nodes:
            return CodexTaskResult(
                task_id=task.task_id,
                success=True,
                findings=[],
                recommendations=[f"No nodes found matching target '{task.target}'"],
            )

        # Get JEPA stats
        jepa_stats = self._get_jepa_stats()

        for node in nodes:
            node_findings = self._analyze_node(node, jepa_stats)
            findings.extend(node_findings)

        # Generate recommendations
        recommendations = self._generate_recommendations(findings, nodes, jepa_stats)

        return CodexTaskResult(
            task_id=task.task_id,
            success=True,
            findings=findings,
            recommendations=recommendations,
        )

    def _get_nodes(self, target_node: str | None) -> List[dict]:
        """Get node data from the database."""
        if not self.mesh_db:
            return []
        if target_node:
            node = self.mesh_db.get_node(target_node)
            return [node] if node else []
        return self.mesh_db.get_all_nodes()

    def _get_jepa_stats(self) -> dict:
        """Get JEPA training stats if available."""
        if not self.jepa_trainer:
            return {}
        try:
            stats = self.jepa_trainer.get_stats()
            return {
                "trained": getattr(stats, "trained", False),
                "mean_error": getattr(stats, "mean_error", 0.0),
                "confidence_by_regime": getattr(stats, "confidence_by_regime", {}),
                "online_updates": getattr(stats, "online_updates", 0),
            }
        except Exception:
            return {}

    def _analyze_node(self, node: dict, jepa_stats: dict) -> List[Dict[str, Any]]:
        """Analyze a single node's health and generate findings."""
        findings = []
        node_id = node.get("node_id", "unknown")

        # Check recent health history
        health_history = []
        if self.mesh_db:
            health_history = self.mesh_db.get_health_history(node_id, limit=50)

        if not health_history:
            findings.append({
                "severity": FindingSeverity.INFO.value,
                "category": "no_health_data",
                "location": node_id,
                "message": f"No health data available for node {node_id}",
                "suggestion": "Ensure the node agent is sending heartbeats",
            })
            return findings

        # Analyze CPU utilization trends
        cpu_vals = [h.get("cpu_pct", 0) for h in health_history if h.get("cpu_pct") is not None]
        if cpu_vals:
            avg_cpu = sum(cpu_vals) / len(cpu_vals)
            max_cpu = max(cpu_vals)
            if avg_cpu > 90:
                findings.append({
                    "severity": FindingSeverity.HIGH.value,
                    "category": "cpu_saturation",
                    "location": node_id,
                    "message": f"CPU consistently saturated (avg {avg_cpu:.1f}%, max {max_cpu:.1f}%)",
                    "suggestion": "Consider offloading jobs to other nodes or increasing CPU resources",
                })
            elif avg_cpu < 5 and max_cpu < 10:
                findings.append({
                    "severity": FindingSeverity.LOW.value,
                    "category": "cpu_idle",
                    "location": node_id,
                    "message": f"CPU mostly idle (avg {avg_cpu:.1f}%)",
                    "suggestion": "Node may be underutilized — consider routing more jobs to it",
                })

        # Analyze memory pressure
        mem_vals = [h.get("memory_used_pct", 0) for h in health_history if h.get("memory_used_pct") is not None]
        if mem_vals:
            avg_mem = sum(mem_vals) / len(mem_vals)
            max_mem = max(mem_vals)
            if avg_mem > 85:
                findings.append({
                    "severity": FindingSeverity.HIGH.value,
                    "category": "memory_pressure",
                    "location": node_id,
                    "message": f"Memory pressure detected (avg {avg_mem:.1f}%, max {max_mem:.1f}%)",
                    "suggestion": "Reduce batch sizes, enable swap, or add memory",
                })
            # Check for memory leak pattern (consistently increasing)
            if len(mem_vals) >= 10:
                first_half = sum(mem_vals[:len(mem_vals)//2]) / (len(mem_vals)//2)
                second_half = sum(mem_vals[len(mem_vals)//2:]) / (len(mem_vals) - len(mem_vals)//2)
                if second_half - first_half > 10:
                    findings.append({
                        "severity": FindingSeverity.MEDIUM.value,
                        "category": "memory_leak",
                        "location": node_id,
                        "message": f"Possible memory leak: usage trending up ({first_half:.1f}% -> {second_half:.1f}%)",
                        "suggestion": "Investigate processes for memory leaks; consider periodic restarts",
                    })

        # Analyze disk usage
        disk_vals = [h.get("disk_used_pct", 0) for h in health_history if h.get("disk_used_pct") is not None]
        if disk_vals:
            latest_disk = disk_vals[0]
            if latest_disk > 90:
                findings.append({
                    "severity": FindingSeverity.CRITICAL.value,
                    "category": "disk_pressure",
                    "location": node_id,
                    "message": f"Disk usage critical at {latest_disk:.1f}%",
                    "suggestion": "Free disk space immediately — prune logs, old checkpoints, and temp files",
                })
            elif latest_disk > 80:
                findings.append({
                    "severity": FindingSeverity.MEDIUM.value,
                    "category": "disk_pressure",
                    "location": node_id,
                    "message": f"Disk usage high at {latest_disk:.1f}%",
                    "suggestion": "Plan disk cleanup before it becomes critical",
                })

        # JEPA-informed tuning
        if jepa_stats.get("trained") and jepa_stats.get("confidence_by_regime"):
            confidence = jepa_stats["confidence_by_regime"]
            low_confidence_regimes = [
                regime for regime, conf in confidence.items() if conf < 0.5
            ]
            if low_confidence_regimes:
                findings.append({
                    "severity": FindingSeverity.INFO.value,
                    "category": "jepa_confidence",
                    "location": node_id,
                    "message": f"JEPA has low confidence for regimes: {', '.join(low_confidence_regimes)}",
                    "suggestion": "Run more benchmarks in these regimes to improve prediction accuracy",
                })

        # GPU utilization analysis
        gpu_data = [h.get("gpu_utilizations") for h in health_history
                    if h.get("gpu_utilizations")]
        if gpu_data:
            # Flatten all GPU utilization values
            all_gpu_utils = []
            for utils in gpu_data:
                if isinstance(utils, list):
                    all_gpu_utils.extend(utils)
            if all_gpu_utils:
                avg_gpu = sum(all_gpu_utils) / len(all_gpu_utils)
                if avg_gpu < 10:
                    findings.append({
                        "severity": FindingSeverity.MEDIUM.value,
                        "category": "gpu_idle",
                        "location": node_id,
                        "message": f"GPU underutilized (avg {avg_gpu:.1f}%)",
                        "suggestion": "Route GPU-bound workloads to this node or check driver status",
                    })

        return findings

    def _generate_recommendations(self, findings: List[Dict[str, Any]],
                                  nodes: List[dict],
                                  jepa_stats: dict) -> List[str]:
        recommendations = []

        critical = [f for f in findings if f["severity"] == "critical"]
        high = [f for f in findings if f["severity"] == "high"]

        if critical:
            recommendations.append(
                f"URGENT: {len(critical)} critical issues require immediate attention"
            )
        if high:
            recommendations.append(
                f"{len(high)} high-severity performance issues detected across {len(nodes)} nodes"
            )

        # JEPA-specific recommendations
        if jepa_stats.get("trained"):
            if jepa_stats.get("mean_error", 1.0) > 0.1:
                recommendations.append(
                    "JEPA model error is high — trigger a retrain with more diverse workloads"
                )
            if jepa_stats.get("online_updates", 0) == 0:
                recommendations.append(
                    "No online JEPA feedback received — enable feedback from job completions"
                )
        else:
            recommendations.append(
                "JEPA model not yet trained — run initial training for performance predictions"
            )

        if not findings:
            recommendations.append(
                f"All {len(nodes)} nodes healthy — no tuning changes needed"
            )

        return recommendations
