"""
🧠 Session Memory — Reader + Writer + Feedback Editor nodes for ComfyUI.

Architecture (no graph cycle):

  🧠 Reader ──session_memory──► system prompt ──► LLM ──► Summarizer
      ↑                                                        │
      │                                                   💾 Writer
      │                                                        │
      │                                              ⭐ Feedback Editor
      │                                                        │
      └───────────────── sessions/{session_id}.json ◄──────────┘

Writer saves the run, then Feedback Editor shows it for rating/notes.
Reader loads everything (entries + feedback) at the start of the next Run.
"""

import json
import logging
import os
import re
import shutil
import threading
import time

from aiohttp import web
from server import PromptServer

logger = logging.getLogger(__name__)

# ── Session file helpers ───────────────────────────────────────────────
SESSIONS_DIR = os.path.join(os.path.dirname(__file__), "sessions")

# Threading events for blocking mode
_pending_feedback: dict[str, threading.Event] = {}


def _safe_id(session_id: str) -> str:
    """Sanitize session_id to a safe filename component."""
    return "".join(c for c in session_id if c.isalnum() or c in "-_") or "default"


def _session_path(session_id: str) -> str:
    os.makedirs(SESSIONS_DIR, exist_ok=True)
    return os.path.join(SESSIONS_DIR, f"{_safe_id(session_id)}.json")


def _load(session_id: str) -> dict:
    path = _session_path(session_id)
    if os.path.exists(path):
        try:
            with open(path, "r", encoding="utf-8") as f:
                data = json.load(f)
            # Migrate legacy string entries to {prompt, summary} dicts
            entries = data.get("entries", [])
            migrated = False
            for i, entry in enumerate(entries):
                if isinstance(entry, str):
                    entries[i] = {"prompt": entry, "summary": entry}
                    migrated = True
            if migrated:
                data["entries"] = entries
                logger.info(f"[Session Memory] Migrated {sum(1 for e in entries if isinstance(e, dict))} legacy entries in '{session_id}'.")
            return data
        except (json.JSONDecodeError, OSError) as e:
            logger.warning(f"[Session Memory] Could not read {path}: {e}")
    return {"run_count": 0, "entries": [], "feedback": {}}


def _save(session_id: str, data: dict) -> None:
    path = _session_path(session_id)
    try:
        with open(path, "w", encoding="utf-8") as f:
            json.dump(data, f, ensure_ascii=False, indent=2)
    except OSError as e:
        logger.error(f"[Session Memory] Could not write {path}: {e}")


def _parse_prompts_from_entry(entry_text: str) -> list[dict]:
    """Extract prompt headers, concepts, and full text from a run entry string.

    Supports three formats (tried in order):
    1. XML-style ``<prompt>text</prompt>`` blocks (Visual DNA Synthesizer)
    2. Legacy ``PROMPT N:`` markers
    3. Fallback: whole entry treated as a single Prompt 1
    """
    prompts = []

    # ── Strategy 1: XML-style <prompt> tags ──────────────────────────
    xml_blocks = re.findall(r'<prompt>\s*(.*?)\s*</prompt>', entry_text, re.DOTALL)
    if xml_blocks:
        for i, block in enumerate(xml_blocks, start=1):
            prompts.append({
                "number": i,
                "concept": f"Prompt {i}",
                "fullText": block.strip(),
            })
        return prompts

    # ── Strategy 2: PROMPT N: markers ────────────────────────────────
    prompt_blocks = re.split(r'(?=PROMPT \d+:)', entry_text)
    # Markers that indicate the start of a trailing summary section (not prompt content)
    _SUMMARY_MARKERS = re.compile(
        r'^(SESSION\s+TALLY|TALLY|SUMMARY|SESSION_FEEDBACK|RUN\s+\d+:)', re.MULTILINE
    )
    for block in prompt_blocks:
        m = re.match(r'PROMPT (\d+):', block)
        if m:
            prompt_num = int(m.group(1))
            concept_match = re.search(r'concept:\s*(.+)', block)
            concept = (concept_match.group(1).strip()
                       if concept_match else f"Prompt {prompt_num}")
            # Trim trailing session-level sections from the block
            full_text = block
            trail = _SUMMARY_MARKERS.search(full_text)
            if trail:
                full_text = full_text[:trail.start()]
            full_text = full_text.strip()
            prompts.append({
                "number": prompt_num,
                "concept": concept,
                "fullText": full_text,
            })

    if prompts:
        return prompts

    # ── Fallback: whole entry as Prompt 1 ────────────────────────────
    if entry_text.strip():
        full_text = entry_text.strip()
        trail = _SUMMARY_MARKERS.search(full_text)
        if trail:
            full_text = full_text[:trail.start()]

        prompts.append({
            "number": 1,
            "concept": "Result",
            "fullText": full_text.strip()
        })

    return prompts



