app.py

from __future__ import annotations

import html
import json
import os
import re
import time
import uuid

try:
    import spaces
except ImportError:
    class _Spaces:
        @staticmethod
        def GPU(function=None, *args, **kwargs):
            if function is not None and callable(function):
                return function
            def decorator(fn):
                return fn
            return decorator
    spaces = _Spaces()

import gradio as gr

from llm_service import GeminiService
from memory_engine import (
    active_memories,
    apply_memories,
    build_scene,
    context_overlap,
    detect_generation_overrides,
    detect_memory_reversals,
    extract_edit_candidates,
    find_memory,
    incorporate_llm_proposals,
    learn_from_edit,
    memory_context,
    memory_context_score,
    memory_summary,
    normalize_memory,
    trim_history,
)
from semantic_memory import backend_status, ml_diagnostics, retrieve_memories
from storage import MemoryStore
from evaluation_runner import EvaluationRunner, load_csv, write_reports
from routing_policy import candidate_ids as route_candidate_ids, local_memory_decision, should_escalate_to_llm


INITIAL_NOTES = [
    {
        "title": "Welcome",
        "content": (
            "Welcome to Lexi.\n\n"
            "Your personal writing space."
        ),
    },
    {"title": "Meeting Notes", "content": ""},
    {"title": "Ideas", "content": ""},
]

store = MemoryStore()
gemini = GeminiService()


@spaces.GPU
def zero_gpu_probe():
    return None


FILLER_RE = re.compile(
    r"\b(?:um+|uh+|er+|ah+|hmm+)\b[\s,]*",
    re.IGNORECASE,
)
REPEATED_WORD_RE = re.compile(
    r"\b([A-Za-z]+)(?:\s+\1\b)+",
    re.IGNORECASE,
)


def _decode_browser_state(value, default):
    if value is None:
        return default
    if isinstance(value, str):
        text = value.strip()
        if not text:
            return default
        try:
            return json.loads(text)
        except Exception:
            return value
    return value


def normalize_notes(notes):
    notes = _decode_browser_state(notes, None)
    if not isinstance(notes, list):
        return [dict(note) for note in INITIAL_NOTES]
    cleaned = []
    seen_titles = set()
    for note in notes:
        if not isinstance(note, dict):
            continue
        title = str(note.get("title", "")).strip()
        content = note.get("content", "")
        if not title or title in seen_titles:
            continue
        if not isinstance(content, str):
            content = str(content)
        cleaned.append({"title": title, "content": content})
        seen_titles.add(title)
    return cleaned or [dict(note) for note in INITIAL_NOTES]


def normalize_meta(meta):
    meta = _decode_browser_state(meta, None)
    if not isinstance(meta, dict):
        meta = {}
    meta.setdefault("current_note", "Welcome")
    meta.setdefault("baselines", {})
    meta.setdefault("last_edit", None)
    meta.setdefault("edit_history", [])
    meta.setdefault("generation_history", [])
    meta.setdefault("last_trace", {})
    meta.setdefault("last_model_metrics", {})
    meta.setdefault("last_generation", {})
    return meta


def note_titles(notes):
    return [note["title"] for note in notes]


def get_note(notes, title):
    for note in notes:
        if note["title"] == title:
            return note
    return None


def radio_update(notes, value=None):
    choices = note_titles(notes)
    if value is None and choices:
        value = choices[0]
    return gr.Radio(
        choices=choices,
        value=value,
        label="Notes",
        show_label=True,
    )


def activity_text(message):
    return f"> {message}"


def memory_count_text(memories):
    summary = memory_summary(memories)
    return (
        f"**{summary['active_count']} active** · "
        f"{summary['candidate_count']} candidates · "
        f"{summary['ambiguous_count']} ambiguous"
    )


def llm_status_markdown():
    diagnostics = ml_diagnostics()
    return (
        f"{gemini.status()}  \n"
        f"ML: {backend_status()}  \n"
        f"ML model loaded: {diagnostics['minilm_loaded']} · fallback: {diagnostics['fallback_available']}  \n"
        "Routing: local first → ML for unresolved/contextual cases → Gemini only for genuine ambiguity"
    )



def run_evaluation_ui(file_path):
    if not file_path:
        return (
            "Upload a CSV evaluation file first.",
            [],
            {"error": "No CSV uploaded."},
            [],
        )

    try:
        rows = load_csv(file_path)
    except Exception as exc:
        return (
            f"CSV error: {exc}",
            [],
            {"error": str(exc)},
            [],
        )

    try:
        report = EvaluationRunner(gemini=gemini).run_rows(
            rows,
            policy="auto",
        )
    except Exception as exc:
        return (
            f"Evaluation error: {exc}",
            [],
            {"error": str(exc)},
            [],
        )

    json_path, csv_path = write_reports(
        report,
        os.path.join("/tmp", "lexi_evaluation"),
    )

    summary = dict(report["summary"])
    summary["policy"] = "auto"
    display_columns = [
        "chain_id",
        "step",
        "action",
        "status",
        "passed",
        "semantic_event",
        "route",
        "ml_used",
        "llm_used",
        "actual_output",
        "failures",
    ]
    table = [
        [item.get(column, "") for column in display_columns]
        for item in report["results"]
    ]

    label = (
        f"Completed {summary['total_steps']} steps: "
        f"{summary['passed_steps']} passed, "
        f"{summary['failed_steps']} failed. "
        f"Auto routes: {json.dumps(summary.get('route_counts', {}))}. "
        f"Runtime {summary['total_runtime_ms']:.1f} ms."
    )
    return label, table, summary, [json_path, csv_path]


def _status_badge(status: str) -> str:
    labels = {
        "active": "ACTIVE",
        "candidate": "CANDIDATE",
        "ambiguous": "AMBIGUOUS",
        "rejected": "REJECTED",
    }
    return labels.get(status, status.upper())


