import json
import os
import time
from typing import Dict, Any, List

LOG_DIR = os.path.join(os.path.dirname(__file__), "logs")
NOTEBOOK_FILE = os.path.join(LOG_DIR, "notebook.json")

class LabNotebook:
    def __init__(self):
        self._ensure_log_dir()
        
    def _ensure_log_dir(self):
        if not os.path.exists(LOG_DIR):
            os.makedirs(LOG_DIR)
        if not os.path.exists(NOTEBOOK_FILE):
            with open(NOTEBOOK_FILE, "w") as f:
                json.dump([], f)

    def log_experiment(self, experiment_data: Dict[str, Any]) -> Dict[str, Any]:
        """
        Logs an experiment entry to the persistent notebook.
        """
        entry = {
            "id": f"EXP-{int(time.time()*1000)}",
            "timestamp": time.strftime("%Y-%m-%d %H:%M:%S"),
            "data": experiment_data
        }
        
        try:
            with open(NOTEBOOK_FILE, "r") as f:
                logs = json.load(f)
        except (json.JSONDecodeError, FileNotFoundError):
            logs = []
            
        logs.append(entry)
        
        # Keep only last 100 entries to prevent infinite growth in this demo
        if len(logs) > 100:
            logs = logs[-100:]
            
        with open(NOTEBOOK_FILE, "w") as f:
            json.dump(logs, f, indent=2)
            
        return entry

    def get_logs(self) -> List[Dict[str, Any]]:
        try:
            with open(NOTEBOOK_FILE, "r") as f:
                return json.load(f)
        except:
            return []

lab_notebook = LabNotebook()