reflection_engine.py

"""
SELF-REFLECTION ENGINE
======================
The brain of the self-improving code AI.
Uses Ollama locally to write, reflect, and improve Python code
until it hits a 99.5% benchmark floor.

Engineering philosophy:

  NEVER GIVE UP:
    If a problem is hard, it's a challenge. Not a wall — a ladder.
    When iterations fail, don't accept it. Re-reason with what you learned.
    Escalate strategy. Change the angle of attack. The engine persists
    because the builder persists.

  RADIAL SLOP:
    There is no perfection in engineering — that's why we target 99.5%, not 100%.
    A bolt threads into a hole because of radial slop: the tiny imperfection
    that makes the interface work. Code is the same — some tolerance is by design.
    Chase diminishing returns and you'll over-engineer or hallucinate problems.

  WORK > CODE > TALK:
    Three modes of output — know the difference:
      WORK  — does the code actually run, pass tests, and solve the task?
      CODE  — is the code itself clean, correct, and structured?
      TALK  — chatty prose, philosophical rambling, markdown explanations.
    The engine cares about WORK and CODE. TALK is waste.

  THE LADDER:
    Everything has a logical pattern. When things fail, backtrack from the
    closest point to the break, then work back one rung at a time.
    Follow the steps and nothing can hide. Don't guess — trace.

Architecture:
  1. Generate code
  2. Run it / test it
  3. If it BROKE → troubleshoot (ladder: start at the error, walk backwards)
  4. If it ran → score it
  5. If score < 99.5 → reflect on WHY → improve → repeat
  6. If stuck → re-reason with accumulated failure context → new angle of attack
  7. If score >= 99.5 → ship it (good enough IS good enough)
  8. If max iterations → report what was tried and what the next attack vector would be
"""

import subprocess
import tempfile
import textwrap
import json
import time
import sys
import re
import os
from dataclasses import dataclass, field
from typing import Optional
import requests
from code_whitelist import validate_code

# ── TALK detection ────────────────────────────────────────────────────────────

# Patterns that signal conversational / non-code input
_TALK_PATTERNS = [
    # Greetings
    r"^(hey|hi|hello|yo|sup|howdy|hiya|greetings)\b",
    # Status checks
    r"^are you\b",
    r"^you (are|there|up|awake|alive|with us|listening|around)\b",
    r"^how are you",
    r"^how('?s| is) it going",
    r"^who are you",
    # General knowledge questions (no code signals required)
    r"^what is \w+\??$",
    r"^what are \w+\??$",
    r"^what('?s| is) (up|good|new|your name|happening)",
    r"^what('?s| is) a \w+\??$",
    r"^what('?s| is) an \w+\??$",
    r"^what does \w+ mean\??$",
    r"^(describe|explain|define) \w+\??$",
    r"^tell me about\b",
    r"^do you know (about|what)?\b",
    # Acknowledgments / reactions
    r"^thanks?\b",
    r"^thank you",
    r"^good (morning|afternoon|evening|night|job|one)\b",
    r"^(yes|no|ok|okay|sure|yep|nah|nope|cool|nice|great|awesome|perfect|sounds good)$",
    r"^what can you do",
    r"^(help|help me)$",
    r"^(bye|goodbye|see ya|later|peace|cheers)\b",
    r"^(lol|lmao|haha|ha)\b",
    r"^(please|pls)\b$",
    # System checks
    r"^is (this|that|it) (working|on|alive|running)",
    r"^anybody (there|home|here)",
    r"^can you hear me",
    r"^still (there|alive|awake|with us|running)",
    r"^you still (there|alive|awake|with us)",
    r"^are we good",
    r"^what(')?s crackin",
    # Capability / conversational questions
    r"^(can|could|would|will|do) you\b",
    r"^(is it|are there)\b",
    r"^(should|shall) (i|we)\b",
    r"^do you (like|think|have|want|need|know)\b",
]
_TALK_RE = re.compile("|".join(_TALK_PATTERNS), re.IGNORECASE)

# Keywords that strongly signal a CODE task even if phrased conversationally
_CODE_SIGNALS = [
    "function", "class", "import", "def ", "return", "loop",
    "sort", "parse", "api", "endpoint", "database", "file",
    "algorithm", "test", "assert", "calculate", "convert",
    "build", "create", "implement", "write", "generate",
    "fix", "debug", "refactor", "optimize", "deploy",
]


def classify_task(task: str) -> str:
    """
    Classify input as 'WORK' (code task) or 'TALK' (conversation).
    If it looks like conversation and has no code keywords → TALK.
    """
    stripped = task.strip().rstrip("?!.")
    lower = stripped.lower()
    word_count = len(stripped.split())
    has_code = any(kw in lower for kw in _CODE_SIGNALS)

    # TALK pattern match with no code keywords → TALK (any length)
    if _TALK_RE.search(stripped) and not has_code:
        return "TALK"

    # Very short input (≤ 6 words) with no code signals → likely TALK
    if word_count <= 6 and not has_code:
        return "TALK"

    return "WORK"


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

OLLAMA_URL      = "http://localhost:11434/api/generate"
LM_STUDIO_URL   = "http://localhost:1234/v1/chat/completions"
BACKEND         = os.environ.get("CODE_ENGINE_BACKEND", "auto")  # "ollama", "lmstudio", "native", or "auto"
MODEL           = os.environ.get("CODE_ENGINE_MODEL", "llama3.2")
LM_STUDIO_MODEL = os.environ.get("LM_STUDIO_MODEL", "")
BENCHMARK_FLOOR = float(os.environ.get("BENCHMARK_FLOOR", "99.5"))
MAX_ITERATIONS  = int(os.environ.get("MAX_ITERATIONS", "10"))
TIMEOUT_SECS    = 30                # per code-run timeout

