"""
phase4/explainer.py
====================
Phase 4 — Natural Language Explainability Layer.

Calls Groq (llama-3.3-70b-versatile) to explain why each phase scored
the code the way it did. Three parallel calls — one per phase — each
following a structured prompt:
  1. Understand the code (what it does, language, complexity)
  2. Analyse through the phase's specific detection lens
  3. Justify the actual score returned

Key allocation:
  Phase 2 uses GROQ_API_KEY_1  to GROQ_API_KEY_20  (heavy — 4 rewrites)
  Phase 4 uses GROQ_API_KEY_21 to GROQ_API_KEY_33  (lighter — 3 calls)
  The ranges never overlap so the two phases never compete.

Key rotation:
  Cooldown on per-minute 429s (TPM), burn on daily quota 429s (TPD).
  If all keys fail, explain() returns the original result unchanged —
  detect() always works without Phase 4.

Public API:
  explain(result: dict, code: str) -> dict
    Takes the dict returned by orchestrator.detect(), adds an
    "explanations" key, and returns the enriched dict.
    Always safe — if keys are missing or all calls fail the
    original result dict is returned unchanged.

  build_explanation_block(result: dict) -> str
    Returns a formatted text block for appending to build_report().
    Returns "" if no explanations are present.
"""

from __future__ import annotations

import os
import re
import time
import threading
import concurrent.futures
from typing import Optional

import httpx
from dotenv import load_dotenv

load_dotenv()

# ── Config ────────────────────────────────────────────────────────────────────

GROQ_MODEL    = "llama-3.3-70b-versatile"
GROQ_ENDPOINT = "https://api.groq.com/openai/v1/chat/completions"

# Keys reserved for Phase 4 — do not overlap with Phase 2 (1-20)
KEY_RANGE_START = 21
KEY_RANGE_END   = 33

MAX_OUTPUT_TOKENS = 800
TEMPERATURE       = 0.3
REQUEST_TIMEOUT   = 40


# ── Phase metadata ────────────────────────────────────────────────────────────

_PHASE_META = {
    "p1": {
        "name": "Phase 1 — Stylometric Analysis (Random Forest)",
        "method": (
            "This phase extracts ~82–136 hand-crafted numerical features from "
            "the source code using Tree-sitter (or Python's native AST for Python). "
            "Feature groups include:\n"
            "  • Layout/lexical: line lengths, indentation patterns, character-class "
            "densities, comment ratios, blank-line ratios.\n"
            "  • Identifier style: snake_case vs camelCase vs PascalCase ratios.\n"
            "  • Generic AST: tree depth, node count, branching factor, leaf ratio.\n"
            "  • Semantic AST: densities of function definitions, conditionals, loops, "
            "assignments, try/catch blocks, imports, type annotations, literals.\n"
            "  • Python-only extras: cyclomatic complexity, maintainability index, "
            "keyword densities, docstring presence.\n"
            "These features are fed to a per-language Random Forest "
            "(300 estimators, balanced class weights). "
            "Output is a probability in [0, 1] where 1.0 = certainly AI."
        ),
        "score_key":   "p1",
        "threshold":   0.5,
        "score_label": "P1 (stylometry score)",
    },
    "p2": {
        "name": "Phase 2 — Rewrite-Similarity Analysis (SimCSE + LLM)",
        "method": (
            "This phase implements the method from Ye et al. (AAAI 2025, "
            "'Uncovering LLM-Generated Code: A Zero-Shot Synthetic Code Detector "
            "via Code Rewriting'). The pipeline is:\n"
            "  1. Strip all comments, docstrings, and blank lines from the code.\n"
            "  2. Ask an LLM (Llama-3.3-70b via Groq) to produce 4 independent "
            "rewrites of the same logic using different variable names, control flow, "
            "and idioms.\n"
            "  3. Embed the original and all rewrites using SimCSE-GraphCodeBERT "
            "(CLS-pooled, L2-normalised, 768-d vectors).\n"
            "  4. Compute mean cosine similarity between the original and each rewrite.\n"
            "  5. Map mean similarity → P(AI) via a calibrated logistic regression.\n"
            "Core insight: AI-generated code has a consistent latent style that "
            "LLM rewrites preserve (high cosine similarity). Human code is more "
            "idiosyncratic — rewrites diverge more (lower similarity)."
        ),
        "score_key":   "p2",
        "threshold":   0.5,
        "score_label": "P2 (rewrite-similarity score)",
    },
    "p3": {
        "name": "Phase 3 — Neural Classifier (CodeT5p + Binary Head)",
        "method": (
            "This phase uses the model from Gurioli et al. (SANER 2025, "
            "'Is This You, LLM?'). Architecture:\n"
            "  • Encoder: Salesforce/codet5p-770m (T5 encoder only, 770M parameters).\n"
            "  • Head: Linear(1024→768) → ReLU → Dropout(0.2) → Linear(768→1) → Sigmoid.\n"
            "  • Trained on H-AIRosettaMP: ~121k balanced samples across 10 languages "
            "(human: Rosetta Code; AI: StarCoder2).\n"
            "Outputs cluster near 1.0 due to sigmoid saturation so a calibrated "
            "threshold of 0.77 is used (not 0.5). Code is chunked into 512-token "
            "windows if needed; chunk scores are averaged. "
            "Final p_ai = 1 − sigmoid_output. p_ai > 0.77 → AI verdict."
        ),
        "score_key":   "p3",
        "threshold":   0.77,
        "score_label": "P3 (neural classifier score, calibrated threshold 0.77)",
    },
}


