"""
marcus_autonomous.py — Autonomous Exploration Mode
====================================================
Project  : Marcus | YS Lootah Technology
Hardware : Unitree G1 EDU + Jetson Orin NX

Marcus moves freely through the office, identifies areas and objects,
builds a live map of what it sees, and saves everything to a session folder.

Commands
--------
  auto on        — start autonomous exploration
  auto off       — stop autonomous exploration
  auto status    — show current exploration state
  auto save      — save current map snapshot
  auto summary   — LLaVA summary of what was explored

How it works
------------
  1. Marcus moves forward, scanning with YOLO every 0.4s
  2. Every N steps: LLaVA assesses the scene (area type, objects, notes)
  3. Odometry records position at each observation
  4. All data saved to map/map_XXX_DATE/ folder
  5. Robot avoids walls by turning when blocked

Files saved
-----------
  ~/Models_marcus/map/map_001_2026-04-05/
    observations.json   — [{step, time, x, y, area_type, objects, observation}]
    path.json           — [{x, y, heading, t}] — full path walked
    summary.txt         — auto-generated LLaVA summary
    frames/             — camera captures at interesting points

Import in marcus_llava.py
-------------------------
  from marcus_autonomous import AutonomousMode
  auto = AutonomousMode(
      get_frame_fn=get_frame,
      send_vel_fn=send_vel,
      gradual_stop_fn=gradual_stop,
      yolo_sees_fn=yolo_sees,
      yolo_summary_fn=yolo_summary,
      yolo_all_classes_fn=yolo_all_classes,
      yolo_closest_fn=yolo_closest,
      odom_fn=_odom_pos,
      call_llava_fn=_call_llava,
      patrol_prompt=PATROL_PROMPT,
      mem=mem,
      models_dir=MODELS_DIR,
  )
  auto.enable()    # start exploration
  auto.disable()   # stop exploration
  auto.status()    # print current state
"""

import json
import time
import threading
import os
import base64
from datetime import datetime
from pathlib import Path


# ══════════════════════════════════════════════════════════════════════════════
# CONFIGURATION
# ══════════════════════════════════════════════════════════════════════════════

LLAVA_EVERY_N_STEPS  = 5      # call LLaVA every N movement steps (saves GPU)
YOLO_CHECK_INTERVAL  = 0.4    # seconds between YOLO checks
FORWARD_DURATION     = 1.5    # seconds per forward step
TURN_DURATION        = 2.0    # seconds to turn when obstacle detected
PERSON_STOP_DIST     = 0.3    # stop if person closer than this (relative)
MAX_OBSERVATIONS     = 500    # max observations before auto-stop
SAVE_FRAMES          = True   # save camera frames at interesting points
INTERESTING_COOLDOWN = 5.0    # seconds between saving "interesting" frames


