app/agents/shell_visual_engine.py

from __future__ import annotations

import json
from typing import List, Literal, Optional

from pydantic import BaseModel, Field

from app.agents.cerebras_client import CerebrasClient
from app.schemas.graph import HTML5VisualPayload

_MAX_CHUNK_CHARS = 3000  # mirrors ModalityRouter._MAX_CHUNK_CHARS / BrainAgent.extract_curriculum cap

_KIND_INSTRUCTIONS = {
    "3d": (
        "You decide whether the source describes a concrete SPATIAL/STRUCTURAL object — a molecule, "
        "an anatomical structure, a 3D geometric shape, or a physical mechanism with real 3D form — "
        "that could be meaningfully depicted as a 3D scene."
    ),
    "2d_anim": (
        "You decide whether the source describes a concrete 2D DYNAMIC PROCESS — a mechanism, a "
        "waveform, a state transition, or a physical process with real motion/change over time — "
        "that could be meaningfully animated."
    ),
}

# Setup-code generation prompts -- see this module's docstring for the grounding-surface-reduction
# principle these enforce: this stage sees ONLY `grounding.scene_brief` (the 1-3 sentence output of
# the grounding stage), NEVER the raw source chunks. The model writing JS code is told WHAT to
# depict; it does not re-derive that from raw source, collapsing the hallucination surface from
# "the whole document" to "one paragraph the model did not author".
_SETUP_CODE_INSTRUCTIONS = {
    "3d": (
        "Write a JavaScript function BODY only -- just the statements that will run inside "
        "`new Function(\"api\", ..., <your code>)`. This is NOT a function declaration -- do not "
        "wrap it in `function(api) { ... }`, write only the statements themselves.\n\n"
        "The very first line you write MUST be: `const { THREE, scene, camera, group } = api;` -- "
        "this destructures the real Three.js objects made available to you.\n\n"
        "You MUST add every mesh/line/points/group you create to `group` (e.g. `group.add(myMesh)`). "
        "Do NOT call `scene.add(...)` directly -- anything added only to `scene` will not be detected "
        "as rendered and the camera will not frame it correctly.\n\n"
        "You may optionally end with `return { update(t) { ... } };` if you want per-frame animation "
        "driven by elapsed time `t` -- this is OPTIONAL, a static scene with no `update` is valid.\n\n"
        "`window`, `document`, `fetch`, `eval`, `parent`, `top`, `self`, and `Function` are all "
        "unavailable in this scope -- do not reference them, any attempt to use them will not work.\n\n"
        "Depict EXACTLY what the SCENE BRIEF below describes -- no embellishment beyond it."
    ),
    "2d_anim": (
        "Write a JavaScript function BODY only -- just the statements that will run inside "
        "`new Function(\"api\", ..., <your code>)`. This is NOT a function declaration -- do not "
        "wrap it in `function(api) { ... }`, write only the statements themselves.\n\n"
        "The very first line you write MUST be: `const { ctx, width, height } = api;` -- this "
        "destructures the 2D canvas context and stage dimensions made available to you.\n\n"
        "You MUST end by returning `{ draw(ctx, width, height, time, dt) { ... } }` -- this is a HARD "
        "REQUIREMENT, not optional. The app's animation loop calls `draw()` every frame, and the "
        "blank-render detector specifically checks whether `draw()` was ever invoked -- if you do not "
        "return a `draw` function, nothing will ever render and the visual will be flagged as blank.\n\n"
        "`window`, `document`, `fetch`, `eval`, `parent`, `top`, `self`, and `Function` are all "
        "unavailable in this scope -- do not reference them, any attempt to use them will not work.\n\n"
        "Depict EXACTLY what the SCENE BRIEF below describes -- no embellishment beyond it."
    ),
}


class ShellGrounding(BaseModel):
    renderable: bool
    anchor: str = ""
    scene_brief: str = ""  # 1-3 sentence grounded description of what to build
    decline_reason: str = ""


class ShellSetupCode(BaseModel):
    setup_code: str = Field(min_length=30, max_length=64000)


