agent/loop.py

from __future__ import annotations

import json
import logging
from typing import Generator

logger = logging.getLogger(__name__)


def _ev(event: str, **kw) -> dict:
    return {"event": event, **kw}


def stream_turn(user_message: str, state: dict) -> Generator[dict, None, None]:
    """
    Streaming agent turn. Yields event dicts consumed by app.py on_chat().

    Event types:
      {"event": "step",     "step": str,  "label": str}
      {"event": "thinking", "text": str}
      {"event": "result",   "state": dict, "reply": str, "code": str, "spec": dict|None}
      {"event": "error",    "state": dict, "message": str}
    """
    import copy
    from agent.intent   import classify
    from agent.healer   import heal
    from agent.memory   import append_user, append_assistant
    from agent.patcher  import patch as patch_spec
    from core.templates import render_page, render_app
    from core.sandbox   import safe_exec
    from core.validator import validate_page_spec
    from storage.db     import save_component
    from core.generator import MODEL_ID, get_device, _generate_messages, parse_json, _split_thinking
    from agent.planner  import _load_agent_system, _ensure_page_schema

    state = copy.deepcopy(state)
    state = append_user(state, user_message)

    current_spec = state.get("current_spec")
    messages     = state["messages"]

    # ── classify ─────────────────────────────────────────────────────────────
    yield _ev("step", step="classifying", label="classifying intent…")
    intent = classify(user_message, current_spec)
    logger.info("intent=%s", intent)

    # ── explain shortcut ──────────────────────────────────────────────────────
    if intent == "EXPLAIN":
        yield _ev("step", step="planning", label="generating explanation…")
        reply = _explain(user_message, current_spec, messages)
        state = append_assistant(state, reply)
        state["_reply"] = reply
        yield _ev("result", state=state, reply=reply, code="", spec=current_spec)
        return

    # ── plan or patch ─────────────────────────────────────────────────────────
    if intent in ("NEW",) or current_spec is None:
        step_label = "planning page…"
    elif intent == "RESTYLE":
        step_label = "restyling CSS…"
    else:
        step_label = f"patching page ({intent})…"

    yield _ev("step", step="planning", label=step_label)

    try:
        if intent == "NEW" or current_spec is None:
            system  = _load_agent_system()
            history = [m for m in messages if m["role"] != "system"]

            prompt_msgs = (
                [{"role": "system", "content": system}]
                + history
                + [{"role": "user", "content": user_message}]
            )

            raw = _generate_messages(prompt_msgs, max_tokens=2048)

            thinking, json_raw = _split_thinking(raw)
            if not json_raw.strip():
                json_raw = raw
            if thinking:
                yield _ev("thinking", text=thinking)

            spec = _ensure_page_schema(parse_json(json_raw))

        else:
            # EDIT / ADD / REMOVE / REPLACE / WIRE / FIX / RESTYLE
            spec = patch_spec(current_spec, user_message, intent, messages)

    except Exception as exc:
        logger.exception("plan/patch failed")
        err_msg = f"I couldn't understand that request ({exc}). Try rephrasing."
        state = _fail(state, err_msg)
        yield _ev("error", state=state, message=err_msg)
        return

    # ── validate page spec ────────────────────────────────────────────────────
    ok_spec, val_errors = validate_page_spec(spec)
    if not ok_spec:
        logger.warning("Page spec validation errors: %s", val_errors)
        # non-fatal — still attempt render; healer may fix exec errors

    # ── render ────────────────────────────────────────────────────────────────
    yield _ev("step", step="rendering", label="rendering page…")
    try:
        code = render_page(spec)
    except Exception as exc:
        logger.exception("render_page failed")
        # Try legacy render_app as fallback
        try:
            code = render_app(spec)
        except Exception as exc2:
            err_msg = f"Render error: {exc2}"
            state = _fail(state, err_msg)
            yield _ev("error", state=state, message=err_msg)
            return

    # ── sandbox ───────────────────────────────────────────────────────────────
    yield _ev("step", step="validating", label="validating in sandbox…")
    exec_ok, exec_err = safe_exec(code)

    # ── heal ──────────────────────────────────────────────────────────────────
    heal_count = 0
    if not exec_ok:
        yield _ev("step", step="healing", label="auto-healing error…")
        healed_code, healed_ok = heal(code, exec_err, messages, spec=spec)
        heal_count = 1
        if healed_ok:
            code    = healed_code
            exec_ok = True
        else:
            _trace(user_message, code, spec, False, False, exec_err, get_device(), MODEL_ID)
            err_msg = _friendly_error(exec_err)
            state   = _fail(state, err_msg, spec=spec, code=code)
            yield _ev("error", state=state, message=err_msg)
            return

    # ── save ──────────────────────────────────────────────────────────────────
    yield _ev("step", step="saving", label="saving…")
    try:
        save_component(user_message, code, spec)
    except Exception:
        logger.exception("save_component failed")

    _trace(user_message, code, spec, ok_spec, True, None, get_device(), MODEL_ID)

    # ── commit ────────────────────────────────────────────────────────────────
    page      = spec.get("page", spec)
    reasoning = spec.get("reasoning", "")
    reply     = _format_reply(page, reasoning, intent)

    state = append_assistant(state, json.dumps(spec))
    state["current_spec"]  = spec
    state["current_code"]  = code
    state["iteration"]    += 1
    state["heal_attempts"] = heal_count
    state["last_error"]    = None
    state["_reply"]        = reply

    yield _ev("result", state=state, reply=reply, code=code, spec=spec)