def _memory_label(memory: dict) -> str:
    status = str(memory.get("status", "candidate")).upper()
    source = str(memory.get("source", ""))
    target = str(memory.get("target", ""))
    kind = str(memory.get("kind", "other")).replace("_", " ")
    confidence = float(memory.get("confidence", 0.0) or 0.0)
    evidence = int(memory.get("evidence_count", 0) or 0)
    return (
        f"{status} · {source} → {target} · {kind} · "
        f"{confidence:.0%} · {evidence} evidence"
    )


def memory_picker_update(memories, value=None):
    memories = [normalize_memory(memory) for memory in (memories or [])]
    memories = [m for m in memories if m.get("status") != "rejected"]
    memories.sort(
        key=lambda m: (
            {"active": 0, "candidate": 1, "ambiguous": 2}.get(m.get("status"), 3),
            -float(m.get("confidence", 0.0) or 0.0),
            str(m.get("source", "")).lower(),
        )
    )
    choices = [
        (_memory_label(memory), memory["id"])
        for memory in memories
    ]
    if value not in {memory["id"] for memory in memories}:
        value = None
    return gr.Dropdown(
        choices=choices,
        value=value,
        label="Select memory to forget",
        info="Deleting a memory removes its stored evidence too.",
        interactive=bool(choices),
    )


def _context_lines(items, empty_label):
    items = items or []
    if not items:
        return f'<div class="lexi-muted">{html.escape(empty_label)}</div>'
    parts = []
    for item in items[:4]:
        context = html.escape(str(item.get("context", "")))
        note = html.escape(str(item.get("note_title", "")))
        tags = ", ".join(html.escape(str(tag)) for tag in item.get("tags", [])[:4])
        meta = " · ".join(part for part in [note, tags] if part)
        parts.append(
            '<div class="lexi-evidence">'
            f'<div>{context}</div>'
            + (f'<div class="lexi-muted">{meta}</div>' if meta else "")
            + '</div>'
        )
    return "".join(parts)


def render_memory_panel(memories):
    summary = memory_summary(memories or [])
    groups = [
        ("active", "Active memories", summary["active_count"]),
        ("candidate", "Candidates", summary["candidate_count"]),
        ("ambiguous", "Ambiguous", summary["ambiguous_count"]),
    ]
    sections = []
    for status, title, count in groups:
        items = [
            memory
            for memory in summary["memories"]
            if memory.get("status") == status
        ]
        if not items:
            continue
        cards = []
        for memory in items:
            confidence = float(memory.get("confidence", 0.0) or 0.0)
            positive = int(memory.get("positive_evidence_count", 0) or 0)
            negative = int(memory.get("negative_evidence_count", 0) or 0)
            explicit = int(memory.get("explicit_evidence_count", 0) or 0)
            kind = html.escape(str(memory.get("kind", "other")).replace("_", " "))
            source = html.escape(str(memory.get("source", "")))
            target = html.escape(str(memory.get("target", "")))
            reason = html.escape(str(memory.get("last_reason", "")))
            aliases = ", ".join(html.escape(str(alias)) for alias in memory.get("aliases", [])[:6])
            exceptions = memory.get("exceptions", [])
            exception_count = len(exceptions)
            context_summary = (
                f'{len(memory.get("positive_contexts", []))} positive · '
                f'{len(memory.get("negative_contexts", []))} negative'
            )
            detail_parts = []
            if aliases:
                detail_parts.append(
                    f'<div class="lexi-detail"><strong>Aliases</strong><br>{aliases}</div>'
                )
            if reason:
                detail_parts.append(
                    f'<div class="lexi-detail"><strong>Why</strong><br>{reason}</div>'
                )
            if exception_count:
                detail_parts.append(
                    f'<div class="lexi-detail"><strong>Exceptions</strong> · {exception_count}</div>'
                )
            detail_parts.append(
                '<div class="lexi-detail"><strong>Positive contexts</strong></div>'
                + _context_lines(memory.get("positive_contexts"), "No positive context recorded.")
                + '<div class="lexi-detail"><strong>Negative contexts</strong></div>'
                + _context_lines(memory.get("negative_contexts"), "No negative context recorded.")
            )
            cards.append(
                '<div class="lexi-memory-card">'
                '<div class="lexi-memory-top">'
                f'<span class="lexi-badge lexi-{status}">{_status_badge(status)}</span>'
                f'<span class="lexi-kind">{kind}</span>'
                '</div>'
                f'<div class="lexi-memory-map"><span>{source}</span> <span class="lexi-arrow">→</span> <strong>{target}</strong></div>'
                f'<div class="lexi-memory-meta">Confidence {confidence:.0%} · +{positive} · -{negative} · {explicit} explicit · {context_summary}</div>'
                '<details class="lexi-details">'
                '<summary>Evidence, context & exceptions</summary>'
                + ''.join(detail_parts)
                + '</details>'
                '</div>'
            )
        sections.append(
            f'<div class="lexi-memory-section"><div class="lexi-section-title">{title} <span>{count}</span></div>'
            + "".join(cards)
            + '</div>'
        )
    if not sections:
        body = '<div class="lexi-empty">Lexi has not learned any personal vocabulary yet.</div>'
    else:
        body = "".join(sections)
    return (
        '<div class="lexi-memory-scroll">'
        f'{body}'
        '</div>'
    )


def _trace_value(value, empty="None"):
    if value is None or value == "":
        return html.escape(empty)
    if isinstance(value, bool):
        return "yes" if value else "no"
    return html.escape(str(value))


