import os
from pathlib import Path


def _load_key_from_env(key_name: str = "GEMINI_API_KEY") -> str | None:
    """
    Try to find the API key from multiple sources (works locally + Colab):
      1. Directly passed in (handled upstream)
      2. os.environ / .env file
      3. Google Colab secrets (google.colab.userdata)
    """
    # 1. Standard env var / .env
    try:
        from dotenv import load_dotenv
        # walk up looking for .env
        for parent in [Path.cwd(), Path.cwd().parent, Path(__file__).resolve().parents[1]]:
            env_file = parent / ".env"
            if env_file.exists():
                load_dotenv(env_file)
                break
    except Exception:
        pass

    key = os.getenv(key_name)
    if key:
        return key

    # 2. Google Colab secrets
    try:
        from google.colab import userdata
        key = userdata.get(key_name)
        if key:
            return key
    except Exception:
        pass

    return None


# ── try new SDK first, fall back to old ──────────────────────────────
try:
    from google import genai as _new_genai
    _SDK = "new"
except ImportError:
    _new_genai = None
    _SDK = None

if _SDK is None:
    try:
        import google.generativeai as _old_genai
        _SDK = "old"
    except ImportError:
        _old_genai = None
        _SDK = None


class GeminiAdvisor:

    def __init__(self, api_key: str = None):
        """
        api_key : pass the key directly (ideal for Colab / cloud notebooks).
                  If None, auto-detected from .env / env vars / Colab secrets.
        """
        key = api_key or _load_key_from_env()

        if not key:
            print("⚠ Gemini disabled – no API key found (pass api_key= or set GEMINI_API_KEY)")
            self.enabled = False
            return

        if _SDK is None:
            print("⚠ Gemini client missing – run:  pip install google-genai")
            self.enabled = False
            return

        try:
            if _SDK == "new":
                self._client = _new_genai.Client(api_key=key)
                self._model  = "gemini-1.5-flash"
                self._mode   = "new"
            else:
                _old_genai.configure(api_key=key)
                self._client = _old_genai.GenerativeModel("gemini-1.5-flash")
                self._mode   = "old"

            self.enabled = True
            print("🤖 Gemini Enabled ✓")
        except Exception as e:
            print(f"⚠ Gemini init failed: {e}")
            self.enabled = False

    # ------------------------------------------------------------------
    
    def ask(self, prompt: str) -> str | None:
        """Generic single-prompt call (SDK-agnostic)."""
        try:
            if self._mode == "new":
                r = self._client.models.generate_content(
                    model=self._model, contents=prompt
                )
            else:
                r = self._client.generate_content(prompt)
            return r.text.strip()
        except Exception as e:
            return f'{{"tip": "Gemini API Error: {str(e)[:150]}"}}'

    # ------------------------------------------------------------------
    def ask_rich(self, metrics: dict, code: str = None, error=None) -> dict | None:
        """Rich context-aware prompt with all metrics. Returns a parsed JSON dict."""
        cpu      = metrics.get("cpu_util_avg",    0)
        gpu      = metrics.get("gpu_util_avg",    0)
        cpu_temp = metrics.get("cpu_temp_c",      0)
        gpu_temp = metrics.get("gpu_temp_c",      0)
        ram_pct  = metrics.get("ram_used_pct",    0)
        vram_pct = metrics.get("gpu_mem_used_pct",0)
        energy   = float(metrics.get("energy_kwh", 0)) * 1000   # mWh
        co2      = metrics.get("co2_g",           0)
        runtime  = metrics.get("runtime_sec",     0)
        hw_type  = metrics.get("hardware_type",  "CPU")

        error_line = f"\nRuntime error: {str(error)[:200]}" if error else ""
        code_line = f"\nExecuted Code:\n{code[:800]}\n" if code else ""

        prompt = f"""You are an expert ML efficiency advisor for the EcoML IDE widget.

System: {hw_type} | CPU: {cpu:.1f}% | GPU: {gpu:.1f}% | RAM: {ram_pct:.1f}% | Runtime: {runtime:.2f}s | Energy: {energy:.2f}mWh {error_line}{code_line}

Provide an actionable insight to improve efficiency.
CRITICAL INSTRUCTION: You must be EXTREMELY concise. MAXIMUM 1 short sentence per field. Do not write paragraphs.

If you can suggest a concrete code change (e.g., changing a hyperparameter, offloading to GPU, or fixing an error), you MUST return a structured JSON object with EXACTLY these keys:
{{
  "detection": "Briefly state what is wrong (e.g., CPU bottleneck detected (95%))",
  "reason": "Why it causes inefficiency (e.g., CPU is overloaded causing overhead)",
  "action": "What to do instead (e.g., Switch to GPU or reduce estimators)",
  "impact": "Estimated impact (e.g., ~40% lower runtime, ~30% less CO₂)",
  "before": "exact old line of code",
  "after": "exact new line of code"
}}
If you cannot suggest a specific code fix, omit 'before' and 'after':
{{
  "detection": "...",
  "reason": "...",
  "action": "...",
  "impact": "..."
}}

Respond with RAW JSON ONLY. No markdown formatted blocks. No prefix.
"""
        import json
        import re
        
        res = self.ask(prompt)
        if not res: return None
        
        # ROBUST JSON EXTRACTION: Find the JSON dictionary even if the AI adds text
        match = re.search(r'(\{.*\})', res, re.DOTALL)
        if match:
            try:
                return json.loads(match.group(1))
            except json.JSONDecodeError:
                pass
                
        # Fallback if it completely failed to generate valid JSON
        return {"tip": res.strip()[:150]}

    # ------------------------------------------------------------------
    def generate_fix(self, err: Exception, code: str | None) -> str | None:
        """Generate 1-line fix for a runtime error."""
        code_snippet = (code or "")[:600]
        prompt = f"""You are a senior Python/ML debugger.

Error: {err}

Code: {code_snippet}

Reply with ONE concise fix in plain text (no markdown). Max 120 chars.
"""
        return self.ask(prompt)
