utils.py

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

import re
import io
import os
import time
import asyncio
import logging
import zipfile
import hashlib
from typing import List, Dict, Optional, AsyncIterator
from collections import deque
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
# ═══════════════════════════════════════════════════════

def sse_event(data: dict) -> str:
    """Format dict as SSE event."""
    import json
    return f"data: {json.dumps(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:
            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:
            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}