class AutonomousMode:
    """
    Autonomous office exploration.

    Thread-safe. enable() starts a background thread. disable() stops it.
    All observations saved to disk automatically.

    State machine:
      IDLE → EXPLORING → IDLE
      Any state → STOPPING → IDLE (on disable())
    """

    def __init__(self, get_frame_fn, send_vel_fn, gradual_stop_fn,
                 yolo_sees_fn, yolo_summary_fn, yolo_all_classes_fn,
                 yolo_closest_fn, odom_fn, call_llava_fn,
                 patrol_prompt: str, mem=None, models_dir: str = None):

        self._get_frame       = get_frame_fn
        self._send_vel        = send_vel_fn
        self._gradual_stop    = gradual_stop_fn
        self._yolo_sees       = yolo_sees_fn
        self._yolo_summary    = yolo_summary_fn
        self._yolo_all_classes = yolo_all_classes_fn
        self._yolo_closest    = yolo_closest_fn
        self._odom_pos        = odom_fn
        self._call_llava      = call_llava_fn
        self._patrol_prompt   = patrol_prompt
        self._mem             = mem
        if models_dir is None:
            models_dir = str(Path(__file__).resolve().parent.parent.parent / "Data" / "Brain")
        self._models_dir      = Path(models_dir)

        # State
        self._enabled     = False
        self._thread      = None
        self._lock        = threading.Lock()

        # Current exploration session
        self._map_dir         = None
        self._observations    = []
        self._path            = []
        self._step            = 0
        self._start_time      = None
        self._last_interesting = 0.0

        # Turn tracking — alternate left/right when blocked
        self._last_turn = "left"

        # Stats
        self._area_counts = {}   # {area_type: count}
        self._all_objects = set()

    # ── PUBLIC API ─────────────────────────────────────────────────────────────

    def enable(self):
        """Start autonomous exploration."""
        with self._lock:
            if self._enabled:
                print("  [Auto] Already running — use 'auto off' to stop first")
                return

            self._enabled    = True
            self._step       = 0
            self._start_time = time.time()
            self._observations = []
            self._path         = []
            self._area_counts  = {}
            self._all_objects  = set()
            self._map_dir      = self._create_map_dir()

            self._thread = threading.Thread(
                target=self._explore_loop,
                daemon=True,
                name="auto-explore"
            )
            self._thread.start()

        print(f"\n  [Auto] Exploration started")
        print(f"  [Auto] Map folder: {self._map_dir}")
        print(f"  [Auto] Type 'auto off' to stop\n")

    def disable(self):
        """Stop autonomous exploration and save results."""
        with self._lock:
            if not self._enabled:
                print("  [Auto] Not running")
                return
            self._enabled = False

        print("\n  [Auto] Stopping exploration...")

        # Wait for thread to finish
        if self._thread and self._thread.is_alive():
            self._thread.join(timeout=5.0)

        self._gradual_stop()
        self._save_session()
        self._print_summary()

    def status(self):
        """Print current exploration state."""
        with self._lock:
            running = self._enabled
            step    = self._step
            obs     = len(self._observations)

        if not running:
            print("  [Auto] Status: IDLE")
            if self._map_dir:
                print(f"  [Auto] Last map: {self._map_dir}")
            return

        elapsed = time.time() - (self._start_time or time.time())
        mins = int(elapsed // 60)
        secs = int(elapsed % 60)

        print(f"  [Auto] Status: EXPLORING")
        print(f"  [Auto] Duration: {mins}m {secs}s")
        print(f"  [Auto] Steps: {step} | Observations: {obs}")

        if self._area_counts:
            areas = ", ".join(f"{k}:{v}" for k, v in sorted(self._area_counts.items()))
            print(f"  [Auto] Areas seen: {areas}")

        if self._all_objects:
            print(f"  [Auto] Objects found: {', '.join(sorted(self._all_objects))}")

        pos = self._odom_pos()
        if pos:
            print(f"  [Auto] Position: x={pos['x']:.2f} y={pos['y']:.2f} heading={pos['heading']:.1f}°")

    def is_enabled(self) -> bool:
        with self._lock:
            return self._enabled

    def save_snapshot(self):
        """Save current state to disk without stopping."""
        self._save_observations()
        self._save_path()
        print(f"  [Auto] Snapshot saved to {self._map_dir}")

    # ── EXPLORATION LOOP ────────────────────────────────────────────────────────

    def _explore_loop(self):
        """
        Main autonomous exploration loop.

        Steps:
          1. Check for person in path (safety stop)
          2. Check YOLO for interesting objects
          3. Move forward
          4. Every N steps: call LLaVA for scene assessment
          5. On obstacle: turn and continue
          6. Log position + observations
        """
        consecutive_blocks = 0   # count consecutive blocked steps

        while self._enabled:
            self._step += 1

            # ── Safety: stop if person too close ─────────────────────────────
            if self._yolo_sees("person"):
                closest = self._yolo_closest("person")
                if closest and closest.distance_estimate == "very close":
                    print(f"  [Auto] Person very close — pausing 2s")
                    self._gradual_stop()
                    time.sleep(2.0)
                    continue

            # ── Record YOLO detections ────────────────────────────────────────
            detected_classes = self._yolo_all_classes()
            for cls in detected_classes:
                self._all_objects.add(cls)

            # ── Record odometry path point ────────────────────────────────────
            pos = self._odom_pos()
            if pos:
                self._path.append({
                    "step": self._step,
                    "t":    round(time.time() - self._start_time, 1),
                    "x":    pos["x"],
                    "y":    pos["y"],
                    "h":    pos["heading"],
                })

            # ── LLaVA scene assessment every N steps ─────────────────────────
            if self._step % LLAVA_EVERY_N_STEPS == 0:
                self._assess_scene(pos, detected_classes)

            # ── Movement decision ─────────────────────────────────────────────
            if consecutive_blocks >= 3:
                # Stuck — turn more aggressively
                print(f"  [Auto] Stuck — turning {self._last_turn} 180°")
                self._turn(self._last_turn, TURN_DURATION * 2)
                consecutive_blocks = 0
                continue

            # Move forward
            blocked = self._move_forward()

            if blocked:
                consecutive_blocks += 1
                # Alternate left/right turns to explore both directions
                turn_dir = "left" if self._last_turn == "right" else "right"
                self._last_turn = turn_dir
                print(f"  [Auto] Obstacle — turning {turn_dir}")
                self._turn(turn_dir, TURN_DURATION)
            else:
                consecutive_blocks = 0

            # ── Max observations check ────────────────────────────────────────
            if len(self._observations) >= MAX_OBSERVATIONS:
                print(f"  [Auto] Max observations ({MAX_OBSERVATIONS}) reached — stopping")
                self._enabled = False
                break

            # No trailing sleep — _move_forward() takes FORWARD_DURATION,
            # _turn() takes TURN_DURATION, and LLaVA assessment is ~1-2s.
            # The body always consumes real wall time, so an extra sleep here
            # would be pure dead time.

        # Clean up
        self._gradual_stop()

    def _move_forward(self) -> bool:
        """
        Move forward for FORWARD_DURATION seconds.
        Returns True if blocked (no actual movement detected).
        """
        # TODO: integrate LiDAR when available for true obstacle detection
        # For now: move and assume not blocked
        t0 = time.time()
        while time.time() - t0 < FORWARD_DURATION and self._enabled:
            self._send_vel(vx=0.25)
            time.sleep(0.05)
        self._send_vel(0, 0, 0)
        time.sleep(0.1)
        return False   # Not blocked — no LiDAR yet

    def _turn(self, direction: str, duration: float):
        """Turn left or right for given duration."""
        vyaw = 0.25 if direction == "left" else -0.25
        t0 = time.time()
        while time.time() - t0 < duration and self._enabled:
            self._send_vel(vyaw=vyaw)
            time.sleep(0.05)
        self._send_vel(0, 0, 0)
        time.sleep(0.2)

    def _assess_scene(self, pos: dict, yolo_classes: set):
        """
        Call LLaVA to understand the current scene.
        Saves observation + optionally a camera frame.
        """
        img = self._get_frame()
        if img is None:
            return

        try:
            raw = self._call_llava(self._patrol_prompt, img, num_predict=120)
            raw_clean = raw.replace("```json", "").replace("```", "").strip()
            s = raw_clean.find("{")
            e = raw_clean.rfind("}") + 1
            d = json.loads(raw_clean[s:e]) if s != -1 and e > 0 else None

            if d is None:
                return

            area_type   = d.get("area_type", "unknown")
            observation = d.get("observation", "")
            objects     = d.get("objects", [])
            interesting = d.get("interesting", False)

            # Update area counts
            self._area_counts[area_type] = self._area_counts.get(area_type, 0) + 1

            # Add objects to global set
            for obj in objects:
                self._all_objects.add(obj)

            # Build observation record
            obs = {
                "step":        self._step,
                "time":        round(time.time() - self._start_time, 1),
                "timestamp":   datetime.now().strftime("%H:%M:%S"),
                "x":           pos["x"] if pos else None,
                "y":           pos["y"] if pos else None,
                "heading":     pos["heading"] if pos else None,
                "area_type":   area_type,
                "objects":     objects,
                "yolo_classes": list(yolo_classes),
                "observation": observation,
                "interesting": interesting,
            }
            self._observations.append(obs)

            print(f"  [Auto] Step {self._step} | {area_type} | {observation[:60]}")
            if objects:
                print(f"  [Auto] Objects: {', '.join(objects)}")

            # Save frame if interesting
            if interesting and SAVE_FRAMES:
                now = time.time()
                if now - self._last_interesting > INTERESTING_COOLDOWN:
                    self._save_frame(img, self._step)
                    self._last_interesting = now

            # Log to session memory
            if self._mem:
                self._mem.log_detection(area_type, "center", "medium")

            # Auto-flush observations every 20 steps
            if self._step % 20 == 0:
                self._save_observations()
                self._save_path()

        except Exception as e:
            print(f"  [Auto] LLaVA assess error: {e}")

    # ── FILE I/O ────────────────────────────────────────────────────────────────

    def _create_map_dir(self) -> Path:
        """Create a new map folder with incremented ID."""
        maps_dir = self._models_dir / "map"
        maps_dir.mkdir(parents=True, exist_ok=True)

        existing = [d for d in maps_dir.iterdir() if d.is_dir() and d.name.startswith("map_")]
        nums = []
        for d in existing:
            parts = d.name.split("_")
            if len(parts) >= 2 and parts[1].isdigit():
                nums.append(int(parts[1]))

        next_num = max(nums) + 1 if nums else 1
        date_str = datetime.now().strftime("%Y-%m-%d")
        map_dir  = maps_dir / f"map_{next_num:03d}_{date_str}"
        map_dir.mkdir(parents=True, exist_ok=True)
        (map_dir / "frames").mkdir(exist_ok=True)
        return map_dir

    def _save_observations(self):
        if not self._map_dir or not self._observations:
            return
        try:
            path = self._map_dir / "observations.json"
            tmp  = path.with_suffix(".tmp")
            with open(tmp, "w", encoding="utf-8") as f:
                json.dump(self._observations, f, indent=2, ensure_ascii=False)
            tmp.replace(path)
        except Exception as e:
            print(f"  [Auto] Save observations error: {e}")

    def _save_path(self):
        if not self._map_dir or not self._path:
            return
        try:
            path = self._map_dir / "path.json"
            tmp  = path.with_suffix(".tmp")
            with open(tmp, "w", encoding="utf-8") as f:
                json.dump(self._path, f, indent=2)
            tmp.replace(path)
        except Exception as e:
            print(f"  [Auto] Save path error: {e}")

    def _save_frame(self, img_b64: str, step: int):
        """Save a camera frame as JPEG."""
        if not self._map_dir or not img_b64:
            return
        try:
            frame_path = self._map_dir / "frames" / f"frame_{step:04d}.jpg"
            with open(frame_path, "wb") as f:
                f.write(__import__("base64").b64decode(img_b64))
        except Exception as e:
            print(f"  [Auto] Save frame error: {e}")

    def _generate_summary(self) -> str:
        """Generate a text summary of the exploration session."""
        elapsed = time.time() - (self._start_time or time.time())
        mins = int(elapsed // 60)
        secs = int(elapsed % 60)

        lines = [
            f"Autonomous Exploration Summary",
            f"==============================",
            f"Map: {self._map_dir.name if self._map_dir else 'unknown'}",
            f"Date: {datetime.now().strftime('%Y-%m-%d %H:%M')}",
            f"Duration: {mins}m {secs}s",
            f"Steps taken: {self._step}",
            f"Observations: {len(self._observations)}",
            f"",
            f"Areas identified:",
        ]
        for area, count in sorted(self._area_counts.items(), key=lambda x: -x[1]):
            lines.append(f"  {area:<20} {count} observations")

        lines.append("")
        lines.append(f"Objects detected:")
        for obj in sorted(self._all_objects):
            lines.append(f"  - {obj}")

        # Add notable observations
        interesting = [o for o in self._observations if o.get("interesting")]
        if interesting:
            lines.append("")
            lines.append(f"Notable observations ({len(interesting)}):")
            for o in interesting[:10]:
                lines.append(f"  [{o['timestamp']}] {o['observation']}")

        return "\n".join(lines)

    def _save_session(self):
        """Save all data to disk."""
        self._save_observations()
        self._save_path()

        if self._map_dir:
            try:
                summary = self._generate_summary()
                with open(self._map_dir / "summary.txt", "w", encoding="utf-8") as f:
                    f.write(summary)
            except Exception as e:
                print(f"  [Auto] Save summary error: {e}")

    def _print_summary(self):
        """Print exploration summary to terminal."""
        elapsed = time.time() - (self._start_time or time.time())
        mins = int(elapsed // 60)
        secs = int(elapsed % 60)

        print(f"\n  [Auto] Exploration complete")
        print(f"  [Auto] Duration: {mins}m {secs}s | Steps: {self._step}")
        print(f"  [Auto] Observations: {len(self._observations)}")

        if self._area_counts:
            print(f"  [Auto] Areas: {dict(sorted(self._area_counts.items()))}")

        if self._all_objects:
            print(f"  [Auto] Objects: {', '.join(sorted(self._all_objects))}")

        if self._map_dir:
            print(f"  [Auto] Saved to: {self._map_dir}\n")