utils.py

"""
RUBRA — utils.py
Helper functions: queue system, rate limit retry, zip creation, code parsing.
"""

import re
import io
import os
import json
import time
import asyncio
import logging
import zipfile
import hashlib
from typing import List, Dict, Optional, AsyncIterator
from collections import deque

try:
    from model_health import health_registry as _health_registry
    _HEALTH_OK = True
except ImportError:
    _HEALTH_OK = False
from dataclasses import dataclass, field
from datetime import datetime

log = logging.getLogger("rubra.utils")

# ═══════════════════════════════════════════════════════
#  REQUEST QUEUE — heavy requests don't block others
# ═══════════════════════════════════════════════════════

@dataclass
class QueuedRequest:
    id:         str
    prompt:     str
    framework:  str
    created_at: float = field(default_factory=time.time)
    status:     str   = "pending"   # pending | processing | done | error
    result:     dict  = field(default_factory=dict)

class RequestQueue:
    """Simple async queue for heavy generation requests."""

    def __init__(self, max_concurrent: int = 2):
        self.queue        = deque()
        self.processing   = {}
        self.results      = {}
        self.max_concurrent = max_concurrent
        self._lock        = asyncio.Lock()

    async def enqueue(self, prompt: str, framework: str) -> str:
        req_id = hashlib.md5(f"{prompt}{time.time()}".encode()).hexdigest()[:8]
        req    = QueuedRequest(id=req_id, prompt=prompt, framework=framework)
        async with self._lock:
            self.queue.append(req)
        log.info(f"[QUEUE] Enqueued {req_id} (depth={len(self.queue)})")
        return req_id

    async def get_status(self, req_id: str) -> dict:
        if req_id in self.results:
            return {"status": "done", "result": self.results[req_id]}
        if req_id in self.processing:
            return {"status": "processing"}
        for r in self.queue:
            if r.id == req_id:
                pos = list(self.queue).index(r)
                return {"status": "pending", "position": pos + 1}
        return {"status": "not_found"}

    async def next(self) -> Optional[QueuedRequest]:
        async with self._lock:
            if self.queue and len(self.processing) < self.max_concurrent:
                req = self.queue.popleft()
                self.processing[req.id] = req
                return req
        return None

    async def complete(self, req_id: str, result: dict):
        async with self._lock:
            self.processing.pop(req_id, None)
            self.results[req_id] = result
            # Keep last 100 results
            if len(self.results) > 100:
                oldest = list(self.results.keys())[0]
                del self.results[oldest]

# Global queue instance
request_queue = RequestQueue(max_concurrent=2)


# ═══════════════════════════════════════════════════════
#  RATE LIMIT HANDLER — smart retry with backoff
# ═══════════════════════════════════════════════════════

class RateLimitError(Exception):
    pass

async def with_retry(
    coro_factory,
    max_retries: int = 3,
    base_delay:  float = 1.0,
    max_delay:   float = 30.0,
    on_retry=None,
) -> any:
    """
    Execute async function with exponential backoff retry.
    Handles rate limits gracefully — user never sees raw errors.
    """
    last_err = None
    for attempt in range(1, max_retries + 1):
        try:
            return await coro_factory()
        except Exception as e:
            last_err = e
            err_str  = str(e).lower()

            # Rate limit / quota → wait and retry
            is_rate = any(t in err_str for t in [
                "429", "rate limit", "quota", "too many requests",
                "overloaded", "capacity", "retry after",
            ])

            if is_rate and attempt < max_retries:
                # Parse retry-after if available
                retry_after = 0
                match = re.search(r'retry.?after[:\s]+(\d+)', err_str)
                if match:
                    retry_after = int(match.group(1))

                delay = max(retry_after, min(base_delay * (2 ** (attempt - 1)), max_delay))
                log.warning(f"[RETRY] Rate limited, waiting {delay:.1f}s (attempt {attempt}/{max_retries})")

                if on_retry:
                    await on_retry(attempt, delay, str(e))

                await asyncio.sleep(delay)
                continue

            # Non-rate-limit error → don't retry
            if not is_rate:
                raise e

    raise last_err


# ═══════════════════════════════════════════════════════
#  ZIP CREATOR
# ═══════════════════════════════════════════════════════

def create_zip(files: List[Dict[str, str]]) -> bytes:
    """
    Create ZIP from list of {path, content} dicts.
    Returns raw bytes for download.
    """
    buf = io.BytesIO()
    with zipfile.ZipFile(buf, mode='w', compression=zipfile.ZIP_DEFLATED) as zf:
        for f in files:
            path    = f.get("path", "file.txt")
            content = f.get("content", "")
            if isinstance(content, str):
                content = content.encode("utf-8")
            zf.writestr(path, content)
    buf.seek(0)
    return buf.read()