# ── Thumbnail helpers ──────────────────────────────────────────────────

def _session_images_dir(session_id: str) -> str:
    """Get the image storage directory for a session."""
    return os.path.join(SESSIONS_DIR, _safe_id(session_id))


def _save_thumbnails(session_id, run_number, num_prompts, batch_count,
                     images_tensor):
    """Save batch images as downscaled JPEG thumbnails.

    Layout: sessions/{sid}/run_{N}/prompt_{P}/img_{K}.jpg
    where K is the batch/pass index (0-based).

    Image order in tensor: [p1, p2, ..., pN, p1, p2, ..., pN, ...]
    distributed round-robin across prompts.

    With ``batch_count > 1`` ComfyUI **re-executes the entire graph** for
    each batch pass, sending the *full* image tensor every time.  We must
    avoid saving duplicate images.  Strategy:

    * On the **first call** (prompt dirs are empty) → save all images.
    * On **subsequent calls** (prompt dirs already populated) → compute
      how many images *per prompt* this call would add.  If the folder
      already holds ``>= expected_per_prompt`` images, skip entirely
      (the images are duplicates from a graph re-execution).
    """
    import numpy as np
    from PIL import Image as PILImage

    base_dir = _session_images_dir(session_id)
    run_dir = os.path.join(base_dir, f"run_{run_number}")

    B = images_tensor.shape[0]

    logger.info(
        f"[Session Feedback] Saving thumbnails: B={B}, "
        f"num_prompts={num_prompts}, batch_count_widget={batch_count}"
    )

    # The user-specified batch_count is the authoritative expected number
    # of images per prompt.  Use it as the dedup threshold:
    #   - Combined tensor (B = num_prompts × batch_count): first call saves
    #     everything correctly; subsequent re-executions are skipped.
    #   - Partial batches (B = num_prompts per call): each call appends
    #     until batch_count is reached, then subsequent calls are skipped.
    expected_per_prompt = max(1, batch_count)

    first_prompt_dir = os.path.join(run_dir, "prompt_1")
    if os.path.isdir(first_prompt_dir):
        existing_count = len([
            f for f in os.listdir(first_prompt_dir)
            if f.startswith("img_") and f.endswith(".jpg")
        ])
        if existing_count >= expected_per_prompt:
            logger.info(
                f"[Session Feedback] Prompt_1 already has {existing_count} "
                f"images (>= {expected_per_prompt} expected per batch_count). "
                f"Skipping duplicate re-execution."
            )
            return

    saved = 0

    # Distribute images round-robin: image[i] → prompt (i % num_prompts)
    for i in range(B):
        p = i % num_prompts          # 0-based prompt index

        prompt_dir = os.path.join(run_dir, f"prompt_{p + 1}")
        os.makedirs(prompt_dir, exist_ok=True)

        # Find next available index for this prompt
        existing = [
            f for f in os.listdir(prompt_dir)
            if f.startswith("img_") and f.endswith(".jpg")
        ]
        next_idx = len(existing)

        # Get single image [H, W, C]
        img_t = images_tensor[i]
        h, w = img_t.shape[0], img_t.shape[1]

        # Downscale to max 1 megapixel
        current_mp = (h * w) / 1_000_000
        if current_mp > 1.0:
            scale = (1.0 / current_mp) ** 0.5
            new_h = max(1, int(h * scale))
            new_w = max(1, int(w * scale))
        else:
            new_h, new_w = h, w

        img_np = (img_t.cpu().numpy() * 255).clip(0, 255).astype(np.uint8)
        pil_img = PILImage.fromarray(img_np)

        if pil_img.mode != "RGB":
            pil_img = pil_img.convert("RGB")

        if (new_h, new_w) != (h, w):
            pil_img = pil_img.resize((new_w, new_h), PILImage.LANCZOS)

        out_path = os.path.join(prompt_dir, f"img_{next_idx}.jpg")
        pil_img.save(out_path, "JPEG", quality=85)
        saved += 1
        logger.debug(
            f"[Session Feedback]   image[{i}] → prompt_{p+1}/img_{next_idx}.jpg"
        )

    logger.info(
        f"[Session Feedback] Saved {saved} thumbnails for Run {run_number} "
        f"in '{session_id}'."
    )