def render_trace(trace):
    trace = trace or {}
    if not trace:
        return '<div class="lexi-empty">No decision has been recorded yet.</div>'

    decision = _trace_value(trace.get("decision"), "unknown")
    stage = _trace_value(trace.get("stage"), "unknown")
    reason = _trace_value(trace.get("reason"), "No reason recorded.")
    mode = _trace_value(trace.get("mode"), "")
    input_text = _trace_value(trace.get("input") or trace.get("after") or "", "")
    timing = trace.get("timing", {}) or {}
    retrieved = trace.get("retrieved_memories", []) or []
    applied = trace.get("applied_changes", []) or []
    scene = trace.get("scene", {}) or {}
    model_metrics = trace.get("model_metrics", {}) or {}

    retrieved_html = []
    for item in retrieved[:6]:
        source = html.escape(str(item.get("source", "")))
        target = html.escape(str(item.get("target", "")))
        score = float(item.get("retrieval_score", 0.0) or 0.0)
        ctx = float(item.get("context_score", 0.0) or 0.0)
        present = bool(item.get("term_present"))
        blocked = bool(item.get("blocked_by_context"))
        flags = []
        if present:
            flags.append("present")
        if blocked:
            flags.append("blocked")
        flag_text = " · ".join(flags)
        retrieved_html.append(
            '<div class="lexi-trace-row">'
            f'<strong>{source} → {target}</strong> · retrieval {score:.0%} · context {ctx:.0%}'
            + (f' · {html.escape(flag_text)}' if flag_text else "")
            + '</div>'
        )
    if not retrieved_html:
        retrieved_html.append('<div class="lexi-muted">No memory candidates were retrieved.</div>')

    applied_html = []
    for change in applied[:8]:
        original = html.escape(str(change.get("original", "")))
        replacement = html.escape(str(change.get("replacement", "")))
        confidence = float(change.get("confidence", 0.0) or 0.0)
        applied_html.append(
            f'<div class="lexi-trace-row"><strong>{original} → {replacement}</strong> · {confidence:.0%}</div>'
        )
    if not applied_html:
        applied_html.append('<div class="lexi-muted">No memory change was applied.</div>')

    tags = scene.get("tags", []) if isinstance(scene, dict) else []
    tag_text = ", ".join(html.escape(str(tag)) for tag in tags) or "none"
    model_text = html.escape(str(model_metrics.get("model") or "local"))
    api_text = html.escape(str(model_metrics.get("api") or "none"))

    return (
        '<div class="lexi-trace-scroll">'
        '<div class="lexi-trace-summary">'
        f'<div><span class="lexi-label">Decision</span><strong>{decision}</strong></div>'
        f'<div><span class="lexi-label">Stage</span>{stage}</div>'
        f'<div><span class="lexi-label">Mode</span>{mode}</div>'
        f'<div><span class="lexi-label">Scene tags</span>{tag_text}</div>'
        '</div>'
        f'<div class="lexi-trace-block"><div class="lexi-section-title">Input</div><div class="lexi-code">{input_text}</div></div>'
        f'<div class="lexi-trace-block"><div class="lexi-section-title">Why</div><div>{reason}</div></div>'
        '<div class="lexi-trace-block"><div class="lexi-section-title">Retrieved memories</div>'
        + "".join(retrieved_html)
        + '</div>'
        '<div class="lexi-trace-block"><div class="lexi-section-title">Applied</div>'
        + "".join(applied_html)
        + '</div>'
        '<div class="lexi-trace-block"><div class="lexi-section-title">Timing</div>'
        f'<div class="lexi-trace-row">Retrieval {_trace_value(timing.get("retrieval_ms"), "0")} ms · LLM {_trace_value(timing.get("llm_ms"), "0")} ms · total {_trace_value(timing.get("total_ms"), "0")} ms</div>'
        f'<div class="lexi-trace-row">Backend {_trace_value(trace.get("retrieval_backend"))} · model {model_text} · API {api_text}</div>'
        '</div>'
        '</div>'
    )


def ensure_workspace(notes, current_note, memories, meta):
    notes = normalize_notes(notes)
    current_note = _decode_browser_state(current_note, "Welcome")
    memories = _decode_browser_state(memories, [])
    if not isinstance(memories, list):
        memories = []
    memories = [normalize_memory(memory) for memory in memories]
    meta = normalize_meta(meta)
    if not get_note(notes, current_note):
        current_note = notes[0]["title"]
    meta["current_note"] = current_note
    for note in notes:
        meta["baselines"].setdefault(note["title"], note["content"])
    return notes, current_note, memories, meta


def load_app(notes, current_note, memories, meta, user_id):
    user_id = _decode_browser_state(user_id, "")
    if not user_id:
        user_id = str(uuid.uuid4())
    notes, current_note, memories, meta = ensure_workspace(
        notes, current_note, memories, meta
    )
    persisted = store.load_memories(user_id)
    memories = persisted
    note = get_note(notes, current_note)
    trace = store.latest_decisions(user_id, limit=1)
    if trace:
        meta["last_trace"] = trace[0]["trace"]
    summary = memory_summary(memories)
    return (
        notes,
        current_note,
        memories,
        meta,
        user_id,
        radio_update(notes, current_note),
        note["content"],
        f"### {current_note}",
        activity_text(
            "Lexi learns lexical preferences from your edits and uses context "
            "before applying them."
        ),
        render_memory_panel(memories),
        memory_count_text(memories),
        memory_picker_update(memories),
        llm_status_markdown(),
        render_trace(meta.get("last_trace", {})),
    )


def create_note(notes, current_note, memories, meta):
    notes, current_note, memories, meta = ensure_workspace(
        notes, current_note, memories, meta
    )
    existing_titles = {note["title"] for note in notes}
    index = 1
    while f"Untitled {index}" in existing_titles:
        index += 1
    title = f"Untitled {index}"
    notes.insert(0, {"title": title, "content": ""})
    meta["baselines"][title] = ""
    meta["current_note"] = title
    meta["last_edit"] = None
    return (
        notes,
        title,
        memories,
        meta,
        radio_update(notes, title),
        "",
        f"### {title}",
        activity_text("New note created."),
        render_memory_panel(memories),
        memory_count_text(memories),
        memory_picker_update(memories),
        "",
    )