# ═══════════════════════════════════════════════════════
#  CODE PARSER — extract code blocks from LLM response
# ═══════════════════════════════════════════════════════

def extract_code_blocks(text: str) -> List[Dict[str, str]]:
    """Extract all fenced code blocks from markdown text."""
    pattern = r'```(\w+)?\n([\s\S]*?)```'
    blocks  = []
    for match in re.finditer(pattern, text):
        lang    = match.group(1) or "text"
        content = match.group(2)
        blocks.append({"language": lang, "content": content})
    return blocks

def extract_file_blocks(text: str) -> List[Dict[str, str]]:
    """
    Extract file blocks from LLM output.
    Looks for patterns like:
      ### src/App.jsx
      ```jsx
      ...code...
      ```
    or:
      // FILE: src/App.jsx
      ```
      ...code...
      ```
    """
    files = []

    # Pattern 1: ### filename\n```lang\n...code...\n```
    pattern1 = r'###\s+([^\n]+)\n```\w*\n([\s\S]*?)```'
    for m in re.finditer(pattern1, text):
        path    = m.group(1).strip().strip('`').strip('"').strip("'")
        content = m.group(2)
        if path and '.' in path.split('/')[-1]:
            files.append({"path": path, "content": content})

    # Pattern 2: // FILE: path or # FILE: path
    pattern2 = r'(?://|#)\s*FILE:\s*([^\n]+)\n```\w*\n([\s\S]*?)```'
    for m in re.finditer(pattern2, text):
        path    = m.group(1).strip()
        content = m.group(2)
        files.append({"path": path, "content": content})

    # Pattern 3: **`path`** or **path**
    pattern3 = r'\*\*`?([^*`\n]+\.\w+)`?\*\*\s*\n```\w*\n([\s\S]*?)```'
    for m in re.finditer(pattern3, text):
        path    = m.group(1).strip()
        content = m.group(2)
        if '.' in path.split('/')[-1]:
            files.append({"path": path, "content": content})

    # Deduplicate by path
    seen  = set()
    dedup = []
    for f in files:
        if f["path"] not in seen:
            seen.add(f["path"])
            dedup.append(f)

    return dedup

def clean_llm_code(code: str) -> str:
    """Remove markdown fences and extra whitespace from code."""
    code = re.sub(r'^```\w*\n?', '', code.strip())
    code = re.sub(r'\n?```$',    '', code.strip())
    return code.strip()


# ═══════════════════════════════════════════════════════
#  STREAMING HELPERS
# ═══════════════════════════════════════════════════════

# Phase 13 hardening (2026-08-14): this used to be a second, completely
# separate SSE formatter from main.py's _sse() — main.py::_sse() had its
# own object-leak normalizer (added Phase 8, expanded Phase 13), but THIS
# function had none at all. /api/generate-project/stream and
# /api/generate-stream (both use sse_event(), not _sse()) were therefore
# NOT covered by any of that hardening — the exact "same class of gap on
# a different endpoint" this audit was looking for. Fixed by making this
# the ONE shared implementation — main.py::_sse() now delegates here
# instead of keeping a second copy of the same logic, so there is
# structurally only one place left to ever have this bug again.
SSE_TEXT_FIELDS = {"content", "message", "error", "preview", "text", "reason", "status_text",
                    "output", "result", "description", "status", "name", "agent",
                    "tool", "command", "intent"}


def _normalize_sse_value(v):
    """Coerce a single value to something safe to render as text.
    Strings/numbers/bools/None pass through untouched (None is valid JSON
    `null` and frontend code already has to handle "field absent"); a
    dict/list gets turned into a real, readable string instead of being
    forwarded raw for the frontend to accidentally str()-coerce into
    "[object Object]"."""
    if v is None or isinstance(v, (str, int, float, bool)):
        return v
    if isinstance(v, dict):
        for k in ("message", "error", "reason", "text", "output"):
            if isinstance(v.get(k), str):
                return v[k]
        try:
            return json.dumps(v, ensure_ascii=False)[:1000]
        except Exception:
            return str(v)[:1000]
    if isinstance(v, list):
        try:
            return json.dumps(v, ensure_ascii=False)[:1000]
        except Exception:
            return str(v)[:1000]
    return str(v)