def _delete_session_images(session_id: str) -> int:
    """Delete all thumbnail images for a session. Returns count removed."""
    img_dir = _session_images_dir(session_id)
    if os.path.isdir(img_dir):
        # Count files before deletion
        count = sum(len(files) for _, _, files in os.walk(img_dir))
        shutil.rmtree(img_dir, ignore_errors=True)
        logger.info(
            f"[Session Memory] Deleted {count} thumbnail images "
            f"for session '{session_id}'."
        )
        return count
    return 0


# ── Feedback formatting ───────────────────────────────────────────────

def _format_feedback_block(feedback: dict) -> str:
    """Format feedback into a SESSION_FEEDBACK block for the system prompt.

    Only includes actionable feedback:
      ★1-2 -> AVOID directive (negative constraint)
      ★4-5 -> REPLICATE directive (positive reference)
      ★3   -> omitted (neutral)
      ★0   -> omitted (unrated)
    """
    if not feedback:
        return ""

    lines = []
    for run_key in sorted(feedback.keys(), key=lambda x: int(x)):
        run_data = feedback[run_key]
        prompts_fb = run_data.get("prompts", {})
        run_lines = []

        for p_key in sorted(prompts_fb.keys(), key=lambda x: int(x)):
            p_data = prompts_fb[p_key]
            rating = p_data.get("rating", 0)
            notes = p_data.get("notes", "").strip()

            if rating == 0 or rating == 3:
                continue

            stars = "\u2605" * rating + "\u2606" * (5 - rating)
            directive = "AVOID" if rating <= 2 else "REPLICATE"

            line = f"  PROMPT {p_key} ({stars} \u2014 {directive})"
            if notes:
                line += f': "{notes}"'
            run_lines.append(line)

        if run_lines:
            lines.append(f"RUN {run_key}:")
            lines.extend(run_lines)

    if not lines:
        return ""

    return "SESSION_FEEDBACK:\n" + "\n".join(lines) + "\n"


# ── API Routes ─────────────────────────────────────────────────────────

@PromptServer.instance.routes.get("/session_feedback/load")
async def _api_load_feedback(request):
    """Load run data + existing feedback for the frontend editor.

    The feedback editor shows the PROMPT text (not the summary).
    """
    session_id = request.query.get("session_id", "default")
    run_number = request.query.get("run_number", "")

    session = _load(session_id)
    entries = session.get("entries", [])
    run_count = session.get("run_count", 0)
    feedback = session.get("feedback", {})

    if run_count == 0:
        return web.json_response({
            "run_count": 0,
            "selected_run": 0,
            "prompts": [],
            "feedback": {},
        })

    # Determine which run to show
    if run_number and run_number.isdigit():
        sel_run = int(run_number)
    else:
        sel_run = run_count

    sel_run = max(1, min(sel_run, run_count))

    idx = sel_run - 1
    entry = entries[idx] if idx < len(entries) else {}
    # Use the prompt field for display in the feedback editor
    if isinstance(entry, dict):
        entry_text = entry.get("prompt", "")
    else:
        entry_text = entry  # legacy string fallback
    prompts = _parse_prompts_from_entry(entry_text)
    run_feedback = feedback.get(str(sel_run), {}).get("prompts", {})

    return web.json_response({
        "run_count": run_count,
        "selected_run": sel_run,
        "prompts": prompts,
        "feedback": run_feedback,
    })


