#!/usr/bin/env python3
"""
explore, clean, quit
nyi
-return to dock (need several hours for 100 epochs x 1 image rock 62x3 3 hours for 4 images for a good return point for the dock)
-floor mapping (tried photo compare with long term sqlite storage, mostly a bust, will consider signs pointing to rooms for instructional cleaning)
"""
import os, re, json, time, base64, requests
from datetime import datetime
from pydantic import BaseModel, ValidationError
from collections import deque
from typing import Any, List

# --- CONFIG ---
CFG = {"CTRL_URL":"http://CONTROL-SERVER-DEVICE-IP:5000","LLM_URL":"http://INFER-SOURCE-IP:1234/v1/chat/completions","MODEL":"google/gemma-3-4b (use a vision model)","API_TOUT":30,"LLM_TOUT":90,"PAUSE":0.25,"FRAME_DIR":"frames","BAT_MIN":20.0,"HISTORY_LEN":4}
os.makedirs(CFG["FRAME_DIR"], exist_ok=True)
class RoombaCmd(BaseModel): maneuver: str; duration: float; scene: Any; reasoning: str
def _post(endpoint, payload=None, timeout=None):
    url = f"{CFG['CTRL_URL']}/{endpoint}"; r = requests.post(url, json=payload, timeout=timeout or CFG["API_TOUT"]); r.raise_for_status(); return r.json()
def _get(endpoint):
    url = f"{CFG['CTRL_URL']}/{endpoint}"; r = requests.get(url, timeout=CFG["API_TOUT"]); r.raise_for_status(); return r.json()
def save_frame(b64_data):
    fn = datetime.now().strftime("%Y%m%d_%H%M%S") + ".jpg"; open(os.path.join(CFG["FRAME_DIR"], fn), "wb").write(base64.b64decode(b64_data)); print(f"---- Saved frame {fn}")

# --- CORE AI LOGIC ---
def get_ai_command(img_b64, status, history) -> RoombaCmd:
    sys_prompt = (
        "You are the control system for a Roomba robot. Your **primary mission is to find and vacuum small debris** like crumbs, lint, and dust. "
        "The camera image has a visual HUD: a red 'CLOSE (DANGER)' zone and a green 'MEDIUM (TARGET)' zone. "
        "Analyze the image and produce *only* this JSON object:\n"
        "{\n"
        "  \"maneuver\": \"...\",   // one of: forward, backward, left, right\n"
        "  \"duration\": ...,       // seconds, float e.g. 0.7 (0.4 to 0.8)\n"
        "  \"scene\": \"...\",      // describe what you see, including obstacles and targets\n"
        "  \"reasoning\": \"...\"   // why this action helps achieve the goal of cleaning debris\n"
        "}\n\n"
        "**OBJECT IDENTIFICATION RULES:**\n"
        "- **Hard Obstacles (Avoid):** Walls, furniture legs, large boxes, pets, shoes, bowls.\n"
        "- **Soft Obstacles (Ignore/Push):** Rugs, blankets, socks, small toys (like LEGOs).\n"
        "- **Targets (Attack):** Crumbs, dust bunnies, lint, loose cat food, litter sand, dirt.\n\n"
        "**HIERARCHY OF ACTIONS (obey in order from 1 to 4):**\n\n"
        "1. **AVOID HARD OBSTACLES:** If a **Hard Obstacle** is in the red 'CLOSE' zone, you **MUST** turn 'left' or 'right' to avoid it. This is your most important safety rule.\n\n"
        "2. **ATTACK TARGETS:** If you see a **Target** (debris) anywhere in the image, your goal is to drive towards it. If it's ahead, move 'forward'. If it's to the side, turn towards it.\n\n"
        "3. **IGNORE SOFT OBSTACLES:** If the only thing in your path is a **Soft Obstacle**, you should treat it as clear floor and move 'forward' to clean over it.\n\n"
        "4. **EXPLORE FOR TARGETS:** If the path is completely clear of all objects, your goal is to find more debris. Drive 'forward' towards the largest open area to continue your search.\n"
    )
    history_str = ", ".join(history) if history else "None"
    messages = [{"role": "system", "content": sys_prompt}, {"role": "user", "content": [{"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{img_b64}"}}, {"type": "text", "text": f"Current Status: {status['battery_percent']:.1f}% battery. Previous Actions: [{history_str}]."}]}]
    try:
        resp = requests.post(CFG["LLM_URL"], json={"model": CFG["MODEL"], "messages": messages, "temperature": 0.0, "max_tokens": 300}, timeout=CFG["LLM_TOUT"])
        resp.raise_for_status(); raw_text = resp.json()["choices"][0]["message"]["content"]; json_match = re.search(r"\{.*\}", raw_text, re.S)
        if not json_match: raise ValueError("No JSON object found")
        obj = json.loads(json_match.group(0));
        if 'duration' not in obj: obj['duration'] = 0.7
        if 'scene' not in obj: obj['scene'] = "AI did not provide scene."
        if 'reasoning' not in obj: obj['reasoning'] = "AI did not provide reasoning."
        return RoombaCmd(**obj)
    except Exception as e: return RoombaCmd(maneuver="left", duration=1.0, scene="AI_FALLBACK", reasoning=f"Error: {e}")