def run_turn(user_message: str, state: dict) -> dict:
    """Non-streaming wrapper — consumes stream_turn and returns final state."""
    final_state = state
    for event in stream_turn(user_message, state):
        if event["event"] in ("result", "error"):
            final_state = event["state"]
    return final_state


def _explain(user_message: str, current_spec, messages: list[dict]) -> str:
    if current_spec is None:
        return "There's no page built yet. Describe an app and I'll build it."
    try:
        from agent.planner import _call_model, _load_agent_system
        system = _load_agent_system()
        prompt = (
            f"The current page spec is:\n{json.dumps(current_spec, indent=2)}\n\n"
            f"User question: {user_message}\n\n"
            "Answer clearly and concisely. No JSON output needed."
        )
        return _call_model(
            [{"role": "system", "content": system}, {"role": "user", "content": prompt}]
        )
    except Exception as exc:
        return f"I can see the current page but couldn't generate an explanation ({exc})."


def _format_reply(page: dict, reasoning: str, intent: str = "") -> str:
    sections   = page.get("sections",  [])
    title      = page.get("title",     "App")
    layout     = page.get("layout",    "")
    n_sections = len(sections)
    n_comps    = sum(len(s.get("components", [])) for s in sections)
    action     = intent.lower() or "built"

    lines = [
        f"**{action.capitalize()}**: {title} · {layout} layout · "
        f"{n_sections} section{'s' if n_sections != 1 else ''} · "
        f"{n_comps} component{'s' if n_comps != 1 else ''}"
    ]
    if reasoning:
        lines.append(f"_{reasoning}_")
    return "\n\n".join(lines)


def _friendly_error(raw_error: str) -> str:
    lower = raw_error.lower()
    if "blocked import" in lower:
        # Extract which module was blocked for a more helpful message
        import re
        m = re.search(r"blocked import[^:]*:\s*(.+)", raw_error, re.IGNORECASE)
        detail = m.group(1).strip() if m else ""
        suffix = f" ({detail})" if detail else ""
        return (
            f"The generated code used a library that isn't available in the sandbox{suffix}. "
            "Try rephrasing — e.g. 'show a plotly scatter chart' instead of 'use sklearn KMeans'."
        )
    if "nameerror"  in lower:
        return "The generated code referenced something that doesn't exist. I tried to fix it but couldn't."
    if "syntaxerror" in lower:
        return "The generated code had a syntax problem I couldn't auto-repair."
    if "importerror" in lower or "modulenot" in lower:
        return "The generated code tried to import something unavailable."
    if "timeout"    in lower or "timed out" in lower:
        return "The generated page took too long to start. Try a simpler request."
    return "The generated page had an error I couldn't fix automatically. Please try rephrasing."


def _fail(state: dict, message: str, spec=None, code: str = "") -> dict:
    from agent.memory import append_assistant
    state = append_assistant(state, message)
    state["_reply"]     = message
    state["last_error"] = message
    if spec is not None:
        state["current_spec"] = spec
    if code:
        state["current_code"] = code
    return state


def _trace(prompt, code, spec, valid, exec_ok, exec_err, device, model_id):
    try:
        from storage.trace_logger import log_turn
        log_turn(
            prompt       = prompt,
            raw_output   = code,
            parsed_spec  = spec,
            valid        = valid,
            exec_success = exec_ok,
            exec_error   = exec_err,
            device       = device,
            model_id     = model_id,
        )
    except Exception:
        logger.exception("trace logging failed")