"""
Session Manager - Save and load analysis sessions
"""

import json
import pickle
import os
from datetime import datetime
import pandas as pd

class SessionManager:
    def __init__(self, session_dir="saved_sessions"):
        self.session_dir = session_dir
        os.makedirs(session_dir, exist_ok=True)
    
    def save_session(self, df, schema, name=None):
        """Save current session"""
        if name is None:
            name = f"session_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
        
        session_data = {
            'name': name,
            'timestamp': datetime.now().isoformat(),
            'data': df.to_dict('records'),
            'columns': list(df.columns),
            'dtypes': df.dtypes.astype(str).to_dict(),
            'schema': schema,
            'shape': df.shape
        }
        
        filepath = os.path.join(self.session_dir, f"{name}.pkl")
        with open(filepath, 'wb') as f:
            pickle.dump(session_data, f)
        
        return name, filepath
    
    def load_session(self, name):
        """Load saved session"""
        filepath = os.path.join(self.session_dir, name)
        if not os.path.exists(filepath):
            # Try with .pkl extension
            filepath = f"{filepath}.pkl"
            if not os.path.exists(filepath):
                return None
        
        with open(filepath, 'rb') as f:
            session_data = pickle.load(f)
        
        # Reconstruct DataFrame
        df = pd.DataFrame(session_data['data'])
        
        return {
            'df': df,
            'schema': session_data['schema'],
            'name': session_data['name'],
            'timestamp': session_data['timestamp']
        }
    
    def list_sessions(self):
        """List all saved sessions"""
        sessions = []
        for file in os.listdir(self.session_dir):
            if file.endswith('.pkl'):
                filepath = os.path.join(self.session_dir, file)
                with open(filepath, 'rb') as f:
                    data = pickle.load(f)
                sessions.append({
                    'name': data['name'],
                    'timestamp': data['timestamp'],
                    'rows': data['shape'][0],
                    'columns': data['shape'][1],
                    'file': file
                })
        return sorted(sessions, key=lambda x: x['timestamp'], reverse=True)
    
    def delete_session(self, name):
        """Delete a saved session"""
        filepath = os.path.join(self.session_dir, name)
        if os.path.exists(filepath):
            os.remove(filepath)
            return True
        return False