def select_note(title, notes):
    notes = normalize_notes(notes)
    note = get_note(notes, title)
    if note is None:
        return "", "### Lexi", title or ""
    return note["content"], f"### {title}", title


def fast_format_asr(text):
    value = re.sub(r"\s+", " ", (text or "").strip())
    value = FILLER_RE.sub(" ", value)
    value = REPEATED_WORD_RE.sub(r"\1", value)
    value = re.sub(r"\s+([,.!?;:])", r"\1", value)
    value = re.sub(r"([,.!?;:])(?=\w)", r"\1 ", value)
    value = re.sub(r"\s{2,}", " ", value).strip()
    if not value:
        return ""
    value = value[0].upper() + value[1:]
    if value[-1] not in ".!?":
        value += "."
    return value


def _append_text(note_content, text):
    if note_content.strip():
        return note_content + "\n\n" + text
    return text


def _retrieve(text, note_title, note_content, memories):
    scene_tail = (note_content or "")[-600:]
    started = time.perf_counter()
    results = retrieve_memories(
        text,
        memories,
        note_title=note_title,
        scene_text=scene_tail,
        top_k=4,
    )
    return results, round((time.perf_counter() - started) * 1000, 2)


def _candidate_ids(retrieved):
    return route_candidate_ids(retrieved)


def _apply_selected(text, memories, selected_ids, note_title, note_content):
    return apply_memories(
        text,
        memories,
        selected_ids=selected_ids,
        note_title=note_title,
        scene_text=(note_content or "")[-600:],
    )


def _finalize_generation(
    note,
    current_note,
    input_text,
    mode,
    rewritten,
    changes,
    retrieved,
    metrics,
    retrieval_ms,
    backend_used,
    meta,
    user_id,
    decision_type,
    reason,
    scene,
):
    note["content"] = _append_text(note["content"], rewritten)
    meta["baselines"][current_note] = note["content"]
    meta["last_edit"] = None
    meta["current_note"] = current_note

    total_ms = round(
        (time.perf_counter() - scene["started_at"]) * 1000,
        2,
    )
    metrics = dict(metrics or {})
    metrics["total_latency_ms"] = total_ms

    trace = {
        "stage": "generation",
        "decision": decision_type,
        "input": input_text,
        "mode": mode,
        "scene": scene.get("scene", {}),
        "retrieved_memories": retrieved,
        "applied_changes": changes,
        "reason": reason,
        "timing": {
            "retrieval_ms": retrieval_ms,
            "llm_ms": metrics.get("latency_ms"),
            "total_ms": total_ms,
        },
        "model_metrics": metrics,
        "retrieval_backend": backend_used,
    }
    meta["last_trace"] = trace
    meta["last_model_metrics"] = metrics
    meta.setdefault("generation_history", []).append(
        {
            "note": current_note,
            "mode": mode,
            "input": input_text,
            "output": rewritten,
            "retrieved": retrieved,
            "changes": changes,
            "metrics": metrics,
            "decision": decision_type,
        }
    )
    meta["generation_history"] = trim_history(meta["generation_history"])

    store.log_decision(
        user_id,
        current_note,
        input_text,
        mode,
        decision_type,
        reason,
        trace,
        metrics,
    )

    meta["last_generation"] = {
        "note_title": current_note,
        "output": rewritten,
        "changes": changes,
        "created_at": time.time(),
    }

    if changes:
        message = "Lexi applied " + ", ".join(
            f'{change["original"]} → {change["replacement"]}'
            for change in changes
        ) + "."
    elif decision_type == "fast_path":
        message = "Lexi used the local path and found no context-supported substitution to apply."
    elif decision_type == "context_skip":
        message = "Lexi kept the wording because the memory was not reliable enough in this scene."
    else:
        message = "Lexi formatted the input without applying a personal lexical memory."

    return note["content"], "", meta, activity_text(message), trace