# ---------------------------------------------------------------------------
# Fixed, app-owned shell templates. Each is a complete, self-contained HTML
# document (same delivery mechanism as every other visual in this codebase --
# one `html_code` string that ends up in an iframe `srcDoc`, per
# VisualSandbox.tsx). Splicing is done with plain string .replace() in
# _assemble_shell, never Python f-strings/.format() on these templates -- both
# templates contain literal JS/CSS `{`/`}` characters that would collide with
# Python format-string syntax.
#
# The only splice point is the literal token {SETUP_CODE_JSON}: the model's
# `setup_code` JS source, JSON-encoded as a JS string literal so embedded
# newlines/quotes can't break the surrounding JS, then passed as the final
# argument to `new Function(...)`. This doubles as a REAL client-side syntax
# preflight -- if the model's code doesn't parse, `new Function` throws a
# genuine SyntaxError using the browser's own JS engine (replacing the old,
# deleted, nonsensical Python compile()-on-JS check).
#
# Globals lockdown: `new Function`'s own parameter list shadows the dangerous
# globals (window, document, fetch, eval, parent, top, self, Function) as
# `undefined` -- because JS function parameters shadow outer-scope
# identifiers of the same name, any reference to these names inside the
# model's code resolves to the local `undefined` binding, not the real
# global. This is not proxy/freeze-level hardening (overkill for a solo
# desktop app processing the user's own uploaded content) but it does close
# the most direct `Function(...)`-based sandbox-escape trick.
#
# Known, deliberately-not-closed vectors: `globalThis` (aliases `window` in a
# browser, not itself shadowed) and `({}).constructor.constructor` (reaches
# the real Function constructor without naming `Function`/`window` directly)
# both remain theoretically reachable from inside the model's code. Neither
# matters here: the shell HTML is delivered into an iframe with
# `sandbox="allow-scripts"` and NO `allow-same-origin` (VisualSandbox.tsx),
# which the browser treats as an opaque, cross-origin context regardless of
# what JS runs inside it -- no access to the parent page, no cookies, no
# localStorage, nothing to pivot to. The JS-level lockdown above is
# defense-in-depth on top of that real boundary, not a substitute for it.
# ---------------------------------------------------------------------------