# ── Key rotation state ────────────────────────────────────────────────────────

_key_lock                        = threading.Lock()
_all_keys: list[str]             = []
_burned_keys: set[int]           = set()
_key_cooldowns: dict[int, float] = {}


def _load_all_keys() -> None:
    global _all_keys
    load_dotenv(override=True)
    keys = []
    for i in range(KEY_RANGE_START, KEY_RANGE_END + 1):
        k = os.getenv(f"GROQ_API_KEY_{i}")
        if k and k.strip():
            keys.append(k.strip())
    _all_keys = keys
    print(f"[Phase4] Loaded {len(_all_keys)} Groq key(s) "
          f"(slots {KEY_RANGE_START}–{KEY_RANGE_END})")


def _keys_available() -> bool:
    global _all_keys
    _load_all_keys()
    return bool(_all_keys)


def _get_best_key_index() -> Optional[int]:
    now = time.time()
    for i in range(len(_all_keys)):
        if i in _burned_keys:
            continue
        if _key_cooldowns.get(i, 0) <= now:
            return i
    best, earliest = None, float("inf")
    for i in range(len(_all_keys)):
        if i in _burned_keys:
            continue
        t = _key_cooldowns.get(i, 0)
        if t < earliest:
            earliest, best = t, i
    return best


def _burn_key(idx: int) -> None:
    with _key_lock:
        _burned_keys.add(idx)
        print(f"[Phase4] Key #{idx + KEY_RANGE_START} BURNED (daily limit). "
              f"({len(_burned_keys)}/{len(_all_keys)} burned)")


def _cooldown_key(idx: int, seconds: float) -> None:
    with _key_lock:
        _key_cooldowns[idx] = time.time() + seconds
        print(f"[Phase4] Key #{idx + KEY_RANGE_START} on cooldown for {seconds:.1f}s")


def _parse_retry_after(msg: str) -> Optional[float]:
    m = re.search(r"try again in (\d+)m([\d.]+)s", msg)
    if m:
        secs = int(m.group(1)) * 60 + float(m.group(2))
        return secs if secs <= 180 else None
    m = re.search(r"try again in ([\d.]+)s", msg)
    if m:
        secs = float(m.group(1))
        return secs if secs <= 180 else None
    return None


def _is_daily_limit(msg: str) -> bool:
    low = msg.lower()
    return any(x in low for x in ("tokens per day", "tpd", "daily limit"))


# ── Prompt builder ────────────────────────────────────────────────────────────

