import os
import code
import sys
import numpy as np
from openai import OpenAI  # pip install openai

try:
    import readline
except ImportError:
    readline = None

# Point the OpenAI client to your local LiteRT-LM server
client = OpenAI(
    base_url="http://127.0.0.1:8000/v1",
    api_key="required-but-ignored-by-local-server" 
)

# Shared global namespace inside the interactive console
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):
    """
    Client-side AI function utilizing the native OpenAI client 
    to talk directly to the background litert-lm server.
    """
    # 1. Capture a 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", "shared_globals", "ai", "readline", "client", "sys"]: 
            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"
            
    print("\n[AI Co-Interpreter]: ", end="")
    sys.stdout.flush()
    
    # 2. Query the server using standard OpenAI Chat Completions with streaming
    try:
        stream = client.chat.completions.create(
            model="gemma-4-E2B-it.litertlm", # Replace with your exact server model ID if needed
            messages=[
                {
                    "role": "system", 
                    "content": "You are an AI co-interpreter inside a live Python runtime. Analyze the snapshot of variables provided, and answer the user query concisely."
                },
                {
                    "role": "user", 
                    "content": f"{namespace_snapshot}\nQuery: {query_string}"
                }
            ],
            stream=True
        )
        
        # 3. Stream the response directly to the console token by token
        for chunk in stream:
            token = chunk.choices[0].delta.content
            if token:
                sys.stdout.write(token)
                sys.stdout.flush()
                
    except Exception as e:
        print(f"\n[Error] Failed to communicate with LiteRT-LM server: {e}")
    
    print("\n")

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.1-OPENAI-API)
====================================================================
System status: OPERATIONAL (Connected to Native litert-lm Server)
Arrow history: READY (Press UP/DOWN arrows).

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

console = code.InteractiveConsole(locals=shared_globals)
console.interact(banner=banner_msg, exitmsg="\n[Shutdown] Closing Client Interpreter.")