@PromptServer.instance.routes.post("/session_feedback/save")
async def _api_save_feedback(request):
    """Save feedback from the frontend editor to the session file."""
    try:
        body = await request.json()
    except Exception:
        return web.json_response({"error": "Invalid JSON"}, status=400)

    session_id = body.get("session_id", "default")
    run_number = str(body.get("run_number", ""))
    prompts_fb = body.get("prompts", {})

    if not run_number:
        return web.json_response({"error": "run_number required"}, status=400)

    session = _load(session_id)
    feedback = session.get("feedback", {})

    if run_number not in feedback:
        feedback[run_number] = {"prompts": {}}

    for p_key, p_data in prompts_fb.items():
        feedback[run_number]["prompts"][str(p_key)] = {
            "rating": int(p_data.get("rating", 0)),
            "notes": str(p_data.get("notes", "")),
        }

    session["feedback"] = feedback
    _save(session_id, session)

    deleted_thumbs = body.get("deleted_thumbs", [])
    if deleted_thumbs:
        del_count = 0
        for thumb_rel in deleted_thumbs:
            # Reconstruct absolute paths safely
            safe_rel = thumb_rel.lstrip("/\\")
            if ".." in safe_rel: 
                continue
            abs_path = os.path.join(SESSIONS_DIR, safe_rel)
            if os.path.isfile(abs_path):
                try:
                    os.remove(abs_path)
                    del_count += 1
                except OSError as e:
                    logger.warning(f"[Session Feedback] Could not delete thumb {thumb_rel}: {e}")
        if del_count > 0:
            logger.info(f"[Session Feedback] Permanently deleted {del_count} selected thumbnails.")

    logger.info(
        f"[Session Feedback] Saved feedback for Run {run_number} "
        f"in session '{session_id}' ({len(prompts_fb)} prompts)."
    )
    return web.json_response({"status": "ok"})


@PromptServer.instance.routes.post("/session_feedback/resume")
async def _api_resume_feedback(request):
    """Resume blocked execution after user saves feedback."""
    try:
        body = await request.json()
    except Exception:
        return web.json_response({"error": "Invalid JSON"}, status=400)

    node_id = str(body.get("node_id", ""))
    event = _pending_feedback.get(node_id)
    if event:
        event.set()
        logger.info(f"[Session Feedback] Resumed execution for node {node_id}.")
        return web.json_response({"status": "resumed"})

    return web.json_response({"status": "not_found"})


@PromptServer.instance.routes.post("/session_feedback/delete_run")
async def _api_delete_run(request):
    """Delete a specific run and its thumbnails."""
    try:
        body = await request.json()
    except Exception:
        return web.json_response({"error": "Invalid JSON"}, status=400)

    session_id = body.get("session_id", "default")
    run_number = str(body.get("run_number", ""))

    if not run_number:
        return web.json_response({"error": "run_number required"}, status=400)

    session = _load(session_id)
    
    del_run_int = int(run_number)
    entries = session.get("entries", [])
    run_count = session.get("run_count", len(entries))
    idx = del_run_int - 1

    # 1. Update entries
    if 0 <= idx < len(entries):
        entries.pop(idx)
    
    session["run_count"] = max(0, run_count - 1)
    
    # Renumber the headers inside the entries to maintain 1..N contiguity
    for i in range(len(entries)):
        if isinstance(entries[i], dict):
            for field in ("prompt", "summary"):
                if field in entries[i]:
                    entries[i][field] = re.sub(
                        r'^---\s*RUN\s+\d+\s*---\n',
                        f'--- RUN {i+1} ---\n',
                        entries[i][field], count=1, flags=re.MULTILINE,
                    )
        else:
            entries[i] = re.sub(r'^---\s*RUN\s+\d+\s*---\n', f'--- RUN {i+1} ---\n', entries[i], count=1, flags=re.MULTILINE)
        
    session["entries"] = entries

    # 2. Update feedback keys
    feedback = session.get("feedback", {})
    new_feedback = {}
    for k_str, v in feedback.items():
        try:
            k_int = int(k_str)
            if k_int < del_run_int:
                new_feedback[str(k_int)] = v
            elif k_int > del_run_int:
                new_feedback[str(k_int - 1)] = v
        except ValueError:
            pass
    session["feedback"] = new_feedback
    _save(session_id, session)

    # 3. Delete the target thumbnail folder
    safe_sid = _safe_id(session_id)
    del_dir = os.path.join(SESSIONS_DIR, safe_sid, f"run_{del_run_int}")
    if os.path.isdir(del_dir):
        shutil.rmtree(del_dir, ignore_errors=True)

    # 4. Shift subsequent thumbnail folders down by 1
    for k in range(del_run_int + 1, run_count + 1):
        old_dir = os.path.join(SESSIONS_DIR, safe_sid, f"run_{k}")
        new_dir = os.path.join(SESSIONS_DIR, safe_sid, f"run_{k-1}")
        if os.path.isdir(old_dir):
            try:
                os.rename(old_dir, new_dir)
            except OSError as e:
                logger.error(f"[Session Feedback] Rename failed: {e}")

    logger.info(f"[Session Feedback] Deleted Run {del_run_int} in session '{session_id}' and shifted remaining runs.")
    return web.json_response({"status": "deleted"})


