import os
import re
import code
import sys
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...")
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):
    """
    Streaming State Inspector function. Automatically hooks a live text token callback
    straight into standard out so the user sees immediate computational progress.
    """
    # 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}"
    )
    
    print("\n[AI Co-Interpreter]: ", end="")
    sys.stdout.flush()

    # The token callback function executed on every step generation
    def token_callback(text_chunk):
        sys.stdout.write(text_chunk)
        sys.stdout.flush()
    
    # Run standalone non-accumulating context conversation block with streaming enabled
    with engine.create_conversation(messages=[system_instruction]) as conversation:
        # Note: Depending on your specific litert_lm wrapper version, the streaming hook
        # is handled either via callback=token_callback or token_callback=token_callback
        conversation.send_message(full_prompt, callback=token_callback)
    
    print("\n") # Append trailing layout line breaks after streaming finishes

shared_globals["ai"] = ai

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

banner_msg = """
====================================================================
     XYFLOW STREAMING DUAL PYTHON INTERPRETER (v5.0-STREAM)
====================================================================
System status: OPERATIONAL (Real-time token callback processing)
Arrow history: READY (Press UP/DOWN arrows).

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

Try typing:
>>> matrix_z[1, 1] = 4.0
>>> ai("I just patched the array. Analyze matrix_z again. Is it stable?")
====================================================================
"""

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