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

# 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])

# 2. Define the Global Namespace and the Shared "AI Button"
# This dictionary holds the variables that BOTH you and the AI can see/modify.
shared_globals = {
    "np": np,
    "current_coordinates": np.array([5.0, -3.5]),
    "matrix_z": np.array([[1.0, 2.0], [3.0, np.nan]])  # Injected a NaN to test inspection!
}

def ai(query_string):
    """
    The 'ai()' function is your direct portal to the model.
    When you call it in your terminal, it builds a localized text snapshot of your 
    active global variables, sends it to Gemma, and updates the console state.
    """
    # 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"]: 
            continue  # Skip framework modules
        
        # Format strings or numpy layouts safely
        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"
            
    # Combine your question with the live RAM snapshot
    full_prompt = (
        f"{namespace_snapshot}\n"
        f"User Query: {query_string}"
    )
    
    # Process through the live conversation context
    response = conversation.send_message(full_prompt)
    ai_output = response["content"][0]["text"]
    
    print(f"\n[AI Co-Interpreter]:\n{ai_output}\n")
    
    # Check if the AI tried to provide any direct parameter overrides (e.g. modifying mu or omega)
    # This maintains the parameter injection feature!
    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

# Register the portal function inside the console so the user can invoke it
shared_globals["ai"] = ai

# 3. Spin up the Interactive Dual-Stream Console
banner_msg = """
====================================================================
      XYFLOW DUAL PYTHON INTERPRETER MANIFOLD (CCT v2.0)
====================================================================
You are in a live shared namespace console with your local SLM.

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 * 2
>>> ai("Look at matrix_z. Is it stable or is there a numerical breakdown?")
====================================================================
"""

# Open the console shell using our shared global dictionary
console = code.InteractiveConsole(locals=shared_globals)
console.interact(banner=banner_msg, exitmsg="\nShutting down dual-interpreter environment. Threads cleared.")