# ── Native backend config (no LM Studio / Ollama needed) ─────────────────────
# Qwen — the coder
NATIVE_HF_REPO     = os.environ.get("NATIVE_HF_REPO", "lmstudio-community/Qwen3-Coder-30B-A3B-Instruct-GGUF")
NATIVE_HF_FILE     = os.environ.get("NATIVE_HF_FILE", "Qwen3-Coder-30B-A3B-Instruct-Q4_K_M.gguf")
NATIVE_MODEL_PATH  = os.environ.get("NATIVE_MODEL_PATH", r"D:\va_data\models\Qwen3-Coder-30B-A3B-Instruct-Q4_K_M.gguf")
# DeepSeek R1 14B — the reasoner (Opus Reasoning 27B when hardware ready)
REASONER_HF_REPO   = os.environ.get("REASONER_HF_REPO", "bartowski/DeepSeek-R1-Distill-Qwen-14B-GGUF")
REASONER_HF_FILE   = os.environ.get("REASONER_HF_FILE", "DeepSeek-R1-Distill-Qwen-14B-Q4_K_M.gguf")
REASONER_MODEL_PATH = os.environ.get("REASONER_MODEL_PATH", r"D:\va_data\models\DeepSeek-R1-Distill-Qwen-14B-Q4_K_M.gguf")
# GPU layers to offload (0 = CPU only, -1 = all layers to GPU)
NATIVE_GPU_LAYERS  = int(os.environ.get("NATIVE_GPU_LAYERS", "0"))
NATIVE_CTX_SIZE    = int(os.environ.get("NATIVE_CTX_SIZE", "8192"))
_native_llm = None    # lazy singleton — Qwen (coder)
_reasoning_llm = None # lazy singleton — DeepSeek R1 (reasoner)

# ── Task mode routing ─────────────────────────────────────────────────────────
# "code" → Qwen for generation, "reasoning" → DeepSeek R1 for thinking/OSINT
_TASK_MODE = "code"  # default

def set_task_mode(mode: str):
    """Set which model handles native calls. 'code' = Qwen, 'reasoning' = DeepSeek R1."""
    global _TASK_MODE
    mode = mode.lower().strip()
    if mode in ("reasoning", "osint", "bellingcat", "investigation"):
        _TASK_MODE = "reasoning"
    else:
        _TASK_MODE = "code"
    print(f"  🧠 Task mode: {_TASK_MODE} ({'DeepSeek R1' if _TASK_MODE == 'reasoning' else 'Qwen Coder'})")

def get_task_mode() -> str:
    return _TASK_MODE

# ── Data structures ───────────────────────────────────────────────────────────

@dataclass
class CodeAttempt:
    iteration:    int
    code:         str
    score:        float
    passed:       bool
    test_output:  str
    errors:       list[str]
    reflections:  list[str] = field(default_factory=list)

@dataclass
class ReflectionResult:
    what_failed:       str
    root_cause:        str
    improvement_plan:  str
    specific_fix:      str
    confidence:        float        # 0-1: how confident the model is in the fix

@dataclass
class TroubleshootResult:
    """The ladder — start at the break, walk backwards rung by rung."""
    error_line:        str          # the exact line / traceback that broke
    ladder:            list[str]    # rungs walked back: each step traced
    root_cause:        str          # the actual bottom rung — why it broke
    fix:               str          # minimal code change to un-break it
    confidence:        float        # 0-1

@dataclass
class ReasoningPlan:
    """Think before you code. Decompose → plan → then generate."""
    sub_problems:      list[str]    # task broken into parts
    approach:          str          # algorithm / data structure / architecture choice
    edge_cases:        list[str]    # things that could bite you
    steps:             list[str]    # ordered implementation steps
    confidence:        float        # 0-1: how clear is the path

# ── Ollama interface ──────────────────────────────────────────────────────────

def _resolve_model_path(model_path, hf_repo, hf_file):
    """Resolve a GGUF path: use local override, check LM Studio cache, or download."""
    if model_path and os.path.isfile(model_path):
        return model_path
    # Check LM Studio cache
    lms_path = os.path.join(
        os.path.expanduser("~"), ".lmstudio", "models",
        hf_repo.replace("/", os.sep), hf_file
    )
    if os.path.isfile(lms_path):
        print(f"  ✓ Found local GGUF: {os.path.basename(lms_path)}")
        return lms_path
    # Download from HuggingFace
    print(f"  ⬇ Downloading {hf_file} from {hf_repo}...")
    from huggingface_hub import hf_hub_download
    path = hf_hub_download(repo_id=hf_repo, filename=hf_file)
    print(f"  ✓ Downloaded to: {path}")
    return path


def _load_native_model():
    """Load Qwen (coder) natively via llama-cpp-python."""
    global _native_llm
    if _native_llm is not None:
        return _native_llm

    from llama_cpp import Llama
    model_path = _resolve_model_path(NATIVE_MODEL_PATH, NATIVE_HF_REPO, NATIVE_HF_FILE)

    print(f"  ⏳ Loading Qwen coder (gpu_layers={NATIVE_GPU_LAYERS}, ctx={NATIVE_CTX_SIZE})...")
    _native_llm = Llama(
        model_path=model_path,
        n_ctx=NATIVE_CTX_SIZE,
        n_gpu_layers=NATIVE_GPU_LAYERS,
        verbose=False
    )
    print(f"  ✓ Qwen coder loaded: {os.path.basename(model_path)}")
    return _native_llm


def _load_reasoning_model():
    """Load DeepSeek R1 (reasoner) natively via llama-cpp-python."""
    global _reasoning_llm
    if _reasoning_llm is not None:
        return _reasoning_llm

    from llama_cpp import Llama
    model_path = _resolve_model_path(REASONER_MODEL_PATH, REASONER_HF_REPO, REASONER_HF_FILE)

    print(f"  ⏳ Loading DeepSeek R1 (gpu_layers={NATIVE_GPU_LAYERS}, ctx={NATIVE_CTX_SIZE})...")
    _reasoning_llm = Llama(
        model_path=model_path,
        n_ctx=NATIVE_CTX_SIZE,
        n_gpu_layers=NATIVE_GPU_LAYERS,
        verbose=False
    )
    print(f"  ✓ DeepSeek R1 loaded: {os.path.basename(model_path)}")
    return _reasoning_llm


def _call_native(prompt: str, system: str = "", temperature: float = 0.3) -> str:
    """Call the appropriate native model based on task mode.
    reasoning/osint → DeepSeek R1, code → Qwen Coder."""
    if _TASK_MODE == "reasoning":
        llm = _load_reasoning_model()
        model_label = "DeepSeek R1"
    else:
        llm = _load_native_model()
        model_label = "Qwen Coder"
    messages = []
    if system:
        messages.append({"role": "system", "content": system})
    messages.append({"role": "user", "content": prompt})

    # Trim prompt if it exceeds context budget
    total_chars = sum(len(m["content"]) for m in messages)
    max_chars = NATIVE_CTX_SIZE * 3  # rough char-to-token ratio
    if total_chars > max_chars:
        messages[-1]["content"] = messages[-1]["content"][:max_chars - 500] + "\n[trimmed to fit context]"

    result = llm.create_chat_completion(
        messages=messages,
        temperature=temperature,
        max_tokens=4096
    )
    text = result["choices"][0]["message"]["content"].strip()

    # DeepSeek R1 wraps output in <think>...</think> — strip it
    if "<think>" in text:
        text = re.sub(r'<think>.*?</think>', '', text, flags=re.DOTALL).strip()

    return text


