import os
import re
import code
import atexit
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...")
# MODIFIED: Attempting to hook into your registered GPU acceleration backend for lightning speed
try:
    engine = litert_lm.Engine(MODEL_PATH, backend=litert_lm.Backend.GPU())
    print("-> Status: Running with GPU Acceleration.")
except Exception:
    engine = litert_lm.Engine(MODEL_PATH, backend=litert_lm.Backend.CPU())
    print("-> Status: GPU fallback failed. Running on Multi-threaded CPU.")

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 accurately and concisely."
)

# Active conversation session
conversation = engine.create_conversation(messages=[system_instruction])

# FIXED: Explicit cleanup hook to gracefully close the session and prevent memory leaks/destructor warnings
def cleanup_session():
    print("\n[Shutdown] Unlinking memory maps and closing active model sessions...")
    try:
        conversation.close()
    except Exception:
        pass
atexit.register(cleanup_session)

# 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 portal function. Evaluates variables locally to prevent context inflation.
    """
    # Dynamic variables readout
    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", "atexit"]: 
            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"
            
    full_prompt = (
        f"{namespace_snapshot}\n"
        f"Query: {query_string}"
    )
    
    # Send current state cleanly
    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 (v3.0-GPU)
====================================================================
System status: OPERATIONAL
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?")
====================================================================
"""

# Early baseline warm-up call to pre-compile the tensor allocations
_ = conversation.send_message("Baseline system handshake warmup.")

console = code.InteractiveConsole(locals=shared_globals)
console.interact(banner=banner_msg, exitmsg="")