@PromptServer.instance.routes.post("/session_feedback/delete_session")
async def _api_delete_session(request):
    """Completely delete a session JSON and all its image assets."""
    try:
        body = await request.json()
    except Exception:
        return web.json_response({"error": "Invalid JSON"}, status=400)

    session_id = body.get("session_id", "default")
    if not session_id:
        return web.json_response({"error": "Session ID required"}, status=400)

    safe_sid = _safe_id(session_id)
    
    # 1. Delete the JSON file
    json_path = os.path.join(SESSIONS_DIR, f"{safe_sid}.json")
    if os.path.isfile(json_path):
        os.remove(json_path)

    # 2. Delete the entire thumbnails folder
    img_dir = os.path.join(SESSIONS_DIR, safe_sid)
    if os.path.isdir(img_dir):
        shutil.rmtree(img_dir, ignore_errors=True)

    logger.warning(f"[Session Feedback] NUKED session '{session_id}' entirely.")
    return web.json_response({"status": "deleted"})


@PromptServer.instance.routes.post("/session_feedback/rename_session")
async def _api_rename_session(request):
    """Rename a session JSON and its thumbnail folder."""
    try:
        body = await request.json()
    except Exception:
        return web.json_response({"error": "Invalid JSON"}, status=400)

    old_id = body.get("old_session_id", "")
    new_id = body.get("new_session_id", "")
    
    if not old_id or not new_id:
        return web.json_response({"error": "old_session_id and new_session_id required"}, status=400)

    safe_old = _safe_id(old_id)
    safe_new = _safe_id(new_id)

    if safe_old == safe_new:
        return web.json_response({"status": "renamed"}) # Nothing to do
        
    old_json = os.path.join(SESSIONS_DIR, f"{safe_old}.json")
    new_json = os.path.join(SESSIONS_DIR, f"{safe_new}.json")
    
    if os.path.exists(new_json):
         return web.json_response({"error": f"Session '{new_id}' already exists."}, status=400)

    # Rename JSON
    if os.path.isfile(old_json):
        try:
            os.rename(old_json, new_json)
        except OSError as e:
            return web.json_response({"error": f"Could not rename file: {e}"}, status=500)
    else:
        # Create an empty session if it didn't exist before renaming
        _save(new_id, {"run_count": 0, "entries": [], "feedback": {}})

    # Rename images directory
    old_dir = os.path.join(SESSIONS_DIR, safe_old)
    new_dir = os.path.join(SESSIONS_DIR, safe_new)
    
    if os.path.isdir(old_dir):
        try:
            os.rename(old_dir, new_dir)
        except OSError as e:
            pass # Non-critical if it fails or doesn't exist

    logger.info(f"[Session Feedback] Renamed session from '{old_id}' to '{new_id}'.")
    return web.json_response({"status": "renamed"})


@PromptServer.instance.routes.get("/session_feedback/thumbnails")
async def _api_list_thumbnails(request):
    """List available thumbnail images for a specific run."""
    session_id = request.query.get("session_id", "default")
    run_number = request.query.get("run_number", "1")

    safe_sid = _safe_id(session_id)
    run_dir = os.path.join(SESSIONS_DIR, safe_sid, f"run_{run_number}")
    result = {}

    if os.path.isdir(run_dir):
        for entry in sorted(os.listdir(run_dir)):
            if not entry.startswith("prompt_"):
                continue
            prompt_path = os.path.join(run_dir, entry)
            if not os.path.isdir(prompt_path):
                continue
            p_num = entry.replace("prompt_", "")
            imgs = sorted(
                f for f in os.listdir(prompt_path)
                if f.lower().endswith((".jpg", ".jpeg", ".png"))
            )
            result[p_num] = [
                f"/session_feedback/thumb?path="
                f"{safe_sid}/run_{run_number}/{entry}/{img}"
                for img in imgs
            ]

    return web.json_response(result)