def _detect_backend() -> str:
    """Auto-detect which backend is running. Tries LM Studio → Ollama → native."""
    if BACKEND != "auto":
        return BACKEND
    # Try LM Studio first (fastest — already running as a server)
    try:
        r = requests.get("http://localhost:1234/v1/models", timeout=3)
        if r.status_code == 200:
            models = r.json().get("data", [])
            if models:
                global LM_STUDIO_MODEL
                if not LM_STUDIO_MODEL:
                    LM_STUDIO_MODEL = models[0].get("id", "")
                print(f"  ✓ LM Studio detected (model: {LM_STUDIO_MODEL})")
                return "lmstudio"
    except Exception:
        pass
    # Try Ollama
    try:
        r = requests.get("http://localhost:11434/api/tags", timeout=3)
        if r.status_code == 200:
            print(f"  ✓ Ollama detected (model: {MODEL})")
            return "ollama"
    except Exception:
        pass
    # Fall back to native — Harvester is a complete unit
    # Check that at least one model file exists (don't eagerly load into RAM)
    try:
        if _TASK_MODE == "reasoning":
            model_path = _resolve_model_path(REASONER_MODEL_PATH, REASONER_HF_REPO, REASONER_HF_FILE)
        else:
            model_path = _resolve_model_path(NATIVE_MODEL_PATH, HF_REPO, HF_FILE)
        if os.path.isfile(model_path):
            return "native"
        raise FileNotFoundError(f"Model not found: {model_path}")
    except Exception as e:
        print(f"\n[ERROR] Native backend failed: {e}")
    print("\n[ERROR] No backend found. Options:")
    print("  1. Start LM Studio and load a model")
    print("  2. Run `ollama serve` then `ollama pull llama3.2`")
    print("  3. Set NATIVE_MODEL_PATH to a .gguf file (or let it download from HuggingFace)")
    sys.exit(1)


def _call_lmstudio(prompt: str, system: str = "", temperature: float = 0.3) -> str:
    """Call LM Studio OpenAI-compatible endpoint."""
    messages = []
    if system:
        messages.append({"role": "system", "content": system})
    messages.append({"role": "user", "content": prompt})
    payload = {
        "model": LM_STUDIO_MODEL or "local-model",
        "messages": messages,
        "temperature": temperature,
        "max_tokens": 4096,
        "stream": False
    }
    for attempt in range(3):
        try:
            r = requests.post(LM_STUDIO_URL, json=payload, timeout=180)
        except requests.exceptions.ReadTimeout:
            print(f"  [LM Studio TIMEOUT] attempt {attempt+1}/3 — trimming prompt")
            messages[-1]["content"] = messages[-1]["content"][:2000] + "\n\n[trimmed — keep response short]"
            if "max_tokens" in payload:
                payload["max_tokens"] = 2048
            continue
        if r.status_code == 200:
            return r.json()["choices"][0]["message"]["content"].strip()
        # Log the error
        err_body = r.text[:500] if r.text else "(empty)"
        print(f"  [LM Studio {r.status_code}] attempt {attempt+1}/3: {err_body}")
        if r.status_code >= 500 or r.status_code == 429:
            time.sleep(2 ** attempt)
            continue
        # 400 — context overflow is the #1 cause with small models
        is_context_overflow = "context size" in err_body.lower() or "context length" in err_body.lower()
        if attempt == 0:
            if is_context_overflow:
                # Aggressive trim — cut to 1500 chars, drop system msg, halve max_tokens
                messages[-1]["content"] = messages[-1]["content"][:1500] + "\n[trimmed for context]"
                messages = [m for m in messages if m["role"] != "system"]
                payload["messages"] = messages
                payload["max_tokens"] = 2048
            else:
                messages[-1]["content"] = messages[-1]["content"][:2500] + "\n[trimmed]"
                if "max_tokens" in payload:
                    payload["max_tokens"] = 2048
            continue
        if attempt == 1:
            # Nuclear trim — bare minimum prompt
            messages[-1]["content"] = messages[-1]["content"][:1000] + "\n[trimmed hard]"
            messages = [m for m in messages if m["role"] != "system"]
            payload["messages"] = messages
            payload.pop("max_tokens", None)
            continue
        # attempt 2 — strip non-ASCII + minimal prompt as last resort
        content = messages[-1]["content"][:800]
        clean = content.encode('ascii', 'ignore').decode('ascii')
        messages[-1]["content"] = clean
        payload["messages"] = messages
        try:
            r2 = requests.post(LM_STUDIO_URL, json=payload, timeout=180)
        except requests.exceptions.ReadTimeout:
            break
        if r2.status_code == 200:
            return r2.json()["choices"][0]["message"]["content"].strip()
    raise RuntimeError("LM Studio: all attempts failed (400/timeout)")


def _call_ollama(prompt: str, system: str = "", temperature: float = 0.3) -> str:
    """Call local Ollama model."""
    payload = {
        "model": MODEL,
        "prompt": prompt,
        "system": system,
        "stream": False,
        "options": {"temperature": temperature}
    }
    r = requests.post(OLLAMA_URL, json=payload, timeout=120)
    r.raise_for_status()
    return r.json().get("response", "").strip()


_ACTIVE_BACKEND = None