def process_input(text, mode, current_note, notes, memories, meta, user_id):
    notes, current_note, memories, meta = ensure_workspace(
        notes, current_note, memories, meta
    )
    note = get_note(notes, current_note)
    text = (text or "").strip()
    if not text:
        return (
            note["content"], "", notes, meta,
            activity_text("Nothing was added."),
            render_memory_panel(memories),
            memory_count_text(memories), memory_picker_update(memories),
            llm_status_markdown(), render_trace(meta.get("last_trace", {})),
        )

    started_at = time.perf_counter()
    scene = {
        "started_at": started_at,
        "scene": build_scene(current_note, note["content"], text),
    }
    retrieved, retrieval_ms = _retrieve(
        text,
        current_note,
        note["content"],
        memories,
    )
    backend_used = retrieved[0]["retrieval_backend"] if retrieved else "local/no-memory"
    selected_ids = _candidate_ids(retrieved)
    decision, local_ids, decision_reason, decision_details = local_memory_decision(retrieved)
    llm_metrics = {"model": None, "api": None, "latency_ms": 0.0}
    llm_used = False
    backend_lower = backend_used.lower()
    ml_used = (
        "minilm" in backend_lower
        or "tf-idf" in backend_lower
        or backend_lower.startswith("ml ")
        or backend_lower.startswith("ml:")
    )

    if mode == "Formatted":
        selected_ids = local_ids if decision == "apply" else set()
        decision_type = "fast_path" if selected_ids else "context_skip"
        reason = decision_reason
    else:
        if decision == "apply":
            selected_ids = local_ids
            decision_type = ("ml_path" if ml_used else "fast_path") if selected_ids else "context_skip"
            reason = decision_reason
        elif decision == "ambiguous" and ml_used and should_escalate_to_llm(
            retrieved,
            decision,
            gemini.available,
        ):
            try:
                decision_result, llm_metrics = gemini.decide_memory_application(
                    text,
                    scene["scene"]["text"],
                    current_note,
                    memory_context(memories, selected_ids, limit=4),
                )
                llm_used = True
                eligible_ids = selected_ids
                proposed_ids = {
                    memory_id
                    for memory_id in decision_result.memory_ids
                    if memory_id in eligible_ids
                }
                if decision_result.decision == "apply" and proposed_ids:
                    selected_ids = proposed_ids
                    decision_type = "context_llm"
                    reason = (
                        "Local and ML signals were ambiguous; Gemini resolved the ambiguity "
                        f"without introducing new memory IDs: {decision_result.rationale}"
                    )
                else:
                    selected_ids = set()
                    decision_type = "context_skip"
                    reason = (
                        "Local and ML signals were ambiguous and Gemini did not find a sufficiently "
                        f"safe memory to apply: {decision_result.rationale}"
                    )
            except Exception as exc:
                llm_metrics = {
                    "model": gemini.model,
                    "api": "interactions_v2",
                    "latency_ms": 0.0,
                    "error": str(exc),
                }
                selected_ids = set()
                decision_type = "context_skip"
                reason = "Context remained ambiguous; Lexi deliberately avoided intervention."
        else:
            selected_ids = set()
            decision_type = "context_skip"
            reason = decision_reason

    rewritten = fast_format_asr(text)
    rewritten, changes = _apply_selected(
        rewritten,
        memories,
        selected_ids,
        current_note,
        note["content"],
    )

    content, input_value, meta, activity, trace = _finalize_generation(
        note, current_note, text, mode, rewritten, changes, retrieved,
        llm_metrics, retrieval_ms, backend_used, meta, user_id,
        decision_type, reason, scene,
    )
    trace["routing"] = {
        "local_decision": decision,
        "local_details": decision_details,
        "ml_used": ml_used,
        "llm_escalated": llm_used,
    }
    if ml_used and not llm_used:
        trace["route"] = "ml"
    elif llm_used:
        trace["route"] = "ml+llm" if ml_used else "llm"
    else:
        trace["route"] = "local"
    meta["last_trace"] = trace
    return (
        content, input_value, notes, meta, activity,
        render_memory_panel(memories), memory_count_text(memories),
        memory_picker_update(memories), llm_status_markdown(), render_trace(trace),
    )


def handle_manual_edit(text, current_note, notes, memories, meta, user_id):
    notes, current_note, memories, meta = ensure_workspace(
        notes, current_note, memories, meta
    )
    note = get_note(notes, current_note)
    if note is None:
        return (
            notes, memories, meta, activity_text("No note selected."),
            render_memory_panel(memories), memory_count_text(memories), memory_picker_update(memories), render_trace(meta.get("last_trace", {}))
        )

    text = text or ""
    baseline = meta["baselines"].get(current_note)
    if baseline is None:
        note["content"] = text
        meta["baselines"][current_note] = text
        return (
            notes, memories, meta,
            activity_text("Lexi is now tracking this note."),
            render_memory_panel(memories), memory_count_text(memories), memory_picker_update(memories),
            render_trace(meta.get("last_trace", {}))
        )

    if text == baseline:
        return (
            notes, memories, meta, activity_text("No manual change detected."),
            render_memory_panel(memories), memory_count_text(memories), memory_picker_update(memories),
            render_trace(meta.get("last_trace", {}))
        )

    previous_edit = meta.get("last_edit")
    last_generation = meta.get("last_generation") or {}

    override_memories, override_event = detect_generation_overrides(
        baseline,
        text,
        memories,
        last_generation,
        current_note,
    )

    negative_memories, negative_event = detect_memory_reversals(
        baseline,
        text,
        memories,
        current_note,
    )

    llm_metrics = {}
    use_llm_learning = False
    if override_event["type"] == "memory_override":
        updated_memories = override_memories
        event = override_event
    elif negative_event["type"] == "memory_negative_evidence":
        updated_memories = negative_memories
        event = negative_event
    else:
        updated_memories, event = learn_from_edit(
            baseline,
            text,
            memories,
            previous_edit=previous_edit,
            last_generation={},
            note_title=current_note,
        )

    note["content"] = text
    meta["baselines"][current_note] = text
    meta["current_note"] = current_note
    meta["last_edit"] = {"before": baseline, "after": text}
    meta["last_generation"] = {}

    meta.setdefault("edit_history", []).append(
        {
            "note": current_note,
            "before": baseline,
            "after": text,
            "decision": event["type"],
            "message": event["message"],
            "candidates": event.get("candidates", []),
            "llm": event.get("llm_decision", {}),
            "metrics": llm_metrics,
            "backend": "gemini" if use_llm_learning else "deterministic",
        }
    )
    meta["edit_history"] = trim_history(meta["edit_history"])

    store.upsert_memories(user_id, updated_memories)

    candidates_by_id = {}
    for item in event.get("candidates", []):
        memory_id = item.get("memory_id")
        if memory_id:
            candidates_by_id.setdefault(memory_id, []).append(item)

    for memory in updated_memories:
        items = candidates_by_id.get(memory["id"], [])
        if not items:
            continue
        negative = event["type"] in {
            "memory_override",
            "memory_conflict",
            "memory_negative_evidence",
        }
        polarity = "negative" if negative else "positive"
        evidence_type = "override" if negative else "manual_edit"
        for item in items:
            store.save_evidence_for_memory(
                user_id,
                memory,
                baseline,
                text,
                evidence_type,
                polarity=polarity,
                note_title=current_note,
                context=item.get("context", text),
                confidence=memory.get("confidence"),
            )

    trace = {
        "stage": "learning",
        "decision": event["type"],
        "before": baseline,
        "after": text,
        "note_title": current_note,
        "memory_state": memory_summary(updated_memories)["memories"],
        "event": event,
        "model_metrics": llm_metrics,
    }
    meta["last_trace"] = trace
    store.log_decision(
        user_id,
        current_note,
        text,
        "manual_edit",
        event["type"],
        event["message"],
        trace,
        llm_metrics,
    )

    return (
        notes,
        updated_memories,
        meta,
        activity_text(event["message"]),
        render_memory_panel(updated_memories),
        memory_count_text(updated_memories),
        memory_picker_update(updated_memories),
        render_trace(trace),
    )


