session_manager.py

import os
import shutil
import uuid
import time
from datetime import datetime
import logging
from typing import Optional, Dict, Any
import glob
import re
import numpy as np

logger = logging.getLogger(__name__)

class WorldState:
    """
    Structured situational memory (The sentinel Mirror).
    Tracks persistent entities, environmental conditions, and mission facts.
    """
    def __init__(self, session_id: str):
        self.session_id = session_id
        self.last_update = time.time()
        
        # ── The Entity Mirror (The core of reality) ──
        # List of Entities. Each entity is a dict:
        # {
        #   "id": int, 
        #   "type": str, 
        #   "attributes": dict, 
        #   "state": str, 
        #   "last_seen": float, 
        #   "confidence": float, 
        #   "history": List[Dict] # Temporal trail of [state, attributes, timestamp]
        # }
        self.entities = [] 
        self.event_history = [] # Global timeline of significant semantic changes
        
        # ── Environmental Context ──
        self.lighting = "unknown"
        self.noise_level = 0.0
        self.spatial_context = "unknown" # e.g., 'corridor', 'outside'
        
        # ── Mission Logic ──
        self.target_status = "searching" # searching, visible, occluded, lost
        self.active_strategies = []
        self.diagnostics = {}
        self.confidence_history = []

    def log_event(self, text: str, event_type: str = "semantic"):
        """Appends a timestamped event to the global history."""
        self.event_history.append({
            "timestamp": time.time(),
            "type": event_type,
            "text": text
        })
        # Increased capacity for long research sessions
        if len(self.event_history) > 200:
            self.event_history.pop(0)

    def get_events_near(self, target_ts: float, count: int = 5) -> list:
        """Returns the `count` closest events to a target timestamp."""
        if not self.event_history:
            return []
        # Sort by distance to target
        scored = [(abs(e["timestamp"] - target_ts), e) for e in self.event_history]
        scored.sort(key=lambda x: x[0])
        return [e for _, e in scored[:count]]

    def get_timeline(self, window_seconds: int = 60, anchor_ts: float = None) -> str:
        """Returns a formatted string of events. 
        Uses anchor_ts (last event time) instead of wall clock to survive restarts.
        """
        # Use the last event's timestamp as anchor, NOT time.time()
        if anchor_ts is None:
            if self.event_history:
                anchor_ts = self.event_history[-1]["timestamp"]
            else:
                anchor_ts = time.time()
        
        start_ts = anchor_ts - window_seconds
        
        # 1. SCENE CONTEXT
        now = time.time()
        active_entities = [e for e in self.entities if now - e.get("last_seen", 0) < 30]
        context_lines = []
        if active_entities:
            type_counts = {}
            for e in active_entities:
                t = e["type"]
                type_counts[t] = type_counts.get(t, 0) + 1
            summary_parts = [f"{count} {t}" for t, count in type_counts.items()]
            context_lines.append(f"[CURRENT SCENE CONTEXT]")
            context_lines.append(f"Currently in view: {', '.join(summary_parts)}")
            context_lines.append("")

        # 2. RECENT EVENTS (anchored to last observation, not wall clock)
        relevant = [e for e in self.event_history if e["timestamp"] >= start_ts]
        event_lines = []
        for e in relevant:
            dt = datetime.fromtimestamp(e["timestamp"]).strftime("%H:%M:%S")
            # Add human-readable time-ago
            ago = int(anchor_ts - e["timestamp"])
            ago_str = f"{ago}s ago" if ago < 60 else f"{ago // 60}m {ago % 60}s ago"
            event_lines.append(f"[{dt}] ({ago_str}) {e['text']}")
            
        if not event_lines:
            event_lines.append(f"No observations recorded in the last {window_seconds}s.")

        output = context_lines + ["[RECENT LOGGED EVENTS]"] + event_lines
        return "\n".join(output)

    def to_dict(self) -> Dict[str, Any]:
        return self.__dict__

    @classmethod
    def from_dict(cls, data: Dict[str, Any]):
        ws = cls(data.get("session_id", "unknown"))
        ws.__dict__.update(data)
        return ws