def ollama(prompt: str, system: str = "", temperature: float = 0.3) -> str:
    """Call local LLM (LM Studio, Ollama, or native). Returns response text."""
    global _ACTIVE_BACKEND
    if _ACTIVE_BACKEND is None:
        _ACTIVE_BACKEND = _detect_backend()
    try:
        if _ACTIVE_BACKEND == "native":
            return _call_native(prompt, system, temperature)
        elif _ACTIVE_BACKEND == "lmstudio":
            return _call_lmstudio(prompt, system, temperature)
        else:
            return _call_ollama(prompt, system, temperature)
    except requests.exceptions.ConnectionError:
        # If a server backend drops, try falling back to native
        if _ACTIVE_BACKEND in ("lmstudio", "ollama"):
            print(f"\n  ⚠ Lost connection to {_ACTIVE_BACKEND} — falling back to native...")
            try:
                _ACTIVE_BACKEND = "native"
                return _call_native(prompt, system, temperature)
            except Exception:
                pass
        print(f"\n[ERROR] Lost connection to {_ACTIVE_BACKEND}")
        print("  → Make sure your LLM backend is still running")
        sys.exit(1)
    except (requests.exceptions.ReadTimeout, requests.exceptions.Timeout):
        print(f"\n[ERROR] LLM timed out — model may be overloaded")
        print("  → Try a smaller model or restart LM Studio")
        raise RuntimeError("LLM timeout — model too slow for this prompt")
    except requests.exceptions.HTTPError as e:
        print(f"\n[ERROR] LLM returned HTTP error: {e}")
        raise RuntimeError(f"LLM backend error: {e}") from e

# ── Language config ────────────────────────────────────────────────────────────

# Supported languages and their generation config
_LANG_CONFIG = {
    "python": {
        "expert": "Python engineer",
        "extension": ".py",
        "fence": "python",
        "system": """You write only clean Python code.
No markdown. No backticks. No prose. No conversation. No talk.
Just raw executable Python starting with imports or function definitions.
Always include test assertions at the bottom of the file.
Remember: WORK (does it run?) and CODE (is it clean?) are all that matter.
TALK (explanations, commentary) is forbidden.""",
        "executable": True,
    },
    "csharp": {
        "expert": "Unity C# engineer (Unity 6, URP, Netcode for GameObjects)",
        "extension": ".cs",
        "fence": "csharp",
        "system": """You write only clean Unity C# code.
No markdown. No backticks. No prose. No conversation. No talk.
Just raw C# starting with using directives.
Use Unity 6 APIs and Netcode for GameObjects where networking is needed.
Remember: WORK (does it compile and function?) and CODE (is it clean?) matter.
TALK (explanations, commentary) is forbidden.""",
        "executable": False,  # no C# compiler in the loop — review-scored
    },
}

# Active language for this run (set via set_language())
_ACTIVE_LANG = "python"

def set_language(lang: str):
    """Set the active language for code generation."""
    global _ACTIVE_LANG
    lang = lang.lower().strip()
    if lang in ("cs", "c#", "unity"):
        lang = "csharp"
    if lang not in _LANG_CONFIG:
        print(f"  [WARN] Unknown language '{lang}' — defaulting to python")
        lang = "python"
    _ACTIVE_LANG = lang
    print(f"  Language: {_ACTIVE_LANG} ({_LANG_CONFIG[lang]['expert']})")

def get_language() -> str:
    return _ACTIVE_LANG

# ── Code generation ───────────────────────────────────────────────────────────

def generate_code(task: str, previous_attempt: Optional[CodeAttempt] = None,
                  reflection: Optional[ReflectionResult] = None,
                  reasoning: Optional[ReasoningPlan] = None) -> str:
    """Generate code for a task in the active language, with optional improvement context."""
    lang = _LANG_CONFIG[_ACTIVE_LANG]
    expert = lang["expert"]
    fence = lang["fence"]

    if previous_attempt and reflection:
        # Improvement pass — give it specific direction
        clean_task = task.split("\n\n[PRIOR KNOWLEDGE FROM MEMORY]")[0] if "[PRIOR KNOWLEDGE" in task else task
        prev_code = previous_attempt.code
        if len(prev_code) > 800:
            prev_code = prev_code[:300] + "\n// ...[middle trimmed]...\n" + prev_code[-500:]
        prompt = f"""Expert {expert}. Fix this code. Output ONLY working {fence}. No talk. No markdown.

TASK: {clean_task}

PREVIOUS CODE (scored {previous_attempt.score:.1f}%):
{prev_code}

WHAT FAILED: {reflection.what_failed}
ROOT CAUSE: {reflection.root_cause}
FIX: {reflection.specific_fix}

Write the complete fixed code. No explanation."""

    else:
        # First attempt — use reasoning plan if available
        reasoning_block = ""
        if reasoning and reasoning.confidence > 0.2:
            steps_str = "\n".join(f"  {j}. {s}" for j, s in enumerate(reasoning.steps, 1))
            edges_str = ", ".join(reasoning.edge_cases[:4]) if reasoning.edge_cases else "none identified"
            reasoning_block = f"""

YOUR REASONING PLAN (follow this):
  Approach: {reasoning.approach}
  Steps:
{steps_str}
  Edge cases to handle: {edges_str}

Implement the plan above. Do not deviate unless the plan is clearly wrong."""

        if _ACTIVE_LANG == "python":
            prompt = f"""You are an expert {expert}. Target: 99.5%% quality — not perfection.
Radial slop: a bolt threads because it's imperfect. Some tolerance is by design.

There are three kinds of output: WORK, CODE, and TALK.
- WORK = does it run, pass tests, solve the task? This is what matters most.
- CODE = is it clean, readable, well-structured? This matters second.
- TALK = prose, explanations, markdown, philosophy. This is WASTE. Zero talk.

TASK: {task}{reasoning_block}

Write clean, efficient Python code to accomplish this task.
- Handle common edge cases (don't invent unlikely scenarios)
- Include sensible error handling — no defensive paranoia
- Write at least 3 meaningful test cases as assert statements at the bottom
- Correctness first, then readability, then performance

Return ONLY the Python code. No explanation. No markdown fences. No talk."""
        else:
            prompt = f"""You are an expert {expert}. Target: 99.5%% quality — not perfection.

There are three kinds of output: WORK, CODE, and TALK.
- WORK = does it compile and function correctly? This is what matters most.
- CODE = is it clean, readable, well-structured? This matters second.
- TALK = prose, explanations, markdown, philosophy. This is WASTE. Zero talk.

TASK: {task}{reasoning_block}

Write clean, production-quality {fence} code.
- Handle common edge cases (don't invent unlikely scenarios)
- Use proper namespace, class structure, and access modifiers
- Correctness first, then readability, then performance

Return ONLY the {fence} code. No explanation. No markdown fences. No talk."""

    system = lang["system"]
    return ollama(prompt, system=system, temperature=0.2)

