code/server/session_manager.py

"""Session management and conversation export utilities.

New module providing:
- Session persistence
- Conversation history management
- Export to multiple formats (JSON, Markdown, HTML)
- Import capabilities
- Template management
"""

from __future__ import annotations

import json
import logging
import os
import time
from dataclasses import dataclass, field, asdict
from datetime import datetime
from pathlib import Path
from typing import Any, Optional

logger = logging.getLogger(__name__)


@dataclass
class ChatMessage:
    """Represents a single chat message."""
    role: str  # "user" or "assistant"
    content: str
    timestamp: float = field(default_factory=time.time)
    token_count: int = 0
    model_used: str = ""
    
    def to_dict(self) -> dict[str, Any]:
        return asdict(self)
    
    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> 'ChatMessage':
        return cls(**data)


@dataclass 
class Session:
    """Represents a user session with full conversation history."""
    session_id: str
    created_at: float = field(default_factory=time.time)
    updated_at: float = field(default_factory=time.time)
    messages: list[ChatMessage] = field(default_factory=list)
    metadata: dict[str, Any] = field(default_factory=dict)
    title: str = ""
    
    def add_message(self, role: str, content: str, **kwargs) -> ChatMessage:
        """Add a message to the session."""
        msg = ChatMessage(role=role, content=content, **kwargs)
        self.messages.append(msg)
        self.updated_at = time.time()
        
        # Auto-generate title from first user message
        if role == "user" and not self.title:
            self.title = content[:50] + "..." if len(content) > 50 else content
            
        return msg
    
    def get_history(self) -> list[dict[str, str]]:
        """Get messages in the format expected by chat_helpers."""
        return [{"role": m.role, "content": m.content} for m in self.messages]
    
    def to_dict(self) -> dict[str, Any]:
        return {
            "session_id": self.session_id,
            "created_at": self.created_at,
            "updated_at": self.updated_at,
            "title": self.title,
            "messages": [m.to_dict() for m in self.messages],
            "metadata": self.metadata,
        }
    
    @classmethod
    def from_dict(cls, data: dict[str, Any]) -> 'Session':
        session = cls(
            session_id=data["session_id"],
            created_at=data.get("created_at", time.time()),
            updated_at=data.get("updated_at", time.time()),
            title=data.get("title", ""),
            metadata=data.get("metadata", {}),
        )
        session.messages = [ChatMessage.from_dict(m) for m in data.get("messages", [])]
        return session