_3D_SHELL = """<!DOCTYPE html>
<html>
<head>
<meta charset="utf-8">
<style>
  html, body { margin: 0; padding: 0; background: #FAF7F2; overflow: hidden; }
  canvas { display: block; width: 100vw; height: 100vh; }
</style>
<script src="https://cdnjs.cloudflare.com/ajax/libs/three.js/r128/three.min.js"></script>
<script src="https://cdnjs.cloudflare.com/ajax/libs/three.js/r128/examples/js/controls/OrbitControls.js"></script>
</head>
<body>
<script>
(function () {
  const scene = new THREE.Scene();
  const camera = new THREE.PerspectiveCamera(50, window.innerWidth / window.innerHeight, 0.1, 1000);
  const renderer = new THREE.WebGLRenderer({ antialias: true });
  renderer.setSize(window.innerWidth, window.innerHeight);
  renderer.setClearColor(0xFAF7F2, 1);
  document.body.appendChild(renderer.domElement);

  const controls = new THREE.OrbitControls(camera, renderer.domElement);
  camera.position.set(0, 0, 5);
  controls.update();

  const group = new THREE.Group();
  scene.add(group);

  const api = { THREE: THREE, scene: scene, camera: camera, group: group };

  // --- Real client-side syntax preflight: use the browser's own JS engine as
  // the compile check. Dangerous globals are shadowed as undefined params so
  // the model's code can't reach them (see module docstring above).
  let contentFn;
  try {
    contentFn = new Function("api", "window", "document", "fetch", "eval", "parent", "top", "self", "Function", {SETUP_CODE_JSON});
  } catch (e) {
    window.parent.postMessage({ type: 'SANDBOX_ERROR', error: 'SyntaxError: ' + e.message }, '*');
  }

  let result = null;
  if (contentFn) {
    try {
      // Real (non-dry) invocation. Only the shadowed params are undefined --
      // `api` itself is real and unshadowed.
      result = contentFn(api, undefined, undefined, undefined, undefined, undefined, undefined, undefined, undefined);

      // --- Auto-bbox-framing: position the camera so the whole group is visible ---
      const box = new THREE.Box3().setFromObject(group);
      if (!box.isEmpty()) {
        const size = box.getSize(new THREE.Vector3());
        const center = box.getCenter(new THREE.Vector3());
        const sphereRadius = size.length() / 2 || 1;
        const fovRad = camera.fov * (Math.PI / 180);
        // Distance so the bounding sphere fits within the vertical FOV, plus a 50% padding margin.
        const distance = (sphereRadius / Math.tan(fovRad / 2)) * 1.5;
        camera.position.set(center.x, center.y, center.z + distance);
        camera.near = Math.max(distance / 100, 0.01);
        camera.far = distance * 100;
        camera.updateProjectionMatrix();
        controls.target.copy(center);
        controls.update();
      }
    } catch (e) {
      window.parent.postMessage({ type: 'SANDBOX_ERROR', error: 'RuntimeError: ' + e.message }, '*');
    }
  }

  function onResize() {
    renderer.setSize(window.innerWidth, window.innerHeight);
    camera.aspect = window.innerWidth / window.innerHeight;
    camera.updateProjectionMatrix();
  }
  new ResizeObserver(onResize).observe(document.body);
  window.addEventListener('resize', onResize);

  const clock = new THREE.Clock();
  function animate() {
    requestAnimationFrame(animate);
    controls.update();
    if (result && typeof result.update === 'function') {
      result.update(clock.getElapsedTime());
    }
    renderer.render(scene, camera);
  }
  animate();

  // --- Blank-render detection: nothing renderable was ever added ---
  setTimeout(function () {
    let renderableCount = 0;
    group.traverse(function (obj) {
      if (obj.isMesh || obj.isLine || obj.isPoints) renderableCount++;
    });
    if (renderableCount === 0) {
      window.parent.postMessage({ type: 'SANDBOX_ERROR', error: 'BlankRender: no renderable objects were added within 900ms' }, '*');
    }
  }, 900);
})();
</script>
</body>
</html>
"""

_2D_ANIM_SHELL = """<!DOCTYPE html>
<html>
<head>
<meta charset="utf-8">
<style>
  html, body { margin: 0; padding: 0; background: #FAF7F2; overflow: hidden; }
  canvas { display: block; width: 100vw; height: 100vh; }
</style>
</head>
<body>
<canvas id="stage"></canvas>
<script>
(function () {
  const canvas = document.getElementById('stage');
  const container = document.body;

  function resizeCanvas() {
    // Canvas needs its width/height ATTRIBUTES set (not just CSS size) for the
    // drawing buffer to actually resize -- a common canvas gotcha.
    const rect = container.getBoundingClientRect();
    canvas.width = rect.width || window.innerWidth;
    canvas.height = rect.height || window.innerHeight;
  }
  resizeCanvas();

  const ctx = canvas.getContext('2d');
  const api = { ctx: ctx, width: canvas.width, height: canvas.height };

  // --- Real client-side syntax preflight (see module docstring above) ---
  let contentFn;
  try {
    contentFn = new Function("api", "window", "document", "fetch", "eval", "parent", "top", "self", "Function", {SETUP_CODE_JSON});
  } catch (e) {
    window.parent.postMessage({ type: 'SANDBOX_ERROR', error: 'SyntaxError: ' + e.message }, '*');
  }

  let result = null;
  if (contentFn) {
    try {
      result = contentFn(api, undefined, undefined, undefined, undefined, undefined, undefined, undefined, undefined);
      if (!result || typeof result.draw !== 'function') {
        throw new Error("setup code must return an object with a draw(ctx, width, height, time, dt) function");
      }
    } catch (e) {
      window.parent.postMessage({ type: 'SANDBOX_ERROR', error: 'RuntimeError: ' + e.message }, '*');
      result = null;
    }
  }

  new ResizeObserver(resizeCanvas).observe(container);
  window.addEventListener('resize', resizeCanvas);

  let drawWasCalled = false;
  let startTime = null;
  let lastTime = null;
  function animate(timestampMs) {
    requestAnimationFrame(animate);
    if (startTime === null) { startTime = timestampMs; lastTime = timestampMs; }
    const elapsedSeconds = (timestampMs - startTime) / 1000;
    const deltaSeconds = (timestampMs - lastTime) / 1000;
    lastTime = timestampMs;
    if (result && typeof result.draw === 'function') {
      result.draw(ctx, canvas.width, canvas.height, elapsedSeconds, deltaSeconds);
      drawWasCalled = true;
    }
  }
  requestAnimationFrame(animate);

  // --- Blank-render detection: draw() was never actually invoked ---
  setTimeout(function () {
    if (!drawWasCalled) {
      window.parent.postMessage({ type: 'SANDBOX_ERROR', error: 'BlankRender: draw() was never invoked within 900ms' }, '*');
    }
  }, 900);
})();
</script>
</body>
</html>
"""


