import os
import sys
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] Loading context-optimized LiteRT-LM Engine on CPU...")
# FIXED: Removed the unsupported 'sampler_params' argument to respect the API signature
engine = litert_lm.Engine(MODEL_PATH, backend=litert_lm.Backend.CPU())

# Shared global variable workspace
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 loop. Creates a clean, fast session for each turn
    to keep CPU execution times consistent.
    """
    # Build clean active global snapshot tracking string
    namespace_snapshot = "RAM 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}:\n{var_val}\n"
        else:
            namespace_snapshot += f"- {var_name}: {var_val}\n"
            
    # CRITICAL BREVITY SYSTEM CONSTRAINTS: Embedded directly into the prompt frame
    full_prompt = (
        "SYSTEM CONSTRAINT: Act as an embedded Python co-interpreter. "
        "Answer the user query in exactly ONE brief sentence or direct assignment. "
        "No greetings, no preambles, no conversational fluff.\n\n"
        f"{namespace_snapshot}\n"
        f"Query: {query_string}"
    )
    
    print("\n[AI Co-Interpreter]: ", end="")
    sys.stdout.flush()
    
    # Open isolated non-accumulating conversation block
    with engine.create_conversation() as conversation:
        stream = conversation.send_message_async(full_prompt)
        for chunk in stream:
            try:
                text_piece = chunk["content"][0]["text"]
                sys.stdout.write(text_piece)
                sys.stdout.flush()
            except (KeyError, IndexError):
                pass
    print("\n")

shared_globals["ai"] = ai

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

banner_msg = """
====================================================================
     XYFLOW HYPER-SPEED DUAL INTERPRETER (v6.2-FINAL)
====================================================================
System status: OPERATIONAL (Context-optimized execution)
Brevity lock:  ACTIVE (Prompt-enforced short responses)

Try standard sequential test:
>>> ai("Look at matrix_z. Is it stable or is there an issue?")
====================================================================
"""

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