class SessionManager:
    """Manages multiple sessions with persistence support.
    
    Features:
    - In-memory session storage
    - JSON file persistence
    - Session search and filtering
    - Automatic cleanup of old sessions
    """
    
    def __init__(self, storage_dir: str | None = None):
        self._sessions: dict[str, Session] = {}
        self._storage_dir = Path(storage_dir) if storage_dir else None
        
        if self._storage_dir:
            self._storage_dir.mkdir(parents=True, exist_ok=True)
            self._load_sessions()
    
    def create_session(self, session_id: str | None = None, **metadata) -> Session:
        """Create a new session."""
        if session_id is None:
            session_id = f"session_{int(time.time() * 1000)}"
            
        session = Session(session_id=session_id, metadata=metadata)
        self._sessions[session_id] = session
        logger.info("Created session: %s", session_id)
        return session
    
    def get_session(self, session_id: str) -> Optional[Session]:
        """Get a session by ID."""
        return self._sessions.get(session_id)
    
    def update_session(self, session_id: str, **updates) -> bool:
        """Update session metadata."""
        session = self.get_session(session_id)
        if not session:
            return False
            
        for key, value in updates.items():
            if hasattr(session, key):
                setattr(session, key, value)
            elif key in session.metadata:
                session.metadata[key] = value
            else:
                session.metadata[key] = value
                
        session.updated_at = time.time()
        return True
    
    def delete_session(self, session_id: str) -> bool:
        """Delete a session."""
        if session_id in self._sessions:
            del self._sessions[session_id]
            
            # Delete persisted file
            if self._storage_dir:
                file_path = self._storage_dir / f"{session_id}.json"
                if file_path.exists():
                    file_path.unlink()
                    
            logger.info("Deleted session: %s", session_id)
            return True
        return False
    
    def list_sessions(
        self, 
        limit: int = 50,
        filter_text: str | None = None,
    ) -> list[dict[str, Any]]:
        """List sessions with optional filtering."""
        sessions = list(self._sessions.values())
        
        # Sort by updated_at descending
        sessions.sort(key=lambda s: s.updated_at, reverse=True)
        
        # Filter by text (search in title and messages)
        if filter_text:
            filter_lower = filter_text.lower()
            sessions = [
                s for s in sessions
                if filter_lower in s.title.lower() or
                   any(filter_lower in m.content.lower() for m in s.messages[:10])
            ]
        
        # Apply limit
        sessions = sessions[:limit]
        
        return [
            {
                "session_id": s.session_id,
                "title": s.title,
                "created_at": s.created_at,
                "updated_at": s.updated_at,
                "message_count": len(s.messages),
            }
            for s in sessions
        ]
    
    def save_session(self, session_id: str) -> bool:
        """Persist a session to disk."""
        if not self._storage_dir:
            return False
            
        session = self.get_session(session_id)
        if not session:
            return False
            
        try:
            file_path = self._storage_dir / f"{session_id}.json"
            with open(file_path, 'w', encoding='utf-8') as f:
                json.dump(session.to_dict(), f, indent=2, ensure_ascii=False)
            return True
        except Exception as e:
            logger.error("Failed to save session %s: %s", session_id, e)
            return False
    
    def load_session(self, session_id: str) -> Optional[Session]:
        """Load a session from disk."""
        if not self._storage_dir:
            return None
            
        try:
            file_path = self._storage_dir / f"{session_id}.json"
            if not file_path.exists():
                return None
                
            with open(file_path, 'r', encoding='utf-8') as f:
                data = json.load(f)
                
            session = Session.from_dict(data)
            self._sessions[session_id] = session
            return session
        except Exception as e:
            logger.error("Failed to load session %s: %s", session_id, e)
            return None
    
    def _load_sessions(self) -> None:
        """Load all sessions from storage directory."""
        if not self._storage_dir:
            return
            
        try:
            for file_path in self._storage_dir.glob("*.json"):
                try:
                    with open(file_path, 'r', encoding='utf-8') as f:
                        data = json.load(f)
                    session = Session.from_dict(data)
                    self._sessions[session.session_id] = session
                except Exception as e:
                    logger.warning("Failed to load session file %s: %s", file_path, e)
        except Exception as e:
            logger.error("Failed to scan sessions directory: %s", e)
    
    def save_all(self) -> int:
        """Save all modified sessions. Returns count saved."""
        count = 0
        for session_id in list(self._sessions.keys()):
            if self.save_session(session_id):
                count += 1
        return count
    
    def cleanup_old_sessions(self, max_age_days: int = 30) -> int:
        """Remove sessions older than max_age_days. Returns count removed."""
        cutoff = time.time() - (max_age_days * 24 * 60 * 60)
        to_remove = [
            sid for sid, s in self._sessions.items()
            if s.updated_at < cutoff
        ]
        
        for sid in to_remove:
            self.delete_session(sid)
            
        logger.info("Cleaned up %d old sessions", len(to_remove))
        return len(to_remove)


# ─── Export Functions ────────────────────────────────────────────────────

def export_to_json(
    session: Session,
    output_path: str,
) -> str:
    """Export session to JSON format.
    
    Args:
        session: The session to export.
        output_path: Path to write the JSON file.
        
    Returns:
        The output path.
    """
    data = session.to_dict()
    
    # Add export metadata
    data["exported_at"] = datetime.now().isoformat()
    data["version"] = "2.1"
    
    with open(output_path, 'w', encoding='utf-8') as f:
        json.dump(data, f, indent=2, ensure_ascii=False)
        
    logger.info("Exported session to JSON: %s", output_path)
    return output_path


def export_to_markdown(
    session: Session,
    output_path: str,
) -> str:
    """Export session to Markdown format.
    
    Args:
        session: The session to export.
        output_path: Path to write the markdown file.
        
    Returns:
        The output path.
    """
    lines = [
        f"# {session.title or 'SoniCoder Conversation'}",
        "",
        f"**Session ID:** `{session.session_id}`",
        f"**Created:** {datetime.fromtimestamp(session.created_at).strftime('%Y-%m-%d %H:%M:%S')}",
        f"**Messages:** {len(session.messages)}",
        "",
        "---",
        "",
    ]
    
    for i, msg in enumerate(session.messages, 1):
        timestamp = datetime.fromtimestamp(msg.timestamp).strftime('%H:%M:%S')
        role_label = "**User**" if msg.role == "user" else "**Assistant**"
        
        lines.append(f"## Message {i} ({role_label}) - {timestamp}")
        lines.append("")
        lines.append(msg.content)
        lines.append("")
        lines.append("---")
        lines.append("")
    
    with open(output_path, 'w', encoding='utf-8') as f:
        f.write('\n'.join(lines))
        
    logger.info("Exported session to Markdown: %s", output_path)
    return output_path