@PromptServer.instance.routes.get("/session_feedback/thumb")
async def _api_serve_thumbnail(request):
    """Serve a single thumbnail image file."""
    rel_path = request.query.get("path", "")
    if not rel_path:
        return web.Response(status=400, text="Missing path")

    full_path = os.path.realpath(os.path.join(SESSIONS_DIR, rel_path))
    sessions_real = os.path.realpath(SESSIONS_DIR)

    # Security: prevent directory traversal
    if not full_path.startswith(sessions_real):
        return web.Response(status=403, text="Forbidden")

    if not os.path.isfile(full_path):
        return web.Response(status=404, text="Not found")

    return web.FileResponse(full_path)


# ── Node: Feedback Editor ─────────────────────────────────────────────
class SessionFeedbackEditorNode:
    """
    ⭐ Session Feedback Editor

    Placed AFTER the Writer. Shows the latest Run for the director to
    rate (1-5 stars) and annotate each prompt. Optionally receives a
    batch of images to save as thumbnails for visual preview.

    Two modes:
      - non_blocking: execution completes, user annotates at leisure
      - blocking: execution pauses until user clicks Save
    """

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "session_id": ("STRING", {
                    "default": "default",
                    "tooltip": "Session identifier. Must match Reader/Writer.",
                }),
                "mode": (["non_blocking", "blocking"], {
                    "default": "non_blocking",
                    "tooltip": "blocking = pauses workflow until you save feedback. "
                               "non_blocking = workflow finishes, annotate later.",
                }),
                "num_prompts": ("INT", {
                    "default": 4,
                    "min": 1,
                    "max": 20,
                    "step": 1,
                    "tooltip": "Number of prompts per run (e.g. 4).",
                }),
                "batch_count": ("INT", {
                    "default": 1,
                    "min": 1,
                    "max": 50,
                    "step": 1,
                    "tooltip": "Times each prompt was sampled. "
                               "Total images = num_prompts × batch_count.",
                }),
            },
            "optional": {
                "run_summary": ("STRING", {
                    "forceInput": True,
                    "tooltip": "Connect from Writer to ensure execution order.",
                }),
                "thumbnail_images": ("IMAGE", {
                    "tooltip": "Batch of images from sampler. "
                               "Total = num_prompts × batch_count.",
                }),
            },
            "hidden": {
                "unique_id": "UNIQUE_ID",
            },
        }

    RETURN_TYPES = ()
    OUTPUT_NODE = True
    FUNCTION = "process"
    CATEGORY = "\U0001f9e0 Memory"

    @classmethod
    def IS_CHANGED(cls, **kwargs):
        return float("NaN")

    def process(self, session_id: str = "default", mode: str = "non_blocking",
                num_prompts: int = 4, batch_count: int = 1,
                run_summary: str = "", thumbnail_images=None,
                unique_id=None):
        session = _load(session_id)
        run_count = session.get("run_count", 0)

        if run_count == 0:
            logger.info("[Session Feedback Editor] No runs yet.")
            return {"ui": {"status": ["no_runs"]}}

        # Save thumbnails if images provided
        if thumbnail_images is not None:
            _save_thumbnails(
                session_id, run_count, num_prompts, batch_count,
                thumbnail_images,
            )

        # Get latest run prompts + existing feedback
        # Feedback editor shows the PROMPT text (not summary)
        entries = session.get("entries", [])
        latest_entry = entries[-1] if entries else {}
        if isinstance(latest_entry, dict):
            entry_text = latest_entry.get("prompt", "")
        else:
            entry_text = latest_entry  # legacy string fallback
        prompts = _parse_prompts_from_entry(entry_text)
        feedback = session.get("feedback", {})
        run_feedback = feedback.get(str(run_count), {}).get("prompts", {})

        # Notify frontend via websocket
        PromptServer.instance.send_sync("session_feedback_update", {
            "node_id": str(unique_id),
            "session_id": session_id,
            "run_count": run_count,
            "selected_run": run_count,
            "prompts": prompts,
            "feedback": run_feedback,
            "mode": mode,
        })

        if mode == "blocking":
            logger.info(
                f"[Session Feedback Editor] Blocking — waiting for feedback "
                f"on Run {run_count} (node {unique_id})..."
            )
            event = threading.Event()
            _pending_feedback[str(unique_id)] = event
            try:
                # Wait up to 10 minutes
                timeout = 600
                start = time.time()
                while not event.is_set():
                    if time.time() - start > timeout:
                        logger.warning(
                            "[Session Feedback Editor] Timeout — "
                            "continuing without feedback."
                        )
                        break
                    event.wait(timeout=1.0)
            finally:
                _pending_feedback.pop(str(unique_id), None)

            logger.info("[Session Feedback Editor] Resumed after feedback.")
        else:
            logger.info(
                f"[Session Feedback Editor] Non-blocking — Run {run_count} "
                f"ready for feedback in UI."
            )

        return {"ui": {"status": ["done"]}}