# ── Code execution & scoring ──────────────────────────────────────────────────

def run_code(code: str) -> tuple[bool, str, list[str]]:
    """
    Execute code in a subprocess sandbox (Python) or review-score it (other languages).
    Returns: (success, output, errors)
    """
    lang = _ACTIVE_LANG
    lang_cfg = _LANG_CONFIG[lang]

    # Safety check — validate against multi-language whitelist
    validation = validate_code(code, language=lang)
    if not validation.safe:
        return False, "", [f"BLOCKED: {b}" for b in validation.blocked]

    # Non-executable languages → LLM review scoring
    if not lang_cfg["executable"]:
        return _review_score_code(code, lang)

    errors = []
    output = ""

    with tempfile.NamedTemporaryFile(mode='w', suffix='.py',
                                     delete=False, encoding='utf-8') as f:
        f.write(code)
        tmp_path = f.name

    try:
        result = subprocess.run(
            [sys.executable, tmp_path],
            capture_output=True, text=True,
            timeout=TIMEOUT_SECS
        )
        output = result.stdout.strip()
        if result.returncode != 0:
            errors.append(result.stderr.strip())
            success = False
        else:
            success = True
    except subprocess.TimeoutExpired:
        errors.append(f"TIMEOUT: code exceeded {TIMEOUT_SECS}s")
        success = False
    except Exception as e:
        errors.append(str(e))
        success = False
    finally:
        os.unlink(tmp_path)

    return success, output, errors


def _review_score_code(code: str, lang: str) -> tuple[bool, str, list[str]]:
    """
    For non-executable languages (C#, etc), use LLM code review as scoring.
    Returns: (success, review_output, errors_found)
    """
    print("  📝 Review-scoring (no compiler available)...")
    # Trim code for context budget
    code_snippet = code[:2000] if len(code) > 2000 else code
    prompt = f"""You are a strict code reviewer for {lang}. Score this code 0-100.

CODE:
{code_snippet}

Score on:
1. Does it look like it would compile? (40 points)
2. Is the logic correct for the stated purpose? (30 points)
3. Is it clean and well-structured? (20 points)
4. Edge cases handled? (10 points)

Respond ONLY with JSON:
{{"score": 85, "errors": ["list of issues found"], "notes": "brief summary"}}"""
    raw = ollama(prompt, temperature=0.1)
    try:
        match = re.search(r'\{.*\}', raw, re.DOTALL)
        if match:
            data = json.loads(match.group())
            score = float(data.get("score", 50))
            errors = data.get("errors", [])
            notes = data.get("notes", "")
            success = score >= BENCHMARK_FLOOR
            return success, f"Review score: {score:.0f}% — {notes}", errors if not success else []
    except Exception:
        pass
    return False, "Review parsing failed", ["Could not parse review response"]


def score_code(code: str, task: str, success: bool,
               output: str, errors: list[str]) -> float:
    """
    WORK-first scoring. Did it run? Did tests pass? That's what matters.
    For non-executable languages, uses the review score from _review_score_code.
    """
    lang_cfg = _LANG_CONFIG[_ACTIVE_LANG]

    # Non-executable languages: extract score from review output
    if not lang_cfg["executable"]:
        match = re.search(r'Review score:\s*(\d+)', output)
        if match:
            return float(match.group(1))
        return 50.0 if success else 30.0

    # Python: execution-based scoring
    if not success:
        if errors and "AssertionError" in str(errors):
            return 35.0   # ran but assertions failed — close
        if errors and "AssertError" in str(errors):
            return 35.0
        return 20.0       # didn't run at all

    if not errors:
        return 100.0

    return 85.0

# ── Self-reflection ───────────────────────────────────────────────────────────

def reflect(task: str, attempt: CodeAttempt) -> ReflectionResult:
    """
    The core of the engine.
    Forces the model to diagnose failure before attempting a fix.
    This is what makes it self-improving rather than just retrying.
    """
    print(f"\n  🧠 Reflecting on iteration {attempt.iteration} (score: {attempt.score:.1f}%)...")
    expert = _LANG_CONFIG[_ACTIVE_LANG]["expert"]
    fence = _LANG_CONFIG[_ACTIVE_LANG]["fence"]

    reflection_prompt = f"""You are a senior {expert} conducting a code review.

ENGINEERING PRINCIPLES:
- Never give up. If it's hard, it's a challenge. Change the angle of attack.
- Radial slop: target 99.5%, not 100%. The tiny tolerance is a feature.

Your job is to evaluate TWO things and ONLY two things:
1. WORK — does the code actually run, pass its tests, and solve the stated task?
2. CODE — is the code clean, readable, and reasonably efficient?

Do NOT produce TALK — no philosophical commentary, no essays about best practices,
no restating the task. Diagnose real issues in the fewest words possible.
If the code works and is clean, say so and move on.

TASK: {task}

CODE (scored {attempt.score:.1f}% — target is 99.5%):
```{fence}
{attempt.code}
```

EXECUTION OUTPUT: {attempt.test_output}
ERRORS: {attempt.errors}

Diagnose only REAL issues. Answer with precision:

1. WHAT_FAILED: What specifically is broken or missing? (If nothing real, say "minor gap")
2. ROOT_CAUSE: Why? What assumption was wrong?
3. IMPROVEMENT_PLAN: Practical steps to reach 99.5% — not 100%. Don't over-engineer.
4. SPECIFIC_FIX: The single most impactful code change. Keep it minimal.
5. CONFIDENCE: How confident are you this fix will work? (0.0 to 1.0)

Respond ONLY with a JSON object (no other text):
{{
  "what_failed": "...",
  "root_cause": "...",
  "improvement_plan": "...",
  "specific_fix": "...",
  "confidence": 0.85
}}"""

    raw = ollama(reflection_prompt, temperature=0.2)

    try:
        raw = re.sub(r'```.*?```', '', raw, flags=re.DOTALL).strip()
        # Find JSON block if surrounded by text
        match = re.search(r'\{.*\}', raw, re.DOTALL)
        if match:
            raw = match.group()
        data = json.loads(raw)
        return ReflectionResult(
            what_failed      = data.get("what_failed", "Unknown failure"),
            root_cause       = data.get("root_cause", "Unknown cause"),
            improvement_plan = data.get("improvement_plan", "Retry"),
            specific_fix     = data.get("specific_fix", "Rewrite"),
            confidence       = float(data.get("confidence", 0.5))
        )
    except Exception:
        return ReflectionResult(
            what_failed      = "Parsing failed — raw output: " + raw[:200],
            root_cause       = "Model did not return valid JSON",
            improvement_plan = "Rewrite code from scratch with stricter constraints",
            specific_fix     = "Full rewrite",
            confidence       = 0.4
        )

