import json
import os
import uuid
from datetime import datetime
from pathlib import Path
from typing import Dict, List, Optional, Any


class TrajectoryStore:
    """Store and retrieve agent trajectories"""

    def __init__(self, storage_dir: str = "./trajectories"):
        self.storage_dir = Path(storage_dir)
        self.storage_dir.mkdir(parents=True, exist_ok=True)

    def save_trajectory(self, episode_data: Dict[str, Any]) -> str:
        """Save a trajectory to storage"""
        # Generate unique ID if not present
        if "id" not in episode_data:
            episode_data["id"] = str(uuid.uuid4())

        # Add timestamp if not present
        if "timestamp" not in episode_data:
            episode_data["timestamp"] = datetime.now().isoformat()

        # Create filename
        filename = f"{episode_data['id']}.json"
        filepath = self.storage_dir / filename

        # Save to file
        with open(filepath, "w") as f:
            json.dump(episode_data, f, indent=2)

        return episode_data["id"]

    def load_trajectory(self, trajectory_id: str) -> Optional[Dict[str, Any]]:
        """Load a trajectory by ID"""
        filepath = self.storage_dir / f"{trajectory_id}.json"

        if not filepath.exists():
            return None

        try:
            with open(filepath, "r") as f:
                return json.load(f)
        except Exception:
            return None

    def list_trajectories(self, limit: int = 100) -> List[Dict[str, Any]]:
        """List recent trajectories"""
        trajectories = []

        # Get all JSON files in storage directory
        files = list(self.storage_dir.glob("*.json"))

        # Sort by modification time (newest first)
        files.sort(key=lambda x: x.stat().st_mtime, reverse=True)

        # Load up to limit trajectories
        for filepath in files[:limit]:
            try:
                with open(filepath, "r") as f:
                    data = json.load(f)
                    trajectories.append(data)
            except Exception:
                # Skip corrupted files
                continue

        return trajectories

    def delete_trajectory(self, trajectory_id: str) -> bool:
        """Delete a trajectory by ID"""
        filepath = self.storage_dir / f"{trajectory_id}.json"

        if filepath.exists():
            filepath.unlink()
            return True
        return False

    def get_storage_stats(self) -> Dict[str, Any]:
        """Get statistics about stored trajectories"""
        files = list(self.storage_dir.glob("*.json"))

        total_reward = 0.0
        count = 0

        for filepath in files:
            try:
                with open(filepath, "r") as f:
                    data = json.load(f)
                    total_reward += data.get("reward", 0.0)
                    count += 1
            except Exception:
                continue

        return {
            "total_trajectories": len(files),
            "average_reward": total_reward / count if count > 0 else 0.0,
            "storage_path": str(self.storage_dir),
        }