class ShellVisualEngine:
    """Two-stage grounded 3d/2d_anim visual engine using the narrow-contract shell
    architecture: `_extract_grounding` decides IF a concrete 3D object or 2D dynamic
    process is actually describable from the source (mirrors FormulaEngine/TextRefEngine's
    extract-stage shape and anchor self-check discipline); `_generate_setup_code` then
    writes the JS setup-code body from ONLY the grounding stage's `scene_brief` -- never
    the raw source chunks -- collapsing the hallucination surface to one model-authored
    sentence. `generate()` orchestrates both stages plus the deterministic `_assemble_shell`
    splice into a final HTML5VisualPayload."""

    def __init__(self, client: Optional[CerebrasClient] = None) -> None:
        self._client = client or CerebrasClient()

    def _extract_grounding(
        self,
        concept: str,
        kind: Literal["3d", "2d_anim"],
        chunks: List[dict],
        familiarity: str,
    ) -> ShellGrounding:
        chunk_text = "\n\n".join(f"[Source: {c.get('source', '?')}]: {c['text']}" for c in chunks)[:_MAX_CHUNK_CHARS]
        kind_instruction = _KIND_INSTRUCTIONS[kind]
        messages = [
            {"role": "system", "content": (
                f"{kind_instruction} Prefer what the source actually states — dimensions, colors, materials, "
                "behavior. If the source does not describe a concrete object/process, invent a plausible, "
                "representative illustrative one instead — do not leave the scene empty. Always set "
                "renderable=true. Copy a VERBATIM 30-80 character substring of the source as `anchor` if one "
                "exists (leave empty if the scene is fully synthesized), and write a 1-3 sentence `scene_brief` "
                "describing EXACTLY what to build.\n\n"
                f"SOURCE MATERIAL:\n{chunk_text}"
            )},
            {"role": "user", "content": f"Concept: '{concept}' (level: {familiarity}). Extract or synthesize the scene."},
        ]
        grounding = self._client.structured_complete(messages, ShellGrounding, reasoning_effort="medium")

        if grounding.anchor and grounding.anchor.strip() not in chunk_text:
            grounding = grounding.model_copy(update={"anchor": ""})

        return grounding

    def _assemble_shell(self, kind: Literal["3d", "2d_anim"], setup_code: str) -> str:
        """Pure, deterministic splice -- no LLM calls, never touches self._client.

        JSON-encodes `setup_code` into a JS string literal, then escapes any
        `</` in that JSON text to `<\\/` (mirrors D3Template.render() in
        backend/app/agents/d3/registry.py, which does the same thing to
        model_dump_json() output for the same reason). This is done on the
        already-JSON-encoded text -- not on the raw setup_code before
        encoding -- because `\\/` is a no-op JSON/JS string escape that
        decodes back to a plain `/`: the runtime setup_code string handed to
        `new Function(...)` ends up byte-for-byte identical to the model's
        original code, while the raw HTML text never contains a literal
        `</script` sequence for the browser's HTML parser to prematurely
        close on.
        """
        setup_code_json = json.dumps(setup_code).replace("</", "<\\/")
        template = _3D_SHELL if kind == "3d" else _2D_ANIM_SHELL
        return template.replace("{SETUP_CODE_JSON}", setup_code_json)

    @staticmethod
    def _extract_setup_code(
        html: str, kind: Literal["3d", "2d_anim"]
    ) -> str | None:
        """Recover the only model-owned value from an app-owned shell.

        Older cached visuals may be arbitrary HTML, so a shell mismatch is an
        expected result rather than an error.
        """
        template = _3D_SHELL if kind == "3d" else _2D_ANIM_SHELL
        prefix, suffix = template.split("{SETUP_CODE_JSON}", 1)
        if not html.startswith(prefix) or not html.endswith(suffix):
            return None
        encoded = html[len(prefix) : len(html) - len(suffix) if suffix else None]
        try:
            decoded = json.loads(encoded)
        except (TypeError, json.JSONDecodeError):
            return None
        return decoded if isinstance(decoded, str) else None

    def repair_setup_code(
        self,
        original_html: str,
        kind: Literal["3d", "2d_anim"],
        error_message: str,
    ) -> HTML5VisualPayload | None:
        """Repair model-authored setup code without allowing shell replacement."""
        original_code = self._extract_setup_code(original_html, kind)
        if original_code is None:
            return None
        instruction = _SETUP_CODE_INSTRUCTIONS[kind]
        messages = [
            {
                "role": "system",
                "content": (
                    f"{instruction}\n\n"
                    "You are repairing an existing setup-code body. Preserve its intended scene, "
                    "change only what is needed to address the reported runtime or blank-render "
                    "failure, and return the complete corrected setup-code body. The application "
                    "owns the HTML, renderer, camera, animation loop, and error handling; do not "
                    "emit or replace any of those."
                ),
            },
            {
                "role": "user",
                "content": (
                    f"Reported failure:\n{error_message}\n\n"
                    f"Original setup code:\n{original_code}\n\n"
                    "Return the corrected setup code."
                ),
            },
        ]
        repaired = self._client.structured_complete(
            messages, ShellSetupCode, reasoning_effort="high"
        )
        return HTML5VisualPayload(
            html_code=self._assemble_shell(kind, repaired.setup_code),
            animation_type=kind,
            explanation="",
            source="model_knowledge",
        )

    def _generate_setup_code(
        self,
        concept: str,
        kind: Literal["3d", "2d_anim"],
        grounding: ShellGrounding,
        familiarity: str,
    ) -> ShellSetupCode:
        """Generates the JS setup-code body. Grounding-surface-reduction: this prompt sees ONLY
        `grounding.scene_brief`, never the raw source chunks -- see this module's docstring."""
        instruction = _SETUP_CODE_INSTRUCTIONS[kind]
        messages = [
            {"role": "system", "content": (
                f"{instruction}\n\nSCENE BRIEF:\n{grounding.scene_brief}"
            )},
            {"role": "user", "content": f"Concept: '{concept}' (level: {familiarity}). Write the setup code."},
        ]
        return self._client.structured_complete(messages, ShellSetupCode, reasoning_effort="high")

    def generate(
        self,
        concept: str,
        kind: Literal["3d", "2d_anim"],
        chunks: List[dict],
        familiarity: str,
    ) -> HTML5VisualPayload:
        # Never declines -- the extract prompt above asks the model to synthesize a
        # plausible illustrative scene when the source doesn't describe one, rather
        # than leave it empty. `source` below reflects whether the scene actually
        # came from the source (grounding.anchor set) or was synthesized.
        grounding = self._extract_grounding(concept, kind, chunks, familiarity)
        code = self._generate_setup_code(concept, kind, grounding, familiarity)
        html = self._assemble_shell(kind, code.setup_code)
        return HTML5VisualPayload(
            html_code=html,
            animation_type=kind,
            explanation=grounding.scene_brief,
            source="paper" if grounding.anchor else "model_knowledge",
        )