# ── Troubleshooter (the ladder) ───────────────────────────────────────────────

def troubleshoot(task: str, attempt: CodeAttempt) -> TroubleshootResult:
    """
    The ladder method: start at the break, walk backwards rung by rung.
    Everything has a logical pattern. Follow the steps, nothing hides.
    
    This is NOT reflect(). Reflect is for code that RUNS but scored low.
    Troubleshoot is for code that BROKE — crashed, errored, wouldn't execute.
    """
    print(f"\n  🪜 Troubleshooting iteration {attempt.iteration} (ladder method)...")
    fence = _LANG_CONFIG[_ACTIVE_LANG]["fence"]

    # Build the error context — start at the closest point to the break
    error_text = "\n".join(attempt.errors) if attempt.errors else "Unknown error"
    
    ts_prompt = f"""You are a troubleshooter. Use the LADDER METHOD:

NEVER GIVE UP. If the error looks hard, it's a challenge — not a wall.

THE LADDER: Start at the exact point of failure. Read the error message.
Then walk backwards through the code one step at a time — like climbing
down a ladder rung by rung. Every failure has a logical pattern.
Follow the steps and nothing can hide.

Do NOT guess. Do NOT skip rungs. Trace the actual execution path.

Rules:
- Output is WORK only. Zero TALK.
- Start at the error (top of ladder)
- Walk back through each line that led to it
- Find the root cause (bottom rung)
- Give one minimal fix

TASK: {task}

CODE:
```{fence}
{attempt.code}
```

ERROR (start here — this is the top of the ladder):
{error_text}

OUTPUT (if any): {attempt.test_output}

Walk the ladder. Respond ONLY with JSON — zero talk:
{{
  "error_line": "the exact error or traceback line",
  "ladder": [
    "RUNG 1 (error): what broke",
    "RUNG 2 (one step back): what called it or fed it bad data",
    "RUNG 3 (deeper): the actual source of the problem"
  ],
  "root_cause": "the bottom rung — the real reason",
  "fix": "minimal code change to un-break it",
  "confidence": 0.85
}}"""

    raw = ollama(ts_prompt, temperature=0.2)

    try:
        raw = re.sub(r'```.*?```', '', raw, flags=re.DOTALL).strip()
        match = re.search(r'\{.*\}', raw, re.DOTALL)
        if match:
            raw = match.group()
        data = json.loads(raw)
        result = TroubleshootResult(
            error_line = data.get("error_line", error_text[:200]),
            ladder     = data.get("ladder", ["Could not trace"]),
            root_cause = data.get("root_cause", "Unknown"),
            fix        = data.get("fix", "Rewrite"),
            confidence = float(data.get("confidence", 0.5))
        )
    except Exception:
        result = TroubleshootResult(
            error_line = error_text[:200],
            ladder     = ["Parsing failed — could not walk ladder"],
            root_cause = "Model did not return valid JSON",
            fix        = "Full rewrite",
            confidence = 0.3
        )

    # Print the ladder
    print(f"  🪜 Error: {result.error_line[:100]}")
    for j, rung in enumerate(result.ladder, 1):
        print(f"     ↓ Rung {j}: {rung[:100]}")
    print(f"  ⚡ Root cause: {result.root_cause[:100]}")
    print(f"  🔧 Fix: {result.fix[:100]}")

    return result


def _troubleshoot_to_reflection(ts: TroubleshootResult) -> ReflectionResult:
    """Convert a TroubleshootResult into a ReflectionResult so the
    improvement pass can consume it uniformly."""
    return ReflectionResult(
        what_failed      = ts.error_line,
        root_cause       = ts.root_cause,
        improvement_plan = " → ".join(ts.ladder),
        specific_fix     = ts.fix,
        confidence       = ts.confidence
    )


# ── Reasoning helpers ─────────────────────────────────────────────────────────

def _parse_reasoning(raw: str, clean_task: str) -> ReasoningPlan:
    """Try to parse structured JSON from LLM reasoning output."""
    try:
        # Strip <think>...</think> tags (DeepSeek R1)
        cleaned = re.sub(r'<think>.*?</think>', '', raw, flags=re.DOTALL).strip()
        # Strip code fences
        cleaned = re.sub(r'```.*?```', '', cleaned, flags=re.DOTALL).strip()
        # Find JSON object
        match = re.search(r'\{.*\}', cleaned, re.DOTALL)
        if match:
            cleaned = match.group()
        data = json.loads(cleaned)
        return ReasoningPlan(
            sub_problems = data.get("sub_problems", [clean_task]),
            approach     = data.get("approach", "Direct implementation"),
            edge_cases   = data.get("edge_cases", []),
            steps        = data.get("steps", ["Implement the task"]),
            confidence   = float(data.get("confidence", 0.5))
        )
    except Exception:
        return ReasoningPlan(
            sub_problems = [clean_task],
            approach     = "Direct implementation (structured parse failed)",
            edge_cases   = [],
            steps        = ["Implement the task directly"],
            confidence   = 0.3
        )


