import os
import sys
import code
import litert_lm
import numpy as np

try:
    import readline
except ImportError:
    readline = None

MODEL_PATH = os.path.expanduser("~/.litert-lm/models/gemma-4-E2B-it.litertlm/model.litertlm")

SYSTEM_CONSTRAINT = (
    "Act as an embedded Python co-interpreter. "
    "Answer the user query in exactly ONE brief sentence or direct assignment. "
    "No greetings, no preambles, no conversational fluff."
)

# Variables that should never appear in the namespace snapshot
_INTERNAL_VARS = frozenset({
    "np", "litert_lm", "shared_globals", "ai", "readline",
    "engine", "conversation", "reset_ai", "SYSTEM_CONSTRAINT",
    "_INTERNAL_VARS",
})

print("[Initialization] Loading LiteRT-LM Engine on CPU...")
engine = litert_lm.Engine(
    MODEL_PATH,
    backend=litert_lm.Backend.CPU,
    cache_dir="/tmp/litert-lm-cache",
    enable_speculative_decoding=True,
)

# ── KEY FIX: Create ONE persistent conversation ──────────────────────────
# The KV cache stays warm across ai() calls. Only new query tokens need
# prefilling, instead of re-prefilling the full system prompt every time.
# The system constraints are set once via system_message= (never re-prefilled).
conversation = engine.create_conversation(system_message=SYSTEM_CONSTRAINT)

shared_globals = {
    "np": np,
    "current_coordinates": np.array([5.0, -3.5]),
    "matrix_z": np.array([[1.0, 2.0], [3.0, np.nan]])
}

def _build_namespace_snapshot():
    """Build a short string describing current variable state."""
    snapshot = "Current RAM State:\n"
    for var_name, var_val in list(shared_globals.items()):
        if var_name in _INTERNAL_VARS:
            continue
        if isinstance(var_val, np.ndarray):
            snapshot += f"- {var_name}:\n{var_val}\n"
        else:
            snapshot += f"- {var_name}: {var_val}\n"
    return snapshot

def ai(query_string):
    """
    Fast co-interpreter call. Reuses the persistent conversation so the
    KV cache is warm — only the short query + namespace snapshot are prefilled.
    """
    user_message = f"{_build_namespace_snapshot()}\nQuery: {query_string}"

    print("\n[AI Co-Interpreter]: ", end="")
    sys.stdout.flush()

    for chunk in conversation.send_message_async(user_message):
        try:
            text_piece = chunk["content"][0]["text"]
            sys.stdout.write(text_piece)
            sys.stdout.flush()
        except (KeyError, IndexError):
            pass
    print()

def reset_ai():
    """
    Discard the current conversation and start a fresh one.
    Call this if history grows too long and you want to reset context.
    """
    global conversation
    conversation.close()
    conversation = engine.create_conversation(system_message=SYSTEM_CONSTRAINT)
    print("[AI] Conversation reset — KV cache cleared.")

shared_globals["ai"] = ai
shared_globals["reset_ai"] = reset_ai

if readline:
    readline.parse_and_bind("tab: complete")
    readline.set_history_length(1000)

banner_msg = """
====================================================================
     XYFLOW HYPER-SPEED DUAL INTERPRETER (v7.0-FIXED)
====================================================================
System status: OPERATIONAL (Warm KV-cache execution)
Brevity lock:  ACTIVE (system_message-enforced)

Performance fixes applied:
  • Persistent conversation (warm KV cache)
  • System constraints via system_message= (prefilled once)
  • Speculative decoding enabled
  • Compiled artifact caching enabled

Try:
>>> ai("Look at matrix_z. Is it stable or is there an issue?")
>>> reset_ai()   # optional: clear conversation history
====================================================================
"""

console = code.InteractiveConsole(locals=shared_globals)
console.interact(banner=banner_msg, exitmsg="\n[Shutdown] Safely clearing engine allocations.")