def normalize_sse_payload(data, _depth: int = 0):
    """
    Recursive pass over an SSE event payload — walks any dict/list at any
    depth. A key in SSE_TEXT_FIELDS gets its value normalized (stringified
    if it's a dict/list) wherever it's found; a key NOT in that set whose
    value is itself a dict/list is recursed into (so a legitimately
    structured container like `files: [...]` is preserved as structured
    data, but anything text-shaped nested inside it — or inside any
    future event shape — still gets caught). Depth-capped at 6 as a cheap
    safety net against a pathological/circular payload.
    """
    if _depth > 6:
        return data
    if isinstance(data, dict):
        out = {}
        for key, val in data.items():
            if key in SSE_TEXT_FIELDS:
                out[key] = _normalize_sse_value(val)
            elif isinstance(val, (dict, list)):
                out[key] = normalize_sse_payload(val, _depth + 1)
            else:
                out[key] = val
        return out
    if isinstance(data, list):
        return [normalize_sse_payload(item, _depth + 1) for item in data]
    return data


def sse_event(data: dict) -> str:
    """Format dict as SSE event — normalized, see normalize_sse_payload()."""
    return f"data: {json.dumps(normalize_sse_payload(data), ensure_ascii=False)}\n\n"

def sse_done() -> str:
    return "data: [DONE]\n\n"

async def stream_text_chunks(text: str, chunk_size: int = 15, delay: float = 0.02) -> AsyncIterator[str]:
    """Stream text character by character for demo/fallback."""
    for i in range(0, len(text), chunk_size):
        yield text[i:i+chunk_size]
        await asyncio.sleep(delay)


# ═══════════════════════════════════════════════════════
#  PROMPT BUILDERS
# ═══════════════════════════════════════════════════════

def build_file_prompt(
    user_prompt:   str,
    file_path:     str,
    file_hint:     str,
    framework:     str,
    other_files:   List[str] = None,
) -> str:
    """Build prompt for generating a single file in a project."""
    other = f"\nOther files in project: {', '.join(other_files)}" if other_files else ""
    return f"""Generate the file `{file_path}` for a {framework} project.

Project description: {user_prompt}
File purpose: {file_hint}{other}

Rules:
- Write ONLY the file content, no explanation
- Complete, production-ready code
- No placeholders or TODOs
- Proper imports and exports

Output just the raw file content (no markdown fences):"""

def build_project_prompt(user_prompt: str, framework: str, file_list: List[str]) -> str:
    """Build prompt for generating all project files at once."""
    files_str = "\n".join(f"- {f}" for f in file_list)
    return f"""Generate a complete {framework} project based on this description:
{user_prompt}

Generate these files:
{files_str}

For each file, use this exact format:
### path/to/file.ext
```language
[file content here]
```

Rules:
- Write COMPLETE file contents — no truncation
- Production-ready code, no placeholders
- Follow {framework} best practices
- Include all imports"""


# ═══════════════════════════════════════════════════════
#  LANGUAGE DETECTION FOR EXPLANATION
# ═══════════════════════════════════════════════════════

def detect_code_language(code: str, hint: str = "") -> str:
    if hint: return hint.lower()
    if re.search(r'^\s*(def |import |from .+import|class .+:)', code, re.M): return "python"
    if re.search(r'(const |let |var |=>|async function|require\()', code): return "javascript"
    if re.search(r'(<[A-Z][a-zA-Z]+|jsx|tsx)', code): return "react"
    if re.search(r'(<!DOCTYPE|<html|<div)', code, re.I): return "html"
    if re.search(r'(\{.*:.*\}|@media|\.[\w-]+\s*\{)', code): return "css"
    if re.search(r'(SELECT|INSERT|CREATE TABLE|FROM|WHERE)', code, re.I): return "sql"
    if re.search(r'(#!.*bash|#!/bin/sh|^\s*\w+=)', code, re.M): return "bash"
    return "code"


# ═══════════════════════════════════════════════════════
#  COMPLETION VERIFICATION — "don't stop until it's actually done"
# ═══════════════════════════════════════════════════════
# llm()'s own cascade (in brain.py) only moves to the next model on a HARD
# failure — a 429, a timeout, a dropped connection. It has no idea a model
# that responded with HTTP 200 just gave back an empty string, a refusal,
# or three lines of a truncated <div>. That's a SOFT failure, and it was
# silently being counted as "done". This section closes that gap: every
# generated file gets checked against what it's actually supposed to look
# like, and a soft failure rotates to a genuinely different model — not
# just a second roll of the same dice — before giving up.

_REFUSAL_MARKERS = (
    "i cannot", "i can't", "i'm sorry", "as an ai", "i am unable",
    "i don't have the ability", "cannot provide", "i won't be able",
    "i'm not able to",
)

