import os
import re
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")

print("[Initialization] Instantiating LiteRT-LM Engine on CPU...")
# Explicitly binding to CPU to guarantee stability across drivers
engine = litert_lm.Engine(MODEL_PATH, backend=litert_lm.Backend.CPU())

# Define our shared global namespace
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 ai(query_string):
    """
    Optimized State Inspector function. By creating a temporary, non-accumulating 
    session, the CPU skips calculating historical context tables and handles 
    queries quickly.
    """
    # 1. Capture a clean text map of the active variable space
    namespace_snapshot = "Active State Snapshot:\n"
    for var_name, var_val in list(shared_globals.items()):
        if var_name in ["np", "litert_lm", "shared_globals", "ai", "readline"]: 
            continue
        if isinstance(var_val, np.ndarray):
            namespace_snapshot += f"- {var_name} (Shape {var_val.shape}):\n{var_val}\n"
        else:
            namespace_snapshot += f"- {var_name}: {var_val}\n"
            
    system_instruction = litert_lm.Message.system(
        "You are an AI co-interpreter inside a live Python runtime. "
        "Analyze the snapshot of variables provided, and answer the user query concisely."
    )
    
    full_prompt = (
        f"{namespace_snapshot}\n"
        f"Query: {query_string}"
    )
    
    # 2. Fire a clean session that immediately closes after processing execution
    with engine.create_conversation(messages=[system_instruction]) as conversation:
        response = conversation.send_message(full_prompt)
        ai_output = response["content"][0]["text"]
    
    print(f"\n[AI Co-Interpreter]:\n{ai_output}\n")

shared_globals["ai"] = ai

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

banner_msg = """
====================================================================
      XYFLOW DUAL PYTHON INTERPRETER MANIFOLD (v4.0-CPU-OPT)
====================================================================
System status: OPERATIONAL (Context-optimized CPU execution)
Arrow history: READY (Press UP/DOWN arrows).

Objects:
  - current_coordinates : active trajectory vector
  - matrix_z            : multi-dimensional array
  - ai("query")         : evaluate data variable anomalies

Try typing:
>>> ai("Look at matrix_z. What element is causing structural chaos?")
====================================================================
"""

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