def rename_note(new_title, current_note, notes, meta):
    current_note = _decode_browser_state(current_note, "Welcome")
    notes = normalize_notes(notes)
    meta = normalize_meta(meta)
    new_title = (new_title or "").strip()
    if not current_note:
        return notes, current_note, meta, radio_update(notes, current_note), "### Lexi", "", activity_text("No note selected.")
    if not new_title:
        return notes, current_note, meta, radio_update(notes, current_note), f"### {current_note}", "", activity_text("Enter a note name.")
    if new_title == current_note:
        return notes, current_note, meta, radio_update(notes, current_note), f"### {current_note}", "", activity_text("The note name is unchanged.")
    if get_note(notes, new_title) is not None:
        return notes, current_note, meta, radio_update(notes, current_note), f"### {current_note}", new_title, activity_text("A note with that name already exists.")
    note = get_note(notes, current_note)
    if note is None:
        return notes, current_note, meta, radio_update(notes, current_note), "### Lexi", "", activity_text("Note not found.")
    note["title"] = new_title
    if current_note in meta["baselines"]:
        meta["baselines"][new_title] = meta["baselines"].pop(current_note)
    meta["current_note"] = new_title
    return notes, new_title, meta, radio_update(notes, new_title), f"### {new_title}", "", activity_text("Note renamed.")


def delete_memory(memory_id, user_id, memories, meta):
    user_id = _decode_browser_state(user_id, "")
    memories = _decode_browser_state(memories, [])
    meta = normalize_meta(meta)
    memories = [normalize_memory(memory) for memory in (memories or [])]
    memory_id = str(memory_id or "").strip()
    if not memory_id:
        return memories, meta, activity_text("Select a memory to forget."), render_memory_panel(memories), memory_count_text(memories), memory_picker_update(memories), render_trace(meta.get("last_trace", {}))

    target = next((memory for memory in memories if memory.get("id") == memory_id), None)
    if target is None:
        return memories, meta, activity_text("That memory no longer exists."), render_memory_panel(memories), memory_count_text(memories), memory_picker_update(memories), render_trace(meta.get("last_trace", {}))

    memories = [memory for memory in memories if memory.get("id") != memory_id]
    if user_id:
        store.delete_memory(user_id, memory_id)

    meta["last_generation"] = {}
    meta["last_trace"] = {
        "stage": "memory_management",
        "decision": "memory_deleted",
        "reason": f'User deleted memory {target.get("source", "")} → {target.get("target", "")}.',
        "memory_id": memory_id,
        "memory_mapping": f'{target.get("source", "")} → {target.get("target", "")}',
    }

    message = (
        f'Forgot "{target.get("source", "")}" → "{target.get("target", "")}". '
        "Its stored memory record and evidence were deleted."
    )
    return (
        memories,
        meta,
        activity_text(message),
        render_memory_panel(memories),
        memory_count_text(memories),
        memory_picker_update(memories),
        render_trace(meta["last_trace"]),
    )


def delete_note(current_note, notes, meta):
    current_note = _decode_browser_state(current_note, "Welcome")
    notes = normalize_notes(notes)
    meta = normalize_meta(meta)
    if len(notes) == 1:
        return notes, current_note, meta, radio_update(notes, current_note), notes[0]["content"], f"### {current_note}", "", activity_text("Lexi keeps at least one note.")
    index = next((i for i, note in enumerate(notes) if note["title"] == current_note), None)
    if index is None:
        return notes, current_note, meta, radio_update(notes, current_note), "", "### Lexi", "", activity_text("Note not found.")
    deleted = notes.pop(index)
    meta["baselines"].pop(deleted["title"], None)
    next_index = min(index, len(notes) - 1)
    next_note = notes[next_index]
    current_note = next_note["title"]
    meta["current_note"] = current_note
    meta["last_edit"] = None
    return notes, current_note, meta, radio_update(notes, current_note), next_note["content"], f"### {current_note}", "", activity_text(f'Deleted "{deleted["title"]}".')


def reset_workspace(user_id):
    user_id = _decode_browser_state(user_id, "")
    if user_id:
        store.reset_user(user_id)
    notes = [dict(note) for note in INITIAL_NOTES]
    memories = []
    meta = {
        "current_note": "Welcome",
        "baselines": {note["title"]: note["content"] for note in notes},
        "last_edit": None,
        "edit_history": [],
        "generation_history": [],
        "last_trace": {},
        "last_model_metrics": {},
        "last_generation": {},
    }
    return (
        notes,
        "Welcome",
        memories,
        meta,
        radio_update(notes, "Welcome"),
        notes[0]["content"],
        "### Welcome",
        "",
        activity_text("Lexi has been reset."),
        render_memory_panel([]),
        "**0 active** · **0 candidates** · **0 ambiguous**",
        memory_picker_update([]),
        llm_status_markdown(),
        render_trace({}),
        user_id,
    )