# ── Node: Reader ───────────────────────────────────────────────────────
class SessionMemoryReaderNode:
    """
    🧠 Session Memory Reader

    Placed at the start of the workflow. Reads the accumulated session log
    AND feedback from disk, outputs it for injection into the system prompt.
    """

    @classmethod
    def INPUT_TYPES(cls):
        sessions = _get_available_sessions()
        return {
            "required": {
                "existing_session": (["[CREATE NEW]"] + sessions, {
                    "default": sessions[-1] if sessions else "[CREATE NEW]",
                    "tooltip": "Select an existing session, or [CREATE NEW] to use the text box below."
                }),
                "new_session_name": ("STRING", {
                    "default": "new_session",
                    "tooltip": "If [CREATE NEW] is selected above, this will be your new session_id."
                }),
            },
        }

    RETURN_TYPES = ("STRING", "INT", "STRING")
    RETURN_NAMES = ("session_memory", "run_count", "session_id")
    FUNCTION = "read"
    CATEGORY = "\U0001f9e0 Memory"

    @classmethod
    def IS_CHANGED(cls, **kwargs):
        return float("NaN")

    def read(self, existing_session="[CREATE NEW]", new_session_name="new_session"):
        if existing_session == "[CREATE NEW]":
            session_id = new_session_name.strip()
        else:
            session_id = existing_session

        if not session_id:
            session_id = "unnamed_session"

        session = _load(session_id)
        entries = session.get("entries", [])
        run_count = session.get("run_count", 0)
        feedback = session.get("feedback", {})

        # Reader returns the SUMMARY field (not the prompt).
        # When summary is empty / placeholder, try extracting <summary>
        # tags from the prompt field (Visual DNA Synthesizer format).
        _PLACEHOLDER = "(No summary provided for this run)"
        summaries = []
        for entry in entries:
            if isinstance(entry, dict):
                summary_text = entry.get("summary", "")
                # Detect placeholder / empty summary → extract <summary> tags
                stripped = summary_text.strip().lstrip("- RUN 0123456789\n")
                if not stripped or stripped == _PLACEHOLDER:
                    prompt_text = entry.get("prompt", "")
                    xml_sums = re.findall(
                        r'<summary>\s*(.*?)\s*</summary>',
                        prompt_text, re.DOTALL,
                    )
                    if xml_sums:
                        summary_text = "\n".join(s.strip() for s in xml_sums)
                summaries.append(summary_text)
            else:
                summaries.append(entry)  # legacy string fallback
        session_memory = "\n".join(summaries) if summaries else ""

        # Append feedback block if any
        feedback_block = _format_feedback_block(feedback)
        if feedback_block:
            session_memory = session_memory + "\n" + feedback_block

        logger.info(
            f"[Session Memory Reader] Session '{session_id}' — "
            f"{run_count} runs, {len(session_memory)} chars"
            f"{', with feedback' if feedback_block else ''}."
        )
        return (session_memory, run_count, session_id)