# Build absolute path to sessions folder relative to this file
_BACKEND_ROOT = os.path.dirname(os.path.abspath(__file__))
DEFAULT_SESSIONS_DIR = os.path.join(_BACKEND_ROOT, "sessions")

class SessionManager:
    def __init__(self, base_dir=DEFAULT_SESSIONS_DIR, max_sessions=50, ttl_seconds=86400):
        self.base_dir = base_dir
        self.max_sessions = max_sessions
        self.ttl_seconds = ttl_seconds
        self.active_session_id: Optional[str] = None  # Tracks user's chosen session
        os.makedirs(self.base_dir, exist_ok=True)
        self.cleanup()

    def set_active_session(self, session_id: str):
        """Set the user's chosen active session. Called from the frontend dropdown."""
        if os.path.exists(self.get_session_path(session_id)):
            self.active_session_id = session_id
            logger.info(f"[SESSION] Active session set to: {session_id}")
            return True
        return False

    def get_or_create_session(self, session_id: Optional[str] = None) -> str:
        """Returns existing session or creates a new one.
        Priority: explicit session_id > active_session_id > create new.
        """
        # Sanitize common string representations of null from frontend
        if isinstance(session_id, str):
            if session_id.lower().strip() in ["", "none", "null", "undefined"]:
                session_id = None

        # 1. Use explicit session_id if it exists on disk
        if session_id and os.path.exists(self.get_session_path(session_id)):
            return session_id

        # 2. Use the user's chosen active session (set via dropdown)
        if not session_id and self.active_session_id:
            if os.path.exists(self.get_session_path(self.active_session_id)):
                return self.active_session_id

        # 3. Create a new session as last resort
        new_id = self.create_session(session_id)
        self.active_session_id = new_id  # Auto-activate newly created sessions
        return new_id

    def create_session(self, session_id: Optional[str] = None) -> str:
        """Creates a new session directory and returns the session_id."""
        if not session_id:
            session_id = str(uuid.uuid4())
            
        session_path = os.path.join(self.base_dir, session_id)
        os.makedirs(session_path, exist_ok=True)
        
        # Create subdirs for media types
        os.makedirs(os.path.join(session_path, "frames"), exist_ok=True)
        os.makedirs(os.path.join(session_path, "audio"), exist_ok=True)
        
        # Initialize world state
        self.update_world_state(session_id, WorldState(session_id))
        
        # Touch metadata file to track time
        with open(os.path.join(session_path, "metadata.txt"), "w") as f:
            f.write(f"Created: {datetime.now().isoformat()}")
            
        self.cleanup() # Run cleanup on every new session creation
        return session_id

    def update_world_state(self, session_id: str, state: WorldState):
        """Persists structured state into session storage."""
        session_path = self.get_session_path(session_id)
        if not os.path.exists(session_path): return
        
        state_path = os.path.join(session_path, "world_state.json")
        import json
        with open(state_path, "w") as f:
            json.dump(state.to_dict(), f)

    def get_world_state(self, session_id: str) -> WorldState:
        """Retrieves the latest world state."""
        session_path = self.get_session_path(session_id)
        state_path = os.path.join(session_path, "world_state.json")
        if os.path.exists(state_path):
            import json
            try:
                with open(state_path, "r") as f:
                    data = json.load(f)
                    return WorldState.from_dict(data)
            except: pass
        return WorldState(session_id)

    def update_sensory_memory(self, session_id: str, new_data: Dict[str, Any]):
        """DEPRECATED: Use update_world_state. Keeps compatibility for now."""
        ws = self.get_world_state(session_id)
        ws.diagnostics.update(new_data)
        self.update_world_state(session_id, ws)

    def get_sensory_memory(self, session_id: str) -> Dict[str, Any]:
        """DEPRECATED: Use get_world_state. Returns diagnostics."""
        ws = self.get_world_state(session_id)
        return ws.diagnostics

    def get_session_path(self, session_id: str) -> str:
        return os.path.join(self.base_dir, session_id)

    def save_frame(self, session_id: str, frame_data, timestamp: float):
        """Saves a frame to the session directory (Throttled by Orchestrator)."""
        session_path = self.get_session_path(session_id)
        if not os.path.exists(session_path):
            return False
            
        frames_dir = os.path.join(session_path, "frames")
        os.makedirs(frames_dir, exist_ok=True)
        
        frame_name = f"frame_{int(timestamp*1000)}.jpg"
        frame_path = os.path.join(frames_dir, frame_name)
        
        try:
            # Assuming frame_data is a numpy array (OpenCV) or PIL image
            if hasattr(frame_data, "save"): # PIL
                frame_data.save(frame_path)
            elif isinstance(frame_data, np.ndarray):
                import cv2
                # OpenCV handles numpy arrays directly
                success = cv2.imwrite(frame_path, frame_data)
                if not success:
                    logger.warning(f"[SESSION] cv2.imwrite failed for {frame_path}")
                    return False
            else:
                logger.warning(f"[SESSION] Unsupported frame data type: {type(frame_data)}")
                return False
            return frame_path
        except Exception as e:
            logger.error(f"[SESSION] Error saving frame: {e}")
            return False

    def get_historic_frame(self, session_id: str, timestamp: float) -> Optional[str]:
        """
        Finds the frame closest to the requested timestamp.
        """
        session_path = self.get_session_path(session_id)
        frames_dir = os.path.join(session_path, "frames")
        if not os.path.exists(frames_dir): return None
        
        frames = glob.glob(os.path.join(frames_dir, "frame_*.jpg"))
        if not frames: return None
        
        # Convert filenames back to timestamps for comparison
        target_ms = int(timestamp * 1000)
        closest_frame = min(frames, key=lambda f: abs(int(re.search(r'frame_(\d+)', f).group(1)) - target_ms))
        
        # Only return if within a reasonable window (e.g. 30 seconds)
        found_ms = int(re.search(r'frame_(\d+)', closest_frame).group(1))
        if abs(found_ms - target_ms) > 30000:
            logger.warning(f"[SESSION] No frame found within 30s window of {timestamp}")
            return None
            
        return closest_frame

    def log_session_event(self, session_id: str, text: str, event_type: str = "semantic"):
        """Adds an event to a session's world state and persists it."""
        ws = self.get_world_state(session_id)
        ws.log_event(text, event_type)
        self.update_world_state(session_id, ws)

    def get_session_timeline(self, session_id: str, window_seconds: int = 60) -> str:
        """Retrieves formatted timeline from world state."""
        ws = self.get_world_state(session_id)
        return ws.get_timeline(window_seconds)

    def get_historic_audio_context(self, session_id: str, timestamp: float, window: int = 15) -> str:
        """
        [STUB] Retrieves audio transcripts or events from memory within a window.
        """
        # For now, we'll return the general mission finding context if relevant
        return ""

    def save_audio_segment(self, session_id: str, audio_data, start_time: float):
        """Saves an audio segment to the session directory."""
        session_path = self.get_session_path(session_id)
        if not os.path.exists(session_path):
            return False
            
        audio_name = f"audio_{int(start_time*1000)}.wav"
        audio_path = os.path.join(session_path, "audio", audio_name)
        # logic to save audio file would go here
        return audio_path

    def cleanup(self):
        """Deletes old sessions based on TTL and count limit."""
        try:
            now = time.time()
            sessions = []
            for d in os.listdir(self.base_dir):
                p = os.path.join(self.base_dir, d)
                if os.path.isdir(p):
                    mtime = os.path.getmtime(p)
                    sessions.append((d, p, mtime))
            
            # Sort by time (oldest first)
            sessions.sort(key=lambda x: x[2])
            
            # 1. Cleanup by TTL
            for sid, path, mtime in sessions:
                if now - mtime > self.ttl_seconds:
                    logger.info(f"Cleaning up expired session: {sid}")
                    shutil.rmtree(path)
                    sessions.remove((sid, path, mtime))
            
            # 2. Cleanup by Count Limit
            while len(sessions) > self.max_sessions:
                sid, path, mtime = sessions.pop(0)
                logger.info(f"Cleaning up session (limit reached): {sid}")
                shutil.rmtree(path)
                
        except Exception as e:
            logger.error(f"Error during SessionManager cleanup: {e}")

# Global instance
session_manager = SessionManager()