"""Session manager implementation"""

import json
import os
import uuid
from typing import List, Optional, Dict, Any
from diona.ai.client.simplecli.models import Message


class Session:
    """Session data model"""
    
    def __init__(self, session_id: str, name: str):
        """Initialize a session
        
        Args:
            session_id: Session ID
            name: Session name
        """
        self.session_id = session_id
        self.name = name
        self.messages: List[Message] = []
        self.created_at = os.path.getmtime(__file__)  # Simplified timestamp
    
    def add_message(self, message: Message):
        """Add a message to the session"""
        self.messages.append(message)
    
    def clear_messages(self):
        """Clear all messages"""
        self.messages = []
    
    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary"""
        return {
            "session_id": self.session_id,
            "name": self.name,
            "messages": [msg.to_dict() for msg in self.messages],
            "created_at": self.created_at
        }
    
    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> 'Session':
        """Create session from dictionary"""
        session = cls(data['session_id'], data['name'])
        session.messages = [Message(msg['role'], msg['content']) for msg in data['messages']]
        session.created_at = data.get('created_at', os.path.getmtime(__file__))
        return session


class SessionManager:
    """Session manager"""
    
    def __init__(self):
        """Initialize session manager"""
        self.sessions: Dict[str, Session] = {}
        self.current_session_id: Optional[str] = None
        self.session_dir = os.path.join(os.path.expanduser("~"), ".diona", "sessions")
        os.makedirs(self.session_dir, exist_ok=True)
        self.load_sessions()
        
        # Create default session if no sessions exist
        if not self.sessions:
            self.create_session("Default")
        # Set current session to first session if not set
        elif not self.current_session_id:
            self.current_session_id = next(iter(self.sessions.keys()))
    
    def create_session(self, name: str) -> Session:
        """Create a new session
        
        Args:
            name: Session name
            
        Returns:
            Created session
        """
        session_id = str(uuid.uuid4())
        session = Session(session_id, name)
        self.sessions[session_id] = session
        self.current_session_id = session_id
        self.save_sessions()
        return session
    
    def switch_session(self, session_id: str) -> Optional[Session]:
        """Switch to a different session
        
        Args:
            session_id: Session ID
            
        Returns:
            Switched session or None if not found
        """
        if session_id in self.sessions:
            self.current_session_id = session_id
            return self.sessions[session_id]
        return None
    
    def add_user_message(self, content: str):
        """Add a user message to the current session"""
        if self.current_session_id:
            session = self.sessions[self.current_session_id]
            session.add_message(Message("user", content))
            self.save_sessions()
    
    def add_assistant_message(self, content: str):
        """Add an assistant message to the current session"""
        if self.current_session_id:
            session = self.sessions[self.current_session_id]
            session.add_message(Message("assistant", content))
            self.save_sessions()
    
    def get_context(self, max_tokens: int = 4096) -> List[Message]:
        """Get context with token limit
        
        Args:
            max_tokens: Maximum tokens for context
            
        Returns:
            List of messages within token limit
        """
        if not self.current_session_id:
            return []
        
        session = self.sessions[self.current_session_id]
        messages = session.messages
        
        # Simple token estimation (1 token ~ 4 chars)
        total_tokens = 0
        context_messages = []
        
        # Start from the end and add messages until token limit
        for msg in reversed(messages):
            msg_tokens = len(msg.content) // 4
            if total_tokens + msg_tokens <= max_tokens:
                context_messages.insert(0, msg)
                total_tokens += msg_tokens
            else:
                break
        
        return context_messages
    
    def clear_history(self):
        """Clear current session history"""
        if self.current_session_id:
            session = self.sessions[self.current_session_id]
            session.clear_messages()
            self.save_sessions()
    
    def get_token_count(self) -> int:
        """Estimate current context token count
        
        Returns:
            Estimated token count
        """
        if not self.current_session_id:
            return 0
        
        session = self.sessions[self.current_session_id]
        total_tokens = 0
        for msg in session.messages:
            total_tokens += len(msg.content) // 4
        return total_tokens
    
    def export_session(self, format: str = "json") -> str:
        """Export current session
        
        Args:
            format: Export format (json or markdown)
            
        Returns:
            Exported content
        """
        if not self.current_session_id:
            return ""
        
        session = self.sessions[self.current_session_id]
        
        if format == "json":
            return json.dumps(session.to_dict(), indent=2, ensure_ascii=False)
        elif format == "markdown":
            markdown = f"# Session: {session.name}\n\n"
            for msg in session.messages:
                if msg.role == "user":
                    markdown += f"> {msg.content}\n\n"
                elif msg.role == "assistant":
                    markdown += f"{msg.content}\n\n"
                elif msg.role == "system":
                    markdown += f"*System:* {msg.content}\n\n"
            return markdown
        return ""
    
    def load_sessions(self):
        """Load sessions from disk"""
        session_file = os.path.join(self.session_dir, "sessions.json")
        if os.path.exists(session_file):
            try:
                with open(session_file, 'r', encoding='utf-8') as f:
                    data = json.load(f)
                    for session_data in data:
                        session = Session.from_dict(session_data)
                        self.sessions[session.session_id] = session
            except Exception as e:
                print(f"Error loading sessions: {e}")
    
    def save_sessions(self):
        """Save sessions to disk"""
        session_file = os.path.join(self.session_dir, "sessions.json")
        try:
            data = [session.to_dict() for session in self.sessions.values()]
            with open(session_file, 'w', encoding='utf-8') as f:
                json.dump(data, f, indent=2, ensure_ascii=False)
        except Exception as e:
            print(f"Error saving sessions: {e}")
    
    def get_current_session(self) -> Optional[Session]:
        """Get current session"""
        if self.current_session_id:
            return self.sessions.get(self.current_session_id)
        return None
    
    def list_sessions(self) -> List[Dict[str, Any]]:
        """List all sessions"""
        return [{
            "session_id": session.session_id,
            "name": session.name,
            "message_count": len(session.messages),
            "is_current": session.session_id == self.current_session_id
        } for session in self.sessions.values()]