css = """
html, body { margin: 0; padding: 0; }
.gradio-container { max-width: none !important; padding: 0 !important; }
#main-content { min-height: 100vh; }
#note-title { padding: 22px 35px 10px 35px; }
#note-editor textarea {
    height: calc(100vh - 270px) !important;
    border: none !important;
    box-shadow: none !important;
    font-size: 17px !important;
    line-height: 1.75 !important;
    padding: 30px 45px !important;
    resize: none !important;
}
#input-row { padding: 15px 30px 25px 30px; }
#input-box textarea { min-height: 55px !important; max-height: 180px !important; }
#activity, #status { padding: 0 40px; }
#memory-panel, #trace-panel {
    border: 1px solid rgba(255,255,255,0.08);
    border-radius: 10px;
    padding: 10px;
    background: rgba(255,255,255,0.02);
}
.lexi-memory-scroll, .lexi-trace-scroll {
    max-height: 420px;
    overflow-y: auto;
    padding-right: 4px;
}
.lexi-memory-section { margin-bottom: 14px; }
.lexi-section-title {
    font-weight: 700;
    font-size: 13px;
    text-transform: uppercase;
    letter-spacing: .04em;
    margin: 5px 0 8px;
}
.lexi-section-title span { opacity: .6; margin-left: 4px; }
.lexi-memory-card {
    padding: 10px;
    margin: 7px 0;
    border: 1px solid rgba(255,255,255,0.08);
    border-radius: 9px;
    background: rgba(255,255,255,0.03);
}
.lexi-memory-top { display: flex; align-items: center; gap: 7px; margin-bottom: 7px; }
.lexi-badge {
    display: inline-block;
    padding: 2px 6px;
    border-radius: 999px;
    font-size: 10px;
    font-weight: 700;
}
.lexi-active { background: rgba(90, 190, 120, .20); }
.lexi-candidate { background: rgba(220, 180, 60, .18); }
.lexi-ambiguous { background: rgba(210, 90, 90, .18); }
.lexi-kind { opacity: .65; font-size: 11px; }
.lexi-memory-map { font-size: 15px; line-height: 1.35; }
.lexi-arrow { opacity: .55; padding: 0 4px; }
.lexi-memory-meta, .lexi-muted { font-size: 11px; opacity: .62; margin-top: 5px; line-height: 1.45; }
.lexi-detail { font-size: 11px; margin-top: 8px; line-height: 1.45; }
.lexi-details { margin-top: 8px; font-size: 11px; }
.lexi-details summary { cursor: pointer; opacity: .72; user-select: none; }
.lexi-evidence {
    padding: 6px 7px;
    margin-top: 4px;
    border-left: 2px solid rgba(255,255,255,0.13);
    background: rgba(255,255,255,0.02);
}
.lexi-empty { padding: 8px 4px; opacity: .65; font-size: 12px; line-height: 1.5; }
.lexi-trace-summary {
    display: grid;
    grid-template-columns: 1fr 1fr;
    gap: 7px;
    margin-bottom: 10px;
}
.lexi-trace-summary > div {
    padding: 7px;
    border: 1px solid rgba(255,255,255,0.07);
    border-radius: 7px;
    font-size: 11px;
}
.lexi-label { display: block; opacity: .55; margin-bottom: 2px; text-transform: uppercase; font-size: 9px; letter-spacing: .04em; }
.lexi-trace-block { margin: 9px 0; }
.lexi-trace-row { font-size: 11px; line-height: 1.45; padding: 4px 0; }
.lexi-code {
    white-space: pre-wrap;
    word-break: break-word;
    padding: 7px;
    border-radius: 7px;
    background: rgba(0,0,0,.18);
    font-size: 11px;
    line-height: 1.45;
}
#memory-delete-row { margin-top: 10px; }
@media (max-width: 900px) {
    .gradio-container { min-width: 0 !important; }
    #note-editor textarea { padding: 20px !important; }
}
"""

browser_secret = os.getenv("LEXI_BROWSER_SECRET", "lexi-development-secret")