def _extract_reasoning_from_prose(raw: str, clean_task: str) -> ReasoningPlan:
    """
    When JSON parse fails, don't give up — read the prose.
    Extract whatever structure we can from natural language reasoning.
    This is blind spot identification: the model DID think, it just didn't format.
    """
    # Strip think tags but KEEP the content — that's where the reasoning is
    text = re.sub(r'</?think>', '', raw).strip()

    # Extract sub-problems: look for numbered lists or bullet points
    sub_problems = []
    for m in re.finditer(r'(?:^|\n)\s*(?:\d+[\.\)]\s*|[-•]\s*)(.+)', text):
        item = m.group(1).strip()
        if len(item) > 10 and len(item) < 200:
            sub_problems.append(item)
    if not sub_problems:
        sub_problems = [clean_task]

    # Extract approach: look for strategy keywords
    approach = "Direct implementation"
    approach_patterns = [
        r'(?:approach|strategy|plan|method|algorithm)[:\s]+(.+?)(?:\n|$)',
        r'(?:I would|We should|The idea is|The key is)[:\s]*(.+?)(?:\n|$)',
        r'(?:use|using|leverage|implement with)\s+(.+?)(?:\n|$)',
    ]
    for pat in approach_patterns:
        m = re.search(pat, text, re.IGNORECASE)
        if m:
            approach = m.group(1).strip()[:200]
            break

    # Extract edge cases: look for warning/risk language
    edge_cases = []
    edge_patterns = [
        r'(?:edge case|corner case|risk|caveat|watch out|careful|tricky)[:\s]*(.+?)(?:\n|$)',
        r'(?:what if|might fail|could break|problem if)[:\s]*(.+?)(?:\n|$)',
    ]
    for pat in edge_patterns:
        for m in re.finditer(pat, text, re.IGNORECASE):
            edge_cases.append(m.group(1).strip()[:100])

    # Confidence: based on how much structure we extracted
    has_approach = approach != "Direct implementation"
    has_subs = len(sub_problems) > 1
    has_edges = len(edge_cases) > 0
    signals = sum([has_approach, has_subs, has_edges])

    if signals >= 2:
        confidence = 0.65  # good reasoning, just bad formatting
    elif signals == 1:
        confidence = 0.5   # partial understanding
    else:
        confidence = 0.35  # genuine uncertainty — honest low score

    return ReasoningPlan(
        sub_problems = sub_problems[:10],
        approach     = approach,
        edge_cases   = edge_cases[:5],
        steps        = sub_problems[:10],  # use sub-problems as steps
        confidence   = confidence
    )


# ── Reasoning engine (think before you code) ──────────────────────────────────

def reason(task: str) -> ReasoningPlan:
    """
    Chain-of-thought reasoning BEFORE code generation.
    Decomposes the task, picks an approach, spots edge cases, and builds a plan.
    This is the difference between 'generate and hope' and 'think then build'.
    """
    print(f"\n  🧩 Reasoning about the task...")
    expert = _LANG_CONFIG[_ACTIVE_LANG]["expert"]
    fence = _LANG_CONFIG[_ACTIVE_LANG]["fence"]

    # Strip memory context for cleaner reasoning
    clean_task = task.split("\n\n[PRIOR KNOWLEDGE")[0].strip() if "[PRIOR KNOWLEDGE" in task else task

    reason_prompt = f"""You are an expert {expert}. THINK about this task before writing any code.

PRINCIPLE: If a problem is hard, it's a challenge — not a blocker.
Never give up on a hard problem. Decompose it. Find the angle of attack.
The harder it looks, the more value there is in solving it.

TASK: {clean_task}

Break it down. Plan your approach. Identify what could go wrong.
If the task seems complex, break it into smaller winnable pieces.
Do NOT write code. Just THINK.

Respond ONLY with JSON — zero talk:
{{
  "sub_problems": ["problem 1", "problem 2", "problem 3"],
  "approach": "the algorithm, data structure, or architecture you'd use and WHY",
  "edge_cases": ["edge case 1", "edge case 2"],
  "steps": ["step 1: ...", "step 2: ...", "step 3: ..."],
  "confidence": 0.85
}}"""

    raw = ollama(reason_prompt, temperature=0.3)

    plan = _parse_reasoning(raw, clean_task)

    # Self-awareness: if parse failed, don't give up — ask differently
    if plan.confidence <= 0.3 and plan.approach.startswith("Direct implementation"):
        print("  🔄 First reasoning attempt unclear — rephrasing and retrying...")
        retry_prompt = f"""You are an expert {expert}. A previous attempt to reason about this task
produced unclear output. Let's try again with a simpler structure.

TASK: {clean_task}

Think step by step. What are the 2-3 main sub-problems?
What approach would you use? What could go wrong?

Respond with ONLY this JSON — nothing else, no thinking tags, no prose:
{{"sub_problems": ["first piece", "second piece"], "approach": "your strategy", "edge_cases": ["risk 1"], "steps": ["step 1", "step 2"], "confidence": 0.7}}"""
        raw2 = ollama(retry_prompt, temperature=0.2)
        plan2 = _parse_reasoning(raw2, clean_task)
        if plan2.confidence > plan.confidence:
            plan = plan2
            print("  ✓ Retry produced better reasoning")
        else:
            # Still can't parse structured JSON — extract what we can from the raw text
            plan = _extract_reasoning_from_prose(raw, clean_task)
            if plan.confidence > 0.3:
                print("  ✓ Extracted reasoning from unstructured response")

    # Print the plan
    print(f"  🧩 Approach: {plan.approach[:120]}")
    print(f"  🧩 Sub-problems: {len(plan.sub_problems)}")
    for j, step in enumerate(plan.steps, 1):
        print(f"     {j}. {step[:100]}")
    if plan.edge_cases:
        print(f"  ⚠️  Edge cases: {', '.join(e[:60] for e in plan.edge_cases[:4])}")
    print(f"  🎲 Confidence: {plan.confidence:.0%}")

    return plan


# ── Main engine loop ──────────────────────────────────────────────────────────