def _build_prompt(phase_key: str, code: str, result: dict) -> str:
    meta      = _PHASE_META[phase_key]
    score     = result.get(meta["score_key"])
    language  = result.get("language", "unknown")
    threshold = (
        result.get("p3_threshold", meta["threshold"])
        if phase_key == "p3"
        else meta["threshold"]
    )

    score_pct   = f"{int(score * 100)}%" if score is not None else "N/A"
    verdict_str = (
        "AI-GENERATED"
        if (score is not None and score > threshold)
        else "HUMAN-WRITTEN"
    )

    extra_ctx = ""
    if phase_key == "p2":
        parts = []
        if result.get("mean_sim")    is not None: parts.append(f"Mean cosine similarity: {result['mean_sim']:.4f}")
        if result.get("n_rewrites")  is not None: parts.append(f"Successful rewrites: {result['n_rewrites']}/4")
        if result.get("suspected_model"):          parts.append(f"Suspected origin model: {result['suspected_model']}")
        if parts:
            extra_ctx = "\nAdditional Phase 2 metrics:\n" + "\n".join(f"  • {p}" for p in parts)

    if phase_key == "p3" and result.get("p3_confidence"):
        extra_ctx = f"\nPhase 3 model confidence level: {result['p3_confidence']}"

    return f"""You are an expert code analyst explaining why an AI-code detection system scored a piece of code the way it did.

════════════════════════════════════════
DETECTION PHASE
════════════════════════════════════════
{meta['name']}

════════════════════════════════════════
HOW THIS PHASE WORKS
════════════════════════════════════════
{meta['method']}

════════════════════════════════════════
SCORE RECORDED
════════════════════════════════════════
  • {meta['score_label']}: {score} ({score_pct})
  • Threshold for AI verdict: {threshold}
  • Phase verdict: {verdict_str}
  • Programming language: {language}{extra_ctx}

════════════════════════════════════════
CODE UNDER ANALYSIS
════════════════════════════════════════
```{language}
{code}
```

════════════════════════════════════════
YOUR TASK
════════════════════════════════════════
Write a detailed explanation in exactly three clearly labelled sections:

**[1. CODE UNDERSTANDING]**
Describe what this code does functionally. Cover: the algorithm or task it solves, the language features it uses, its structural complexity (size, control flow depth, abstraction level), and any notable patterns or idioms present.

**[2. PHASE ANALYSIS]**
Analyse the code strictly through the lens of {meta['name']}. Be specific — reference actual lines, identifiers, patterns, or structures from the code. Explain which concrete signals in this code are consistent with AI authorship and which are consistent with human authorship, according to what this phase measures. Do not be generic — tie every observation to something visible in the code.

**[3. SCORE JUSTIFICATION]**
Explain why the score of {score} ({score_pct}) makes sense given your analysis. Connect the specific signals you identified to the numeric outcome. If the phase verdict conflicts with other phases or your intuition, acknowledge it and explain why. Be direct and analytical.

Keep the total response under 550 words. Write in coherent paragraphs within each section — no bullet spam."""


# ── Groq call with key rotation ───────────────────────────────────────────────

def _call_groq(prompt: str, preferred_key_idx: int) -> Optional[str]:
    global _all_keys
    if not _all_keys:
        _load_all_keys()
    if not _all_keys:
        return None

    max_attempts = len(_all_keys) + 1
    key_idx      = preferred_key_idx % len(_all_keys)

    for attempt in range(max_attempts):
        # Rotate away from burned / cooling keys
        with _key_lock:
            now = time.time()
            if key_idx in _burned_keys or _key_cooldowns.get(key_idx, 0) > now:
                key_idx = _get_best_key_index()
                if key_idx is None:
                    print("[Phase4] All keys unavailable")
                    return None

        # Wait out cooldown if needed
        wait = 0.0
        with _key_lock:
            cooldown_until = _key_cooldowns.get(key_idx, 0)
            if cooldown_until > time.time():
                wait = cooldown_until - time.time()

        if wait > 0:
            print(f"[Phase4] Waiting {wait:.1f}s for key cooldown...")
            time.sleep(wait + 0.3)

        headers = {
            "Authorization": f"Bearer {_all_keys[key_idx]}",
            "Content-Type":  "application/json",
        }
        payload = {
            "model":       GROQ_MODEL,
            "messages":    [{"role": "user", "content": prompt}],
            "max_tokens":  MAX_OUTPUT_TOKENS,
            "temperature": TEMPERATURE,
        }

        try:
            with httpx.Client(timeout=REQUEST_TIMEOUT) as client:
                resp = client.post(GROQ_ENDPOINT, headers=headers, json=payload)

            if resp.status_code == 200:
                data = resp.json()
                text = (
                    data.get("choices", [{}])[0]
                        .get("message", {})
                        .get("content", "")
                )
                return text.strip() or None

            elif resp.status_code == 429:
                body = resp.text
                if _is_daily_limit(body):
                    _burn_key(key_idx)
                else:
                    wait_secs = _parse_retry_after(body) or 60.0
                    _cooldown_key(key_idx, wait_secs + 1.0)
                with _key_lock:
                    key_idx = _get_best_key_index()
                if key_idx is None:
                    return None

            elif resp.status_code in (401, 403):
                print(f"[Phase4] Key #{key_idx + KEY_RANGE_START} auth error — burning")
                _burn_key(key_idx)
                with _key_lock:
                    key_idx = _get_best_key_index()
                if key_idx is None:
                    return None

            else:
             body = resp.text
             if resp.status_code == 400 and "organization_restricted" in body:
                print(f"[Phase4] Key #{key_idx + KEY_RANGE_START} restricted — burning")
                _burn_key(key_idx)
                with _key_lock:
                    key_idx = _get_best_key_index()
                if key_idx is None:
                    return None
                continue 
             else:
              print(f"[Phase4] Unexpected status {resp.status_code}: {body[:300]}")
              return None

        except (httpx.TimeoutException, httpx.RequestError) as e:
            print(f"[Phase4] Request error: {e}")
            _cooldown_key(key_idx, 10.0)
            with _key_lock:
                key_idx = _get_best_key_index()
            if key_idx is None:
                return None

    return None