with gr.Blocks(title="Lexi") as demo:
    notes_browser = gr.BrowserState(
        INITIAL_NOTES,
        storage_key="lexi_notes_v7",
        secret=browser_secret,
    )
    current_note_browser = gr.BrowserState(
        "Welcome",
        storage_key="lexi_current_note_v7",
        secret=browser_secret,
    )
    memories_browser = gr.BrowserState(
        [],
        storage_key="lexi_memories_v7",
        secret=browser_secret,
    )
    meta_browser = gr.BrowserState(
        {},
        storage_key="lexi_meta_v7",
        secret=browser_secret,
    )
    user_id_browser = gr.BrowserState(
        "",
        storage_key="lexi_user_id_v7",
        secret=browser_secret,
    )

    with gr.Sidebar(open=True, width=410):
        gr.Markdown("# Lexi")
        gr.Markdown("A context-aware personal word memory")
        new_note = gr.Button("+ New note", variant="secondary")
        note_list = gr.Radio(
            choices=note_titles(INITIAL_NOTES),
            value="Welcome",
            label="Notes",
        )
        rename_box = gr.Textbox(label="Rename note", placeholder="Enter a new name")
        with gr.Row():
            rename_button = gr.Button("Rename", variant="secondary")
            delete_button = gr.Button("Delete", variant="stop")

        gr.Markdown("---")
        gr.Markdown("### Memory")
        memory_count = gr.Markdown("**0 active** · **0 candidates** · **0 ambiguous**")
        with gr.Accordion("What Lexi remembers", open=False):
            memory_panel = gr.HTML("<div class=\"lexi-empty\">Lexi has not learned any personal vocabulary yet.</div>", elem_id="memory-panel")
            memory_picker = gr.Dropdown(
                choices=[],
                label="Select memory to forget",
                info="Deleting a memory removes its stored evidence too.",
                interactive=False,
            )
            with gr.Row(elem_id="memory-delete-row"):
                delete_memory_button = gr.Button("Forget selected memory", variant="stop")
        with gr.Accordion("Decision trace", open=False):
            trace_panel = gr.HTML("<div class=\"lexi-empty\">No decision has been recorded yet.</div>", elem_id="trace-panel")
        status = gr.Markdown("", elem_id="status")
        reset_button = gr.Button("Reset Lexi", variant="stop")

    with gr.Column(elem_id="main-content"):
        with gr.Tabs():
            with gr.Tab("Lexi"):
                note_title = gr.Markdown("### Welcome", elem_id="note-title")
                note_editor = gr.Textbox(
                    value=INITIAL_NOTES[0]["content"],
                    show_label=False,
                    lines=20,
                    elem_id="note-editor",
                    container=False,
                )
                activity = gr.Markdown("> Lexi is ready.", elem_id="activity")
                with gr.Row(elem_id="input-row"):
                    mode = gr.Radio(
                        ["Raw ASR", "Formatted"],
                        value="Raw ASR",
                        show_label=False,
                        scale=2,
                    )
                    input_box = gr.Textbox(
                        placeholder="Paste what the recognizer produced...",
                        show_label=False,
                        lines=2,
                        scale=8,
                        elem_id="input-box",
                    )
                    send_button = gr.Button("→", variant="primary", scale=1)

            with gr.Tab("Evaluation Lab"):
                gr.Markdown(
                    "## Evaluation Lab\n"
                    "Upload a CSV containing reproducible memory journeys. Each `chain_id` keeps its own state, so you can test learn → apply → override → reapply sequences."
                )
                with gr.Row():
                    evaluation_file = gr.File(
                        label="Evaluation CSV",
                        file_types=[".csv"],
                        type="filepath",
                        scale=3,
                    )
                    run_evaluation_button = gr.Button(
                        "Run Auto evaluation",
                        variant="primary",
                        scale=1,
                    )
                evaluation_status = gr.Markdown(
                    "Upload a CSV and run Auto. Lexi chooses the execution path for each step."
                )
                evaluation_table = gr.Dataframe(
                    headers=[
                        "chain_id", "step", "action", "status", "passed",
                        "event", "route", "ml_used", "llm_used",
                        "actual_output", "failures"
                    ],
                    datatype=[
                        "str", "number", "str", "str", "bool",
                        "str", "str", "bool", "bool", "str", "str"
                    ],
                    value=[],
                    interactive=False,
                    wrap=True,
                    max_height=520,
                )
                with gr.Row():
                    evaluation_summary = gr.JSON(
                        label="Run summary",
                        value={},
                    )
                    evaluation_reports = gr.Files(
                        label="Download reports",
                        file_count="multiple",
                    )
                gr.Markdown(
                    "### CSV columns\n"
                    "Required: `chain_id`, `step`, `action`, `mode`. A generation uses `input`; an edit uses `before` and `after`. The CSV describes only the user journey and expected behavior. It does not contain `use_llm` or `use_ml`; Lexi decides routing automatically. Semantic events include `learn`, `ignore`, `negative_evidence`, `apply`, and `do_not_apply`. Use semicolons for multiple output assertions."
                )

    demo.load(
        load_app,
        inputs=[notes_browser, current_note_browser, memories_browser, meta_browser, user_id_browser],
        outputs=[notes_browser, current_note_browser, memories_browser, meta_browser, user_id_browser, note_list, note_editor, note_title, activity, memory_panel, memory_count, memory_picker, status, trace_panel],
    )

    note_list.change(
        select_note,
        inputs=[note_list, notes_browser],
        outputs=[note_editor, note_title, current_note_browser],
    )

    note_editor.blur(
        handle_manual_edit,
        inputs=[note_editor, current_note_browser, notes_browser, memories_browser, meta_browser, user_id_browser],
        outputs=[notes_browser, memories_browser, meta_browser, activity, memory_panel, memory_count, memory_picker, trace_panel],
    )

    send_button.click(
        process_input,
        inputs=[input_box, mode, current_note_browser, notes_browser, memories_browser, meta_browser, user_id_browser],
        outputs=[note_editor, input_box, notes_browser, meta_browser, activity, memory_panel, memory_count, memory_picker, status, trace_panel],
    )

    input_box.submit(
        process_input,
        inputs=[input_box, mode, current_note_browser, notes_browser, memories_browser, meta_browser, user_id_browser],
        outputs=[note_editor, input_box, notes_browser, meta_browser, activity, memory_panel, memory_count, memory_picker, status, trace_panel],
    )

    new_note.click(
        create_note,
        inputs=[notes_browser, current_note_browser, memories_browser, meta_browser],
        outputs=[notes_browser, current_note_browser, memories_browser, meta_browser, note_list, note_editor, note_title, activity, memory_panel, memory_count, memory_picker, rename_box],
    )

    rename_button.click(
        rename_note,
        inputs=[rename_box, current_note_browser, notes_browser, meta_browser],
        outputs=[notes_browser, current_note_browser, meta_browser, note_list, note_title, rename_box, activity],
    )

    delete_button.click(
        delete_note,
        inputs=[current_note_browser, notes_browser, meta_browser],
        outputs=[notes_browser, current_note_browser, meta_browser, note_list, note_editor, note_title, rename_box, activity],
    )

    delete_memory_button.click(
        delete_memory,
        inputs=[memory_picker, user_id_browser, memories_browser, meta_browser],
        outputs=[memories_browser, meta_browser, activity, memory_panel, memory_count, memory_picker, trace_panel],
    )

    reset_button.click(
        reset_workspace,
        inputs=[user_id_browser],
        outputs=[notes_browser, current_note_browser, memories_browser, meta_browser, note_list, note_editor, note_title, rename_box, activity, memory_panel, memory_count, memory_picker, status, trace_panel, user_id_browser],
    )

    run_evaluation_button.click(
        run_evaluation_ui,
        inputs=[evaluation_file],
        outputs=[evaluation_status, evaluation_table, evaluation_summary, evaluation_reports],
    )


if __name__ == "__main__":
    demo.launch(
        css=css,
        theme=gr.themes.Base(primary_hue="neutral", secondary_hue="neutral"),
    )