def export_to_html(
    session: Session,
    output_path: str,
) -> str:
    """Export session to HTML format with styling.
    
    Args:
        session: The session to export.
        output_path: Path to write the HTML file.
        
    Returns:
        The output path.
    """
    messages_html = []
    
    for msg in session.messages:
        timestamp = datetime.fromtimestamp(msg.timestamp).strftime('%H:%M:%S')
        role_class = "user-message" if msg.role == "user" else "assistant-message"
        role_label = "👤 User" if msg.role == "user" else "🤖 Assistant"
        
        # Escape HTML special characters in content
        content = (
            msg.content
            .replace('&', '&amp;')
            .replace('<', '&lt;')
            .replace('>', '&gt;')
        )
        
        # Convert code blocks to HTML
        import re
        content = re.sub(
            r'```(\w*)\n(.*?)```',
            r'<pre><code class="language-\1">\2</code></pre>',
            content,
            flags=re.DOTALL
        )
        
        messages_html.append(f"""
        <div class="message {role_class}">
            <div class="message-header">
                <span class="role-label">{role_label}</span>
                <span class="timestamp">{timestamp}</span>
            </div>
            <div class="message-content">{content}</div>
        </div>
        """)
    
    html = f"""<!DOCTYPE html>
<html lang="en">
<head>
    <meta charset="UTF-8">
    <meta name="viewport" content="width=device-width, initial-scale=1.0">
    <title>{session.title or 'SoniCoder Conversation'}</title>
    <style>
        body {{
            font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif;
            line-height: 1.6;
            color: #333;
            max-width: 900px;
            margin: 0 auto;
            padding: 20px;
            background: #f5f5f5;
        }}
        h1 {{
            color: #2c3e50;
            border-bottom: 2px solid #3498db;
            padding-bottom: 10px;
        }}
        .meta {{
            color: #666;
            font-size: 0.9em;
        }}
        .message {{
            margin: 20px 0;
            padding: 15px;
            border-radius: 8px;
            box-shadow: 0 2px 4px rgba(0,0,0,0.1);
        }}
        .user-message {{
            background: #e3f2fd;
            border-left: 4px solid #2196f3;
        }}
        .assistant-message {{
            background: #f3e5f5;
            border-left: 4px solid #9c27b0;
        }}
        .message-header {{
            display: flex;
            justify-content: space-between;
            margin-bottom: 10px;
            font-weight: bold;
        }}
        .timestamp {{
            color: #888;
            font-size: 0.85em;
        }}
        pre {{
            background: #263238;
            color: #aed581;
            padding: 15px;
            border-radius: 4px;
            overflow-x: auto;
            font-size: 0.9em;
        }}
        code {{
            font-family: 'Fira Code', 'Consolas', monospace;
        }}
    </style>
</head>
<body>
    <h1>{session.title or 'SoniCoder Conversation'}</h1>
    <div class="meta">
        <p><strong>Session:</strong> {session.session_id}</p>
        <p><strong>Created:</strong> {datetime.fromtimestamp(session.created_at).strftime('%Y-%m-%d %H:%M:%S')}</p>
        <p><strong>Messages:</strong> {len(session.messages)}</p>
    </div>
    <hr>
    {''.join(messages_html)}
</body>
</html>"""
    
    with open(output_path, 'w', encoding='utf-8') as f:
        f.write(html)
        
    logger.info("Exported session to HTML: %s", output_path)
    return output_path


def export_session(
    session: Session,
    format: str = "json",
    output_dir: str = ".",
) -> str:
    """Export a session in the specified format.
    
    Args:
        session: The session to export.
        format: Export format ('json', 'markdown', 'html').
        output_dir: Directory to save the exported file.
        
    Returns:
        Path to the exported file.
    """
    # Ensure output directory exists
    os.makedirs(output_dir, exist_ok=True)
    
    safe_title = "".join(
        c for c in (session.title or "conversation")
        if c.isalnum() or c in (' ', '-', '_')
    ).strip()[:50] or "conversation"
    
    timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
    
    exporters = {
        "json": export_to_json,
        "markdown": export_to_markdown,
        "html": export_to_html,
    }
    
    if format not in exporters:
        raise ValueError(f"Unsupported format: {format}. Supported: {list(exporters.keys())}")
    
    filename = f"{safe_title}_{timestamp}.{format}"
    output_path = os.path.join(output_dir, filename)
    
    return exporters[format](session, output_path)


# Global session manager instance
_global_manager: Optional[SessionManager] = None


def get_session_manager(storage_dir: str | None = None) -> SessionManager:
    """Get or create the global session manager."""
    global _global_manager
    if _global_manager is None:
        _global_manager = SessionManager(storage_dir=storage_dir)
    return _global_manager