# ── Public API ────────────────────────────────────────────────────────────────

def explain(result: dict, code: str) -> dict:
    """
    Generate natural language explanations for each phase that produced a score.

    Takes the dict returned by orchestrator.detect(), runs up to 3 parallel
    Groq calls (one per phase that did NOT return None/abstain), and returns
    the enriched dict with an "explanations" key added.

    Always safe — if keys are missing or all calls fail, returns the
    original result dict unchanged.
    """
    if not _keys_available():
        print("[Phase4] No Groq keys available (slots 21–33) — skipping explanations")
        return result

    phases_to_explain = [
        key for key in ("p1", "p2", "p3")
        if result.get(_PHASE_META[key]["score_key"]) is not None
    ]

    if not phases_to_explain:
        return result

    prompts = {
        phase: _build_prompt(phase, code, result)
        for phase in phases_to_explain
    }

    n_keys = len(_all_keys)
    assigned_keys = {
        phase: (i % n_keys)
        for i, phase in enumerate(phases_to_explain)
    }

    explanations: dict[str, Optional[str]] = {"p1": None, "p2": None, "p3": None}

    def _explain_phase(phase: str) -> tuple[str, Optional[str]]:
        print(f"[Phase4] Generating explanation for {phase.upper()}...")
        text = _call_groq(prompts[phase], assigned_keys[phase])
        if text:
            print(f"[Phase4] {phase.upper()} done ({len(text)} chars)")
        else:
            print(f"[Phase4] {phase.upper()} failed")
        return phase, text

    with concurrent.futures.ThreadPoolExecutor(max_workers=3) as executor:
        futures = {executor.submit(_explain_phase, p): p for p in phases_to_explain}
        for fut in concurrent.futures.as_completed(futures):
            phase, text = fut.result()
            explanations[phase] = text

    if any(v is not None for v in explanations.values()):
        return {**result, "explanations": explanations}

    print("[Phase4] All explanation calls failed — result unchanged")
    return result


# ── Report formatter ──────────────────────────────────────────────────────────

def build_explanation_block(result: dict) -> str:
    """
    Format the explanations into a text block for appending to build_report().
    Returns "" if no explanations are present.
    """
    explanations = result.get("explanations")
    if not explanations:
        return ""

    phase_labels = {
        "p1": "Phase 1 — Stylometric Analysis",
        "p2": "Phase 2 — Rewrite-Similarity Analysis",
        "p3": "Phase 3 — Neural Classifier",
    }

    lines = []
    lines.append("")
    lines.append("=" * 60)
    lines.append("  PHASE-BY-PHASE EXPLANATION  (Phase 4)")
    lines.append("=" * 60)

    for key in ("p1", "p2", "p3"):
        text = explanations.get(key)
        if text is None:
            continue
        lines.append("")
        lines.append("-" * 60)
        lines.append(f"  {phase_labels[key]}")
        lines.append("-" * 60)
        for line in text.splitlines():
            lines.append(f"  {line}")

    lines.append("")
    lines.append("=" * 60)
    return "\n".join(lines)