import ollama, base64, json, time
import pyrealsense2 as rs
import numpy as np, cv2
import zmq

HOLOSOMA_IP = "127.0.0.1"
HOLOSOMA_PORT = 5556

def capture_frame():
    pipeline = rs.pipeline()
    cfg = rs.config()
    cfg.enable_stream(rs.stream.color, 640, 480, rs.format.bgr8, 30)
    pipeline.start(cfg)
    for _ in range(5):
        pipeline.wait_for_frames()
    frames = pipeline.wait_for_frames()
    img = np.asanyarray(frames.get_color_frame().get_data())
    pipeline.stop()
    cv2.imwrite('/tmp/marcus_eye.jpg', img)
    return '/tmp/marcus_eye.jpg'

def ask_qwen(image_path, command):
    with open(image_path, 'rb') as f:
        img_b64 = base64.b64encode(f.read()).decode()
    prompt = f"""أنت ماركس، روبوت ذكي يتنقل داخل المبنى.
You are Marcus, an intelligent indoor navigation robot.

User command: "{command}"

Look at the camera image. Respond with ONLY one line:
FORWARD [0.1 to 1.0 meters]
LEFT [5 to 45 degrees]  
RIGHT [5 to 45 degrees]
STOP [reason]
ARRIVED"""

    response = ollama.chat(
        model='qwen2.5vl:7b',
        messages=[{
            'role': 'user',
            'content': prompt,
            'images': [img_b64]
        }]
    )
    return response['message']['content'].strip().split('\n')[0]

def send_to_robot(action):
    print(f"  Robot action: {action}")
    parts = action.upper().split()
    if not parts:
        return
    cmd = parts[0]
    val = float(parts[1]) if len(parts) > 1 else 0

    if cmd == "FORWARD":
        print(f"  Walking forward {val}m")
    elif cmd == "LEFT":
        print(f"  Turning left {val} degrees")
    elif cmd == "RIGHT":
        print(f"  Turning right {val} degrees")
    elif cmd == "STOP":
        print(f"  Stopping: {' '.join(parts[1:])}")
    elif cmd == "ARRIVED":
        print("  Destination reached!")

if __name__ == "__main__":
    print("=" * 50)
    print("Marcus Navigation Brain")
    print("Powered by Qwen2.5-VL on Jetson Orin NX")
    print("Speaks Arabic + English")
    print("=" * 50)
    print("Type your command (or 'quit'):\n")

    while True:
        try:
            cmd = input("Command: ").strip()
            if cmd.lower() in ['quit', 'exit', 'خروج']:
                print("Marcus shutting down.")
                break
            if not cmd:
                continue
            print("Capturing camera frame...")
            frame = capture_frame()
            print("Qwen2.5-VL thinking...")
            t0 = time.time()
            action = ask_qwen(frame, cmd)
            elapsed = time.time() - t0
            print(f"Decision ({elapsed:.1f}s): {action}")
            send_to_robot(action)
            print()
        except KeyboardInterrupt:
            print("\nStopped.")
            break