def validate_file_content(path: str, content: str) -> "tuple[bool, str]":
    """
    Dependency-free sanity check that a generated file is real and usable —
    not empty, not a refusal, not cut off mid-block, roughly shaped like its
    extension expects. NOT a full compiler — it's the floor that catches the
    actual failure mode reported in practice: empty/truncated/placeholder
    output getting reported to the user as "done" when it plainly wasn't.
    Returns (is_valid, reason_if_invalid).
    """
    if not content or not content.strip():
        return False, "empty content"

    stripped = content.strip()
    low = stripped.lower()

    if len(stripped) < 3:
        return False, "suspiciously short (<3 chars) — likely not real content"

    if any(marker in low[:300] for marker in _REFUSAL_MARKERS):
        return False, "looks like a refusal/apology, not actual file content"

    if stripped.startswith("```") or stripped.endswith("```"):
        return False, "still wrapped in markdown fences — not raw file content"

    ext = path.rsplit(".", 1)[-1].lower() if "." in path else ""

    if ext in ("html", "htm"):
        if "<" not in stripped or ">" not in stripped:
            return False, "doesn't look like HTML — no tags found"
        opens  = len(re.findall(r'<(html|head|body|div|main|section|nav)\b', low))
        closes = len(re.findall(r'</(html|head|body|div|main|section|nav)>', low))
        if opens >= 3 and closes < opens - 2:   # generous slack — this is not a real parser
            return False, f"structural tags look unbalanced ({opens} opened vs {closes} closed) — likely truncated"

    elif ext == "css":
        if stripped.count("{") == 0:
            return False, "no CSS rules found"
        if stripped.count("{") != stripped.count("}"):
            return False, "unbalanced braces — likely truncated"

    elif ext in ("js", "jsx", "ts", "tsx"):
        if stripped.count("{") != stripped.count("}"):
            return False, "unbalanced braces — likely truncated mid-block"
        if stripped.count("(") != stripped.count(")"):
            return False, "unbalanced parentheses — likely truncated mid-call"

    elif ext == "json":
        import json as _json
        try:
            _json.loads(stripped)
        except Exception as e:
            return False, f"invalid JSON: {str(e)[:80]}"

    elif ext == "py":
        try:
            compile(stripped, path or "<generated>", "exec")
        except SyntaxError as e:
            return False, f"Python syntax error: {str(e)[:80]}"

    return True, ""


async def generate_diff_edit_with_rotation(
    path:          str,
    current_content: str,
    instruction:   str,
    framework:     str,
    cascade:       list,
    stream_llm_func,
    max_rounds:    int = 2,
    max_tokens:    int = 8192,
    on_attempt=None,
) -> dict:
    """
    The 'edit, don't regenerate' engine. Asks the model for a SEARCH/REPLACE
    diff (or FULL_REWRITE for pervasive changes) instead of the whole file,
    applies it via merge_engine's 3-tier matcher, and validates the RESULT —
    same model-rotation safety net as generate_with_model_rotation(), plus
    one extra grounding step: a merge failure (ambiguous/not-found SEARCH
    block) re-prompts using merge_engine's auto-generated retry prompt,
    which carries the file's ACTUAL current content, instead of blindly
    repeating the same request that just failed.

    Returns: {"content","ok","reason","model","attempts","tiers"}
    """
    from merge_engine import merge_diff

    def _initial_prompt() -> str:
        return (
            f"Edit the file `{path}` ({framework} project).\n\n"
            f"=== CURRENT CONTENT ===\n{current_content}\n=== END CURRENT CONTENT ===\n\n"
            f"Requested change: \"{instruction}\"\n\n"
            f"Respond with ONE OR MORE SEARCH/REPLACE blocks in this exact format:\n"
            f"<<<<<<< SEARCH\n...exact text copied from CURRENT CONTENT above...\n"
            f"=======\n...replacement...\n>>>>>>> REPLACE\n\n"
            f"Use multiple blocks if the change touches multiple separate places. "
            f"Include at least 2-3 lines of context in each SEARCH block (not a "
            f"single short line) so it's uniquely identifiable in the file. "
            f"Copy the SEARCH text EXACTLY as it appears above — do not paraphrase, "
            f"re-indent, or reconstruct it from memory.\n\n"
            f"If the change is too pervasive for a patch (e.g. a full redesign), "
            f"reply instead with a single line `FULL_REWRITE` followed by the "
            f"complete new file content.\n\n"
            f"Output ONLY the diff/rewrite — no explanation, no markdown fences "
            f"wrapped around the whole thing."
        )

    system_msg  = {"role": "system",
                   "content": f"Expert {framework} developer making a precise, minimal edit to an existing file."}
    prompt_text = _initial_prompt()
    attempts    = 0
    total       = max_rounds * len(cascade)
    last_reason = "no attempts made"

    for _round in range(max_rounds):
        for cfg in cascade:
            if _HEALTH_OK and not _health_registry.is_available(cfg.get("model", "")):
                last_reason = f"{cfg.get('label', cfg.get('model'))}: circuit open, skipped"
                continue
            attempts += 1
            label = cfg.get("label", cfg.get("model", "?"))
            if on_attempt:
                try:
                    await on_attempt(label, attempts, total)
                except Exception:
                    pass

            messages = [system_msg, {"role": "user", "content": prompt_text}]
            diff_text = ""
            try:
                async for tok in stream_llm_func(
                    messages, cfg["url"], cfg["key"], cfg["model"],
                    cfg.get("temp", 0.1), max_tokens
                ):
                    diff_text += tok
            except Exception as e:
                last_reason = f"{label} request failed: {str(e)[:120]}"
                continue

            diff_text = diff_text.strip()
            if diff_text.startswith("```"):
                diff_text = re.sub(r'^```\w*\n?', '', diff_text)
                diff_text = re.sub(r'\n?```$', '', diff_text)

            merge_result = merge_diff(path, current_content, diff_text)
            if not merge_result.success:
                last_reason = f"{label}: merge failed — {merge_result.error}"
                # Ground the NEXT attempt in the real failure + real file
                # content instead of repeating the same prompt blind.
                prompt_text = merge_result.retry_prompt
                continue

            ok, reason = validate_file_content(path, merge_result.content)
            if ok:
                return {"content": merge_result.content, "ok": True, "reason": "",
                        "model": label, "attempts": attempts, "tiers": merge_result.match_tiers}
            last_reason = f"{label}: merged but failed content validation — {reason}"
            prompt_text = _initial_prompt()  # reset to a clean ask, last one wasn't the diff's fault

    return {"content": current_content, "ok": False, "reason": last_reason,
            "model": "exhausted", "attempts": attempts, "tiers": []}


