"""
episodic_store.py

Lazy loader for all_sessions.json. Keeps the full raw-turn data in memory
and returns sessions on demand by session ID, avoiding loading all sessions
into ChatHistory upfront.
"""

import json
from datetime import datetime
from typing import Dict, List, Optional


class EpisodicMemoryStore:
    def __init__(self, all_sessions_path: str):
        print(f"[EpisodicMemoryStore] Loading {all_sessions_path} ...")
        with open(all_sessions_path) as f:
            self._data: Dict[str, List] = json.load(f)
        print(f"[EpisodicMemoryStore] Loaded {len(self._data)} sessions.")

    @staticmethod
    def _parse_date(date_str: str) -> datetime:
        """Convert '2023/04/10 (Mon) 17:50' to datetime."""
        date_part = date_str.split('(')[0].strip()
        time_part = date_str.split(')')[-1].strip()
        return datetime.strptime(date_part + time_part, "%Y/%m/%d%H:%M")

    def get_raw_sessions(
        self,
        sess_ids: List[str],
        date_lookup: Dict[str, str],
        topic_lookup: Optional[Dict[str, List[str]]] = None,
    ) -> List[dict]:
        """
        Return a list of session dicts compatible with ChatHistory(sessions=...).

        Args:
            sess_ids:     Ordered list of session IDs to load.
            date_lookup:  Mapping sess_id -> date string (e.g. '2023/04/10 (Mon) 17:50').
            topic_lookup: Optional mapping sess_id -> list of topic strings.

        Returns:
            List of session dicts, each with keys:
              session_id, session_date, session, topic, timestamp
        """
        sessions = []
        for sid in sess_ids:
            if sid not in self._data:
                continue
            date_str = date_lookup.get(sid, "")
            try:
                ts = self._parse_date(date_str) if date_str else datetime.min
            except Exception:
                ts = datetime.min
            sessions.append({
                "session_id": sid,
                "session_date": date_str,
                "session": self._data[sid],
                "topic": topic_lookup.get(sid, []) if topic_lookup else [],
                "timestamp": ts,
            })
        return sessions