def run_engine(task: str, talk_memory: str = "") -> CodeAttempt:
    """
    Main self-improvement loop.
    Runs until benchmark floor is hit or max iterations reached.
    """
    # ── TALK gate: don't spin up the code loop for conversation ──
    # Strip any enrichment context so classify sees the raw task
    raw_task = task.split("\n\n[PRIOR KNOWLEDGE")[0].strip() if "[PRIOR KNOWLEDGE" in task else task
    task_type = classify_task(raw_task)
    if task_type == "TALK":
        print(f"\n  💬 TALK detected — responding conversationally, no code loop.")
        memory_block = ""
        if talk_memory:
            memory_block = f"\n\n[YOUR MEMORY — reference this when relevant]\n{talk_memory}\n"
        reply = ollama(
            f"The user said: {task}{memory_block}\nRespond naturally and helpfully. Reference your memory if the user asks what you know, what you've learned, or about past work. Keep it concise — a few sentences max.",
            system="You are Harvester, a self-improving AI code engine with persistent memory. You have a database of coding patterns, failure lessons, and task history that you can recall. When asked what you know or remember, reference your memory. Never say you don't remember — you DO have memory. Answer knowledge questions clearly. For greetings, respond warmly. Never output code, function definitions, or test cases. Just plain conversational text.",
            temperature=0.5
        )
        return CodeAttempt(
            iteration=0,
            code=reply.strip(),
            score=100.0,
            passed=True,
            test_output="",
            errors=[],
            reflections=["TALK — conversational response with memory recall"]
        )

    print(f"\n{'='*60}")
    print(f"  SELF-IMPROVING CODE ENGINE")
    print(f"  Model:     {LM_STUDIO_MODEL or MODEL}")
    print(f"  Target:    {BENCHMARK_FLOOR}%")
    print(f"  Max loops: {MAX_ITERATIONS}")
    print(f"{'='*60}")
    print(f"\n  TASK: {task}\n")

    # ── Reason before coding ──
    plan = reason(task)

    attempts: list[CodeAttempt] = []
    current_code = None
    last_reflection = None

    for i in range(1, MAX_ITERATIONS + 1):
        print(f"\n{'─'*60}")
        print(f"  ITERATION {i} / {MAX_ITERATIONS}")
        print(f"{'─'*60}")

        # ── Generate ──
        print("  ✍️  Generating code...")
        last_attempt = attempts[-1] if attempts else None
        code = generate_code(task, last_attempt, last_reflection,
                             reasoning=plan if i == 1 else None)

        # Clean up model output just in case
        code = re.sub(r'^```python\n?', '', code.strip())
        code = re.sub(r'^```\n?', '', code.strip())
        code = re.sub(r'```$', '', code.strip())

        # ── Execute ──
        print("  ⚙️  Running code...")
        success, output, errors = run_code(code)

        # ── Score ──
        print("  📊 Scoring...")
        score = score_code(code, task, success, output, errors)

        attempt = CodeAttempt(
            iteration   = i,
            code        = code,
            score       = score,
            passed      = score >= BENCHMARK_FLOOR,
            test_output = output,
            errors      = errors
        )
        attempts.append(attempt)

        status = "✅ PASS" if attempt.passed else "❌ FAIL"
        print(f"\n  {status} — Score: {score:.1f}%")
        if errors:
            print(f"  Errors: {errors[0][:120]}")
        if output:
            print(f"  Output: {output[:120]}")

        if attempt.passed:
            print(f"\n{'='*60}")
            print(f"  🎯 BENCHMARK HIT on iteration {i}! Score: {score:.1f}%")
            print(f"{'='*60}")
            return attempt

        # ── Troubleshoot or Reflect ──
        if i < MAX_ITERATIONS:
            if not success and errors:
                # Code BROKE — use the ladder: backtrack from the error
                ts_result = troubleshoot(task, attempt)
                last_reflection = _troubleshoot_to_reflection(ts_result)
                attempt.reflections.append(f"[LADDER] {last_reflection.improvement_plan}")
            else:
                # Code ran but scored low — reflect on quality
                last_reflection = reflect(task, attempt)
                attempt.reflections.append(last_reflection.improvement_plan)
            print(f"  💡 Root cause: {(last_reflection.root_cause or 'unknown')[:100]}")
            print(f"  🔧 Fix: {(last_reflection.specific_fix or 'none')[:100]}")
            print(f"  🎲 Confidence: {last_reflection.confidence:.0%}")

            # ── NEVER GIVE UP: Re-reason when stuck ──
            # If confidence is low or we're past the halfway point with no
            # improvement, re-reason with failure context for a new angle.
            scores = [a.score for a in attempts]
            stuck = (len(scores) >= 3 and max(scores[-3:]) <= max(scores) * 1.05)
            low_confidence = last_reflection.confidence < 0.4

            if stuck or low_confidence:
                print(f"\n  🔄 NEVER GIVE UP — re-reasoning with failure context...")
                failure_context = f"""Previous attempts failed. Here's what we know:
- Best score so far: {max(scores):.1f}%
- Last error: {last_reflection.root_cause}
- Fix tried: {last_reflection.specific_fix}
- Attempts: {len(attempts)}

Original task: {task}

CHANGE YOUR ANGLE OF ATTACK. The previous approach isn't working."""
                plan = reason(failure_context)
                # Feed new plan into next iteration
                if plan.confidence > 0.3:
                    print(f"  🔄 New approach: {plan.approach[:100]}")
                    # Override: next generate_code gets the fresh plan
                    # by injecting reasoning into the reflection
                    last_reflection = ReflectionResult(
                        what_failed=last_reflection.what_failed,
                        root_cause=last_reflection.root_cause,
                        improvement_plan=f"NEW APPROACH: {plan.approach}. Steps: {'; '.join(plan.steps[:4])}",
                        specific_fix=f"Follow new reasoning plan: {plan.steps[0] if plan.steps else 'rethink'}",
                        confidence=plan.confidence
                    )

    # Hit max iterations — NEVER GIVE UP: report what was tried and next attack
    best = max(attempts, key=lambda a: a.score)
    all_reflections = []
    for a in attempts:
        all_reflections.extend(a.reflections)
    print(f"\n{'='*60}")
    print(f"  ⚠️  Max iterations reached. Best score: {best.score:.1f}%")
    print(f"  Approaches tried: {len(attempts)}")
    if all_reflections:
        print(f"  Last strategy: {all_reflections[-1][:120]}")
    if last_reflection:
        print(f"  Next attack would be: {last_reflection.specific_fix[:120]}")
    print(f"  This is a challenge, not a wall. Run again to continue.")
    print(f"{'='*60}")
    return best


# ── Output formatter ──────────────────────────────────────────────────────────

def print_final(attempt: CodeAttempt):
    print(f"\n{'='*60}")
    print(f"  FINAL OUTPUT  (iteration {attempt.iteration}, score {attempt.score:.1f}%)")
    print(f"{'='*60}\n")
    print(attempt.code)
    print(f"\n{'─'*60}")
    if attempt.reflections:
        print(f"  Improvement steps taken: {len(attempt.reflections)}")
        for i, r in enumerate(attempt.reflections, 1):
            print(f"  {i}. {r[:80]}")


# ── Entry point ───────────────────────────────────────────────────────────────

if __name__ == "__main__":
    if len(sys.argv) > 1:
        task = " ".join(sys.argv[1:])
    else:
        print("\nSELF-REFLECTION ENGINE — Interactive Mode")
        print("─" * 40)
        task = input("What should I build? → ").strip()
        if not task:
            task = "Write a function that finds all prime numbers up to N using the Sieve of Eratosthenes, with tests"

    result = run_engine(task)
    print_final(result)