async def generate_with_model_rotation(
    messages:      list,
    cascade:       list,
    stream_llm_func,
    path:          str = "",
    max_rounds:    int = 2,
    max_tokens:    int = 8192,
    on_attempt=None,   # optional async callback(model_label, attempt_no, total)
) -> dict:
    """
    Try every model in `cascade`, in order, up to `max_rounds` full passes —
    stopping the instant validate_file_content() passes. A model that
    "succeeds" with garbage content gets skipped to the NEXT model, not
    retried on itself; the same broken prompt + same broken model usually
    just produces the same broken answer twice.

    `cascade` is a list of dicts: {"url","key","model","temp","label"} — the
    same shape as brain.CODING_MODEL_CASCADE.

    Returns: {"content","ok","reason","model","attempts"}
    """
    best_content = ""
    best_reason  = "no attempts made"
    attempts     = 0
    total        = max_rounds * len(cascade)

    for _round in range(max_rounds):
        for cfg in cascade:
            # Phase 7 hardening: apply the same circuit-breaker awareness
            # from the DAG orchestrator (Phase 5) to this second engine —
            # skip a model whose circuit is already open instead of
            # wasting a round trip on it every single round.
            if _HEALTH_OK and not _health_registry.is_available(cfg.get("model", "")):
                best_reason = f"{cfg.get('label', cfg.get('model'))}: circuit open, skipped"
                continue
            attempts += 1
            label = cfg.get("label", cfg.get("model", "?"))
            if on_attempt:
                try:
                    await on_attempt(label, attempts, total)
                except Exception:
                    pass

            content = ""
            try:
                async for tok in stream_llm_func(
                    messages, cfg["url"], cfg["key"], cfg["model"],
                    cfg.get("temp", 0.1), max_tokens
                ):
                    content += tok
            except Exception as e:
                best_reason = f"{label} request failed: {str(e)[:100]}"
                continue

            content = clean_llm_code(content)
            ok, reason = validate_file_content(path, content)
            if ok:
                return {"content": content, "ok": True, "reason": "",
                         "model": label, "attempts": attempts}
            if len(content) > len(best_content):
                best_content, best_reason = content, f"{label}: {reason}"

    # Every model, every round, failed validation — return the best partial
    # result found so the caller can still keep the old file rather than
    # nothing, but flag ok=False so it's reported honestly, not as "done".
    return {"content": best_content, "ok": False, "reason": best_reason,
            "model": "exhausted", "attempts": attempts}