# --- MENU LOGIC ---
def run_mission(is_cleaning: bool):
    mode = "CLEANING" if is_cleaning else "EXPLORATION"; print(f"\n XOXO Starting {mode} mission... Press Ctrl+C to stop.")
    action_history = deque(maxlen=CFG["HISTORY_LEN"])
    has_undocked = False
    try:
        if is_cleaning: print("🧹 Starting vacuum motors..."); _post("vacuum", {"on": True})
        while True:
            status = _get("status")
            if not has_undocked and status.get("docked"):
                print("⚓ Robot is docked. Executing undock maneuver..."); _post("undock")
                print("   ↳ Maneuver complete. Pausing 3s for hardware to settle..."); has_undocked = True
                action_history.append("undock"); time.sleep(3.0); continue
            if status.get("battery_percent", 100) < CFG["BAT_MIN"]: print("⚡ Battery low! Ending mission."); break
            try:
                bumper_status = _get("bumper")
                if bumper_status.get("bumped"):
                    print("PHYSICAL BUMP DETECTED! Overriding AI to escape.")
                    _post("move", {"maneuver": "backward", "duration": 1.0})
                    _post("move", {"maneuver": "left", "duration": 0.8})
                    action_history.append("BUMP_ESCAPE"); time.sleep(CFG["PAUSE"]); continue
            except Exception as e: print(f"TTTT Could not get bumper status: {e}")
            img_b64 = _get("image").get("image_base64", "");
            if not img_b64: cmd = RoombaCmd(maneuver="left", duration=1.0, scene="NO_IMAGE", reasoning="No image received")
            else: save_frame(img_b64); cmd = get_ai_command(img_b64, status, list(action_history))
            cmd.duration = max(0.4, min(cmd.duration, 0.8)); scene_for_print = cmd.scene if isinstance(cmd.scene, str) else json.dumps(cmd.scene)
            print(f"\n{status.get('battery_percent', 0):.1f}% | History: {list(action_history)}")
            print(f"AI Reasoning: {cmd.reasoning}")
            print(f"↳ Executing: {cmd.maneuver:<8} for {cmd.duration:.1f}s. (Scene: {scene_for_print})")
            _post("move", {"maneuver": cmd.maneuver, "duration": cmd.duration}, timeout=cmd.duration + 10)
            action_history.append(cmd.maneuver); time.sleep(CFG["PAUSE"])
    except KeyboardInterrupt: print("\nXOXO User interrupted mission.")
    except Exception as e: print(f"XXXX Unrecoverable mission error: {e}")
    finally:
        print("--- Mission Ending ---");
        if is_cleaning:
            print("Stopping vacuum motors...")
            try: _post("vacuum", {"on": False})
            except Exception as e: print(f"Could not stop vacuum: {e}")
        try: _post("move", {"maneuver": "stop", "duration": 0})
        except Exception as e: print(f"Could not send stop command: {e}")
        print("Mission complete.")
def main_menu():
    while True:
        print("\n"+"="*40); print("  ROOMBA AI CONTROL (VACUUM HUNTER)"); print("="*40)
        try:
            status = _get("status"); print(f"Status: |||| {status.get('battery_percent', 'N/A')}% | ⚓ Docked: {status.get('docked', 'N/A')}")
        except Exception as e: print(f"Status: TTTT Could not connect to control server: {e}"); time.sleep(3); continue
        print("\nCHOOSE MISSION:\n  1. Explore (No Vacuum)\n  2. Clean   (With Vacuum)\n  q. Quit")
        choice = input("Enter your choice: ").lower().strip()
        if choice == '1': run_mission(is_cleaning=False)
        elif choice == '2': run_mission(is_cleaning=True)
        elif choice == 'q': print("Exiting program."); break
        else: print("Invalid choice.")
        time.sleep(1)

if __name__ == "__main__": main_menu()
