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

# FIXED: Import readline to natively hook up arrow key command history
try:
    import readline
except ImportError:
    # Fallback for systems without readline (like bare windows without pyreadline)
    readline = None

# 1. Establish the explicit binary path to your local model
MODEL_PATH = os.path.expanduser("~/.litert-lm/models/gemma-4-E2B-it.litertlm/model.litertlm")

print("[Initialization] Loading LiteRT-LM Engine into memory space...")
engine = litert_lm.Engine(MODEL_PATH, backend=litert_lm.Backend.CPU())

# System instructions to lock the model into a State-Inspection persona
system_instruction = litert_lm.Message.system(
    "You are an AI co-interpreter embedded directly inside a live Python runtime environment. "
    "You have access to the global variable namespace wrapper. When asked about arrays or variables, "
    "inspect them using math or code logic, then explain your findings or suggest adjustments concisely."
)
conversation = engine.create_conversation(messages=[system_instruction])

# FIXED: Warm up the engine silently so the first real user call has ZERO startup lag
print("[Warmup] Priming model weights and cache layers...")
_ = conversation.send_message("Warmup token verification step.")

# 2. Define the Global Namespace and the Shared "AI Button"
shared_globals = {
    "np": np,
    "current_coordinates": np.array([5.0, -3.5]),
    "matrix_z": np.array([[1.0, 2.0], [3.0, np.nan]])  # NaN element intact for debugging
}

def ai(query_string):
    """
    The 'ai()' function grabs a text snapshot of the live RAM namespace,
    pipes it along with your prompt to the local SLM, and handles parameter hooks.
    """
    # Create a clean text representation of what variables are currently in memory
    namespace_snapshot = "Active Variables in Global Workspace:\n"
    for var_name, var_val in shared_globals.items():
        if var_name in ["np", "litert_lm", "shared_globals", "ai", "readline", "sys"]: 
            continue  # Skip structural infrastructure modules
        
        if isinstance(var_val, np.ndarray):
            namespace_snapshot += f"- {var_name} (NumPy Array, Shape {var_val.shape}):\n{var_val}\n"
        else:
            namespace_snapshot += f"- {var_name}: {var_val}\n"
            
    full_prompt = (
        f"{namespace_snapshot}\n"
        f"User Query: {query_string}"
    )
    
    # Send message synchronously through our pre-warmed context loop
    response = conversation.send_message(full_prompt)
    ai_output = response["content"][0]["text"]
    
    print(f"\n[AI Co-Interpreter]:\n{ai_output}\n")
    
    # Catch any direct numeric parameter updates generated by the model
    try:
        mu_match = re.search(r"mu\s*=\s*([\d\.]+)", ai_output)
        if mu_match:
            shared_globals["mu"] = float(mu_match.group(1))
            print(f"-> Code Hook: AI successfully registered 'mu = {shared_globals['mu']}' in workspace.")
    except Exception:
        pass

# Bind the interactive function straight into the global dictionary map
shared_globals["ai"] = ai

# 3. Setup the Interactive Command Line History
if readline:
    # Use standard tab completion inside the prompt for variables
    readline.parse_and_bind("tab: complete")
    # Optional: Set a clean baseline memory capacity for your arrow key scroll-back
    readline.set_history_length(1000)

# 4. Spin up the Interactive Console Manifold
banner_msg = """
====================================================================
      XYFLOW DUAL PYTHON INTERPRETER MANIFOLD (v2.0-HYPER)
====================================================================
System status: FULLY ARMED & PRE-WARMED.
Arrow history: READY (Press UP/DOWN keys to recall commands).

Available objects:
  - current_coordinates : active trajectory vector
  - matrix_z            : an experimental multi-dimensional array
  - ai("your question")  : passes your query + a live RAM snapshot to the SLM

Try typing: 
>>> current_coordinates * 5
>>> ai("Look at matrix_z. Is it stable or is there a numerical breakdown?")
====================================================================
"""

console = code.InteractiveConsole(locals=shared_globals)
console.interact(banner=banner_msg, exitmsg="\nShutting down dual-interpreter environment. Memory maps safely unlinked.")