# ── Node: Writer ───────────────────────────────────────────────────────
class SessionMemoryWriterNode:
    """
    💾 Session Memory Writer

    Placed at the end of the workflow, after the Gemini Summarizer.
    Appends the structured run summary to the session file on disk.
    Outputs session_id so it can chain to the Feedback Editor.
    """

    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "session_id": ("STRING", {
                    "default": "default",
                    "tooltip": "Session identifier. Must match Reader/Feedback.",
                }),
                "prompt": ("STRING", {
                    "multiline": True,
                    "default": "",
                    "tooltip": "The raw prompt sent to the model. Displayed in the Feedback Editor.",
                }),
                "summary": ("STRING", {
                    "multiline": True,
                    "default": "",
                    "tooltip": "Structured summary from the Gemini Summarizer. Returned by the Reader.",
                }),
            },
            "optional": {
                "max_runs": ("INT", {
                    "default": 0,
                    "min": 0,
                    "max": 100,
                    "step": 1,
                    "tooltip": "Sliding window: keep only the last N entries. "
                               "0 = unlimited.",
                }),
            },
        }

    RETURN_TYPES = ("STRING", "STRING")
    RETURN_NAMES = ("run_summary", "session_id")
    OUTPUT_NODE = True
    FUNCTION = "write"
    CATEGORY = "\U0001f9e0 Memory"

    @classmethod
    def IS_CHANGED(cls, **kwargs):
        return float("NaN")

    def write(
        self,
        session_id: str = "default",
        prompt: str = "",
        summary: str = "",
        max_runs: int = 0,
    ):
        prompt_stripped = prompt.strip()
        summary_stripped = summary.strip()

        if not prompt_stripped:
            prompt_stripped = "(No prompt provided for this run)"
        if not summary_stripped:
            summary_stripped = "(No summary provided for this run)"
            logger.info(
                f"[Session Memory Writer] Empty summary — creating a new run with placeholder in session '{session_id}'."
            )

        session = _load(session_id)
        run_count = session.get("run_count", 0) + 1
        entries: list = session.get("entries", [])
        feedback = session.get("feedback", {})

        entry = {
            "prompt": f"--- RUN {run_count} ---\n{prompt_stripped}\n",
            "summary": f"--- RUN {run_count} ---\n{summary_stripped}\n",
        }
        entries.append(entry)

        # Sliding window
        if max_runs > 0 and len(entries) > max_runs:
            trimmed = len(entries) - max_runs
            entries = entries[-max_runs:]
            logger.info(
                f"[Session Memory Writer] Trimmed {trimmed} oldest entries "
                f"(window={max_runs})."
            )

        _save(session_id, {
            "run_count": run_count,
            "entries": entries,
            "feedback": feedback,
        })
        logger.info(
            f"[Session Memory Writer] Written Run {run_count} to session "
            f"'{session_id}' (prompt={len(prompt_stripped)} chars, "
            f"summary={len(summary_stripped)} chars, {len(entries)} entries)."
        )
        return (summary, session_id)


def _get_available_sessions():
    os.makedirs(SESSIONS_DIR, exist_ok=True)
    sessions = []
    for f in os.listdir(SESSIONS_DIR):
        if f.endswith(".json"):
            sessions.append(f[:-5])
    return sorted(sessions) if sessions else ["default"]


class SessionIDSelectorNode:
    """
    Select an existing session id from a list, or type a new one to start fresh.
    """
    @classmethod
    def INPUT_TYPES(cls):
        sessions = _get_available_sessions()
        return {
            "required": {
                "existing_session": (["[CREATE NEW]"] + sessions, {
                    "default": sessions[-1] if sessions else "[CREATE NEW]",
                    "tooltip": "Select an existing session, or [CREATE NEW] to use the text box below."
                }),
                "new_session_name": ("STRING", {
                    "default": "new_session",
                    "tooltip": "If [CREATE NEW] is selected above, this will be your new session_id."
                }),
            }
        }

    RETURN_TYPES = ("STRING",)
    RETURN_NAMES = ("session_id",)
    FUNCTION = "select_id"
    CATEGORY = "🧠 Memory"

    def select_id(self, existing_session, new_session_name):
        if existing_session == "[CREATE NEW]":
            final_id = new_session_name.strip()
        else:
            final_id = existing_session
        return (final_id,)


# ── Registration ───────────────────────────────────────────────────────
NODE_CLASS_MAPPINGS = {
    "SessionIDSelector": SessionIDSelectorNode,
    "SessionFeedbackEditor": SessionFeedbackEditorNode,
    "SessionMemoryReader": SessionMemoryReaderNode,
    "SessionMemoryWriter": SessionMemoryWriterNode,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "SessionIDSelector": "🗑️ (DEPRECATED) Session ID Selector",
    "SessionFeedbackEditor": "\u2b50 Session Feedback Editor",
    "SessionMemoryReader": "\U0001f9e0 Session Memory Reader",
    "SessionMemoryWriter": "\U0001f4be Session Memory Writer",
}