from typing import Dict, Any


# ============================================================
# THRESHOLDS  (tune once, change everywhere)
# ============================================================
CPU_HIGH       = 80    # %
CPU_LOW        = 20    # %
GPU_HIGH       = 85    # %
GPU_LOW        = 15    # %
TEMP_WARN      = 80    # °C
TEMP_CRITICAL  = 90    # °C
RAM_WARN       = 80    # %
RAM_CRITICAL   = 90    # %
GPU_MEM_WARN   = 85    # % of vRAM used
CO2_WARN_G     = 5.0   # g CO₂ per cell run – flag heavy cells
ENERGY_WARN_WH = 0.003 # Wh per cell – flag energy-hungry cells
RUNTIME_SLOW_S = 30    # s  – flag slow cells
RUNTIME_VERY_SLOW=120  # s  – escalate recommendation


class RecommendationEngine:

    def __init__(self, tracker=None, gemini=None):
        self.tracker = tracker
        self.gemini  = gemini

    # ==========================================================
    def recommend(self, m: Dict[str, Any], code=None, error=None):
        """
        Multi-signal recommendation engine.

        Returns dict with:
            recommended_hardware    – short label
            recommended_confidence  – 0.0–1.0
            recommended_reasons     – pipe-separated advice string
        """

        suggestions = []

        # ── pull metrics (safe float conversion) ──────────────
        cpu      = float(m.get("cpu_util_avg",       0) or 0)
        gpu      = float(m.get("gpu_util_avg",       0) or 0)
        cpu_temp = float(m.get("cpu_temp_c",         0) or 0)
        gpu_temp = float(m.get("gpu_temp_c",         0) or 0)
        ram_pct  = float(m.get("ram_used_pct",       0) or 0)
        gpu_mem_pct = float(m.get("gpu_mem_used_pct",0) or 0)
        runtime  = float(m.get("runtime_sec",        0) or 0)
        energy_wh= float(m.get("energy_kwh",         0) or 0) * 1000  # convert to Wh
        co2_g    = float(m.get("co2_g",              0) or 0)
        hw_type  = str(m.get("hardware_type", "CPU"))

        # ── choose a base label ───────────────────────────────
        label = "Balanced – Continue Current Hardware"
        flags = 0                                     # count of active problems

        # ══════════════════════════════════════════════════════
        # RULE 1 – CPU Bottleneck (CPU busy, GPU idle)
        # ══════════════════════════════════════════════════════
        if cpu > CPU_HIGH and gpu < GPU_LOW:
            label = "CPU Bottleneck – Offload to GPU"
            suggestions.append(
                "⚠ CPU saturated while GPU is idle → "
                "migrate compute to GPU (e.g., move tensors to .cuda() or use cuDF)"
            )
            flags += 2

        # ══════════════════════════════════════════════════════
        # RULE 2 – GPU Overloaded
        # ══════════════════════════════════════════════════════
        if gpu > GPU_HIGH:
            label = "GPU Overloaded – Reduce Batch Size or Upgrade"
            suggestions.append(
                f"⚠ GPU utilisation {gpu:.1f}% > {GPU_HIGH}% → "
                "reduce batch size, apply mixed-precision (fp16), or upgrade GPU"
            )
            flags += 2

        # ══════════════════════════════════════════════════════
        # RULE 3 – Balanced GPU (healthy zone)
        # ══════════════════════════════════════════════════════
        if GPU_LOW < gpu < GPU_HIGH and hw_type == "GPU" and cpu < CPU_HIGH:
            label = "GPU Efficient – Continue GPU Usage"
            suggestions.append(
                f"✅ GPU at {gpu:.1f}% (healthy zone) – "
                "maintain current configuration for optimal throughput"
            )

        # ══════════════════════════════════════════════════════
        # RULE 4 – GPU Underutilised (has GPU but barely using it)
        # ══════════════════════════════════════════════════════
        if hw_type == "GPU" and gpu < GPU_LOW and cpu > CPU_HIGH:
            suggestions.append(
                f"💤 GPU almost idle ({gpu:.1f}%) while CPU is busy → "
                "check whether your libraries are using CUDA correctly"
            )
            flags += 1

        # ══════════════════════════════════════════════════════
        # RULE 5a – CPU-only environment (no GPU in use at all)
        #  → always suggest GPU for ML workloads regardless of CPU %
        # ══════════════════════════════════════════════════════
        if hw_type == "CPU":
            label = "CPU Only – Consider GPU for ML Workloads"
            suggestions.append(
                "💡 No GPU detected/active → GPU acceleration could make training "
                f"~10× faster and reduce CO₂ by up to 70% for deep-learning workloads"
            )
            flags += 1

        # ══════════════════════════════════════════════════════
        # RULE 5b – CPU-only AND overloaded
        # ══════════════════════════════════════════════════════
        if hw_type == "CPU" and cpu > CPU_HIGH:
            label = "CPU Bottleneck – Enable GPU Immediately"
            suggestions.append(
                f"⚠ CPU at {cpu:.1f}% with no GPU → "
                "enable CUDA (PyTorch) or TPU (TF/JAX) to offload compute"
            )
            flags += 2

        # ══════════════════════════════════════════════════════
        # RULE 6 – CPU Temperature warning
        # ══════════════════════════════════════════════════════
        if cpu_temp >= TEMP_CRITICAL:
            suggestions.append(
                f"🔥 CPU critically hot ({cpu_temp:.1f}°C) → "
                "check cooling, reduce CPU-bound threads, or throttle back"
            )
            flags += 2
        elif cpu_temp >= TEMP_WARN:
            suggestions.append(
                f"🌡 CPU warm ({cpu_temp:.1f}°C) → "
                "monitor cooling; consider reducing parallelism"
            )
            flags += 1

        # ══════════════════════════════════════════════════════
        # RULE 7 – GPU Temperature warning
        # ══════════════════════════════════════════════════════
        if gpu_temp >= TEMP_CRITICAL:
            suggestions.append(
                f"🔥 GPU critically hot ({gpu_temp:.1f}°C) → "
                "reduce batch size or pause training to cool down"
            )
            flags += 2
        elif gpu_temp >= TEMP_WARN:
            suggestions.append(
                f"🌡 GPU warm ({gpu_temp:.1f}°C) → "
                "ensure case airflow and check fan RPM"
            )
            flags += 1

        # ══════════════════════════════════════════════════════
        # RULE 8 – RAM pressure
        # ══════════════════════════════════════════════════════
        if ram_pct >= RAM_CRITICAL:
            label = "Memory Pressure – Reduce Data Size"
            suggestions.append(
                f"💾 System RAM critically high ({ram_pct:.1f}%) → "
                "use chunked loading (chunksize=), delete unneeded variables, or upgrade RAM"
            )
            flags += 2
        elif ram_pct >= RAM_WARN:
            suggestions.append(
                f"💾 System RAM high ({ram_pct:.1f}%) → "
                "consider reducing dataset size or freeing intermediate objects"
            )
            flags += 1

        # ══════════════════════════════════════════════════════
        # RULE 9 – GPU VRAM pressure
        # ══════════════════════════════════════════════════════
        if gpu_mem_pct >= GPU_MEM_WARN:
            suggestions.append(
                f"📦 GPU VRAM almost full ({gpu_mem_pct:.1f}%) → "
                "use gradient checkpointing, reduce batch size, or move to CPU offload"
            )
            flags += 1

        # ══════════════════════════════════════════════════════
        # RULE 10 – Slow runtime
        # ══════════════════════════════════════════════════════
        if runtime >= RUNTIME_VERY_SLOW:
            suggestions.append(
                f"⏳ Cell took {runtime:.1f}s → "
                "profile with %%prun or torch.profiler; consider vectorisation or batching"
            )
            flags += 1
        elif runtime >= RUNTIME_SLOW_S:
            suggestions.append(
                f"⏱ Slow cell ({runtime:.1f}s) → "
                "check for Python loops that could be vectorised"
            )

        # ══════════════════════════════════════════════════════
        # RULE 11 – High CO₂ / energy footprint
        # ══════════════════════════════════════════════════════
        if co2_g >= CO2_WARN_G:
            suggestions.append(
                f"♻ High carbon footprint ({co2_g:.3f}g CO₂) → "
                "reduce epochs, enable early stopping, or use a smaller model"
            )
            flags += 1
        if energy_wh >= ENERGY_WARN_WH:
            suggestions.append(
                f"⚡ Energy-heavy cell ({energy_wh*1000:.3f} mWh) → "
                "consider quantisation (int8) or pruning to reduce power draw"
            )

        # ══════════════════════════════════════════════════════
        # RULE 12 – Runtime error
        # ══════════════════════════════════════════════════════
        if error:
            import re
            err_str = str(error)
            
            # Intelligent fallback for missing modules
            if "No module named" in err_str:
                match = re.search(r"No module named\s+'([^']+)'", err_str)
                mod = match.group(1) if match else "library"
                suggestions.append(f"❌ Missing Library → Run <code>!pip install {mod}</code> in a new cell to fix this error.")
            elif type(error).__name__ in ["NameError", "AttributeError"]:
                suggestions.append(f"❌ Code Error ({type(error).__name__}) → {err_str[:80]}")
            else:
                suggestions.append(f"❌ Runtime error: {str(error)[:120]}")
                
            flags += 1

        # ══════════════════════════════════════════════════════
        # CONFIDENCE  (linear scale based on how many flags)
        # ══════════════════════════════════════════════════════
        if flags == 0:
            confidence = 1.0
        else:
            confidence = round(min(1.0, 0.45 + flags * 0.15), 2)

        # ══════════════════════════════════════════════════════
        # 🧠 GEMINI – enrich only when something is wrong
        # ══════════════════════════════════════════════════════
        problem_detected = (
            flags > 0 or
            error is not None
        )

        eco_fix_before = None
        eco_fix_after = None

        # ── OFFLINE FALLBACK: ML Parameter Optimizer (n_estimators) ──
        if "n_estimators" in (code or "") and co2_g >= 0.1 and not eco_fix_before:
            import re as _re
            _est_m = _re.search(r"n_estimators\s*=\s*(\d+)", code)
            if _est_m:
                _old_val = int(_est_m.group(1))
                if _old_val > 100:
                    _new_val = max(50, _old_val // 2)
                    eco_fix_before = _est_m.group(0)
                    eco_fix_after  = f"n_estimators={_new_val}"
                    suggestions.append(f"🌱 <strong>Eco Fix Available</strong> → Reduce <code>n_estimators</code> from {_old_val} to {_new_val} to save ~50% energy with minimal impact on accuracy.")

        if self.gemini and self.gemini.enabled and problem_detected:
            ai_data = self.gemini.ask_rich(m, code=code, error=error)
            if ai_data and isinstance(ai_data, dict):
                # Set code fix if Gemini returned before/after
                if "before" in ai_data and "after" in ai_data:
                    eco_fix_before = ai_data["before"]
                    eco_fix_after  = ai_data["after"]

                # Always render structured 4-part card if detection key present
                if "detection" in ai_data:
                    tip_html = f"""
                    <div style='display:flex; flex-direction:column; margin-top:10px; font-size:13px; line-height:1.4;'>
                      <div style='color:var(--color-red); font-weight:700; font-size:14px; margin-bottom:6px;'>
                        🌱 Eco Suggestion
                      </div>
                      <div style='color:var(--text-highlight); font-weight:600; margin-bottom:4px; display:flex; align-items:center;'>
                        <span style='margin-right:6px;'>⚠️</span> {ai_data.get('detection', '')}
                      </div>
                      <div style='color:var(--text-muted); margin-left:14px; padding-left:12px; border-left:2px solid #334155; margin-bottom:12px; font-size:12px;'>
                        ↳ {ai_data.get('reason', '')}
                      </div>
                      <div style='margin-bottom:8px; display:flex; align-items:flex-start;'>
                        <span style='color:#60a5fa; font-weight:700; margin-right:6px;'>👉 Action:</span> 
                        <span style='color:var(--text-highlight);'>{ai_data.get('action', '')}</span>
                      </div>
                      <div style='display:flex; align-items:flex-start;'>
                        <span style='color:#4ade80; font-weight:700; margin-right:6px;'>✨ Impact:</span> 
                        <span style='color:#a7f3d0;'>{ai_data.get('impact', '')}</span>
                      </div>
                    </div>
                    """
                    suggestions.insert(0, tip_html)
                else:
                    # Unstructured fallback
                    tip_str = ai_data.get("tip", str(ai_data))
                    if not tip_str.startswith("Gemini API Error"):
                        suggestions.insert(0, f"<div style='color:var(--color-red); font-weight:700; margin-bottom:4px;'>🌱 Eco Suggestion</div><div style='margin-left:10px;'>{tip_str}</div>")
            elif ai_data and isinstance(ai_data, str):
                if not ai_data.startswith("Gemini API Error"):
                    suggestions.insert(0, f"<div style='color:var(--color-red); font-weight:700; margin-bottom:4px;'>🌱 Eco Suggestion</div><div style='margin-left:10px;'>{ai_data}</div>")

        # ══════════════════════════════════════════════════════
        # COMPARISON DATA  (CPU vs GPU estimates)
        # ══════════════════════════════════════════════════════
        comparison = _build_comparison(m)

        # ══════════════════════════════════════════════════════
        # NO PROBLEMS DETECTED
        # ══════════════════════════════════════════════════════
        if not suggestions:
            return {
                "recommended_hardware":    "Optimal",
                "recommended_confidence":  1.0,
                "recommended_reasons":     "✅ No issues detected – system is running efficiently",
                "comparison":              comparison,
                "eco_fix_before":          None,
                "eco_fix_after":           None,
            }

        return {
            "recommended_hardware":    label,
            "recommended_confidence":  confidence,
            "recommended_reasons":     " | ".join(suggestions),
            "comparison":              comparison,
            "eco_fix_before":          eco_fix_before,
            "eco_fix_after":           eco_fix_after,
        }


# ============================================================
# COMPARISON ESTIMATOR
# ============================================================
# Speedup assumptions (evidence-based heuristics)
# Image classification / CNN training  → GPU ~10× faster
# General ML (sklearn, pandas)        → GPU not helpful (~1×)
# Deep learning inference             → GPU ~5× faster
_GPU_SPEEDUP   = 10.0   # CPU → GPU speedup for ML workloads
_CPU_SLOWDOWN  = 10.0   # GPU → CPU slowdown
_GPU_TDP_W     = 80.0   # typical mid-range laptop GPU TDP (W)
_CPU_TDP_W     = 35.0   # typical laptop CPU TDP under load (W)
EMISSION_FACTOR = 0.475 # kg CO₂ per kWh (India grid average)


def _build_comparison(m: Dict[str, Any]) -> Dict[str, Any]:
    """
    Returns estimated performance metrics for BOTH hardware options
    so the UI can render a CPU vs GPU comparison panel.

    Keys:
        current_hw, alt_hw
        current_runtime, alt_runtime
        current_energy_wh, alt_energy_wh
        current_co2_g, alt_co2_g
        current_power_w, alt_power_w
        current_util, alt_util
        speedup          – ratio alt_runtime / current_runtime (>1 = alt is slower)
        co2_saved_pct    – % CO₂ saved by switching to recommended
    """
    hw_type  = str(m.get("hardware_type", "CPU"))
    runtime  = float(m.get("runtime_sec",  1) or 1)
    energy_wh= float(m.get("energy_kwh",   0) or 0) * 1000
    co2_g    = float(m.get("co2_g",        0) or 0)
    cpu_util = float(m.get("cpu_util_avg", 0) or 0)
    gpu_util = float(m.get("gpu_util_avg", 0) or 0)

    if hw_type == "CPU":  # running on CPU → estimate GPU alternative
        alt_runtime   = runtime / _GPU_SPEEDUP
        alt_power_w   = _GPU_TDP_W
        alt_energy_wh = alt_power_w * (alt_runtime / 3600)
        alt_co2_g     = alt_energy_wh * EMISSION_FACTOR
        alt_util      = min(70.0, cpu_util * 0.6)   # GPU won't be at same % as CPU was
        return dict(
            current_hw        = "CPU",
            alt_hw            = "GPU",
            current_runtime   = round(runtime,   3),
            alt_runtime       = round(alt_runtime, 3),
            current_energy_wh = round(energy_wh, 6),
            alt_energy_wh     = round(alt_energy_wh, 6),
            current_co2_g     = round(co2_g,     6),
            alt_co2_g         = round(alt_co2_g, 6),
            current_power_w   = round(_CPU_TDP_W, 1),
            alt_power_w       = round(_GPU_TDP_W, 1),
            current_util      = round(cpu_util,   1),
            alt_util          = round(alt_util,   1),
            speedup           = round(_GPU_SPEEDUP, 1),
            co2_saved_pct     = round((1 - alt_co2_g / max(co2_g, 1e-9)) * 100, 1),
        )
    else:  # running on GPU → show what CPU would have cost
        alt_runtime   = runtime * _CPU_SLOWDOWN
        alt_power_w   = _CPU_TDP_W
        alt_energy_wh = alt_power_w * (alt_runtime / 3600)
        alt_co2_g     = alt_energy_wh * EMISSION_FACTOR
        alt_util      = min(95.0, gpu_util * 1.4)
        return dict(
            current_hw        = "GPU",
            alt_hw            = "CPU",
            current_runtime   = round(runtime,   3),
            alt_runtime       = round(alt_runtime, 3),
            current_energy_wh = round(energy_wh, 6),
            alt_energy_wh     = round(alt_energy_wh, 6),
            current_co2_g     = round(co2_g,     6),
            alt_co2_g         = round(alt_co2_g, 6),
            current_power_w   = round(_GPU_TDP_W, 1),
            alt_power_w       = round(_CPU_TDP_W, 1),
            current_util      = round(gpu_util,   1),
            alt_util          = round(alt_util,   1),
            speedup           = round(_CPU_SLOWDOWN, 1),
            co2_saved_pct     = round((1 - co2_g / max(alt_co2_g, 1e-9)) * 100, 1),
        )
