"""Session management for agent conversations using SQLite"""

from __future__ import annotations

import os
import uuid
from dataclasses import dataclass, field
from datetime import datetime
from typing import Optional

from agents import Agent, Runner, SQLiteSession, RunConfig

from diona.project.agent.model import BaseModelProvider
from diona.project.agent.sub_registry import get_sub_agent_registry


@dataclass
class SessionInfo:
    """Information about a session"""
    session_id: str
    name: str
    created_at: datetime
    last_active: datetime


class AgentSessionManager:
    """Manager for agent sessions using SQLite storage"""
    
    _instance: Optional[AgentSessionManager] = None
    
    def __new__(cls) -> AgentSessionManager:
        if cls._instance is None:
            cls._instance = super().__new__(cls)
            cls._instance._initialized = False
        return cls._instance
    
    def __init__(self):
        if self._initialized:
            return
        
        self._storage = self._get_storage()
        self._sessions_dir = self._get_sessions_dir()
        self._sessions: dict[str, SessionInfo] = {}
        self._current_session_id: Optional[str] = None
        self._core_agent = None
        self._model_provider = None
        self._initialized = True
        
        self._load_sessions()
    
    def _get_storage(self):
        """Get storage instance"""
        from diona.project.storage import get_storage_instance
        return get_storage_instance()
    
    def _get_sessions_dir(self) -> str:
        """Get sessions directory path"""
        base_path = self._storage.get_path("agent", "sessions", storage_type="user")
        os.makedirs(base_path, exist_ok=True)
        return base_path
    
    def _get_session_db_path(self, session_id: str) -> str:
        """Get database path for a session"""
        return os.path.join(self._sessions_dir, f"{session_id}.db")
    
    def _load_sessions(self) -> None:
        """Load existing sessions from directory"""
        if not os.path.exists(self._sessions_dir):
            return
        
        for filename in os.listdir(self._sessions_dir):
            if filename.endswith(".db"):
                session_id = filename[:-3]
                db_path = self._get_session_db_path(session_id)
                if os.path.exists(db_path):
                    stat = os.stat(db_path)
                    created_at = datetime.fromtimestamp(stat.st_ctime)
                    last_active = datetime.fromtimestamp(stat.st_mtime)
                    
                    self._sessions[session_id] = SessionInfo(
                        session_id=session_id,
                        name=f"Session {session_id[:8]}",
                        created_at=created_at,
                        last_active=last_active
                    )
        
        if self._sessions:
            most_recent = max(self._sessions.values(), key=lambda s: s.last_active)
            self._current_session_id = most_recent.session_id
    
    def set_model_provider(self, model_provider: BaseModelProvider) -> None:
        """Set model provider for sessions"""
        self._model_provider = model_provider
    
    def create_session(self, name: Optional[str] = None) -> str:
        """Create a new session
        
        Args:
            name: Optional session name
            
        Returns:
            Session ID
        """
        session_id = str(uuid.uuid4())
        
        db_path = self._get_session_db_path(session_id)
        session = SQLiteSession(session_id, db_path)
        
        session_info = SessionInfo(
            session_id=session_id,
            name=name or f"Session {session_id[:8]}",
            created_at=datetime.now(),
            last_active=datetime.now()
        )
        
        self._sessions[session_id] = session_info
        self._current_session_id = session_id
        
        return session_id
    
    def get_session(self, session_id: str) -> Optional[SessionInfo]:
        """Get session info
        
        Args:
            session_id: Session ID
            
        Returns:
            SessionInfo or None
        """
        return self._sessions.get(session_id)
    
    def list_sessions(self) -> list[SessionInfo]:
        """List all sessions
        
        Returns:
            List of session info
        """
        return list(self._sessions.values())
    
    def delete_session(self, session_id: str) -> bool:
        """Delete a session
        
        Args:
            session_id: Session ID
            
        Returns:
            True if deleted
        """
        if session_id not in self._sessions:
            return False
        
        db_path = self._get_session_db_path(session_id)
        if os.path.exists(db_path):
            os.remove(db_path)
        
        del self._sessions[session_id]
        
        if self._current_session_id == session_id:
            self._current_session_id = None
        
        return True
    
    def set_current_session(self, session_id: str) -> bool:
        """Set current active session
        
        Args:
            session_id: Session ID
            
        Returns:
            True if set successfully
        """
        if session_id not in self._sessions:
            return False
        
        self._current_session_id = session_id
        return True
    
    def get_current_session_id(self) -> Optional[str]:
        """Get current session ID
        
        Returns:
            Current session ID or None
        """
        return self._current_session_id
    
    def run_with_session(self, user_input: str, model_provider: Optional[BaseModelProvider] = None) -> str:
        """Run agent with current session
        
        Args:
            user_input: User input
            model_provider: Optional model provider
            
        Returns:
            Agent response
        """
        from diona.project.agent import get_core_agent
        
        if self._current_session_id is None:
            self.create_session()
        
        if self._core_agent is None:
            self._core_agent = get_core_agent(
                model_provider=model_provider or self._model_provider
            )
        
        db_path = self._get_session_db_path(self._current_session_id)
        session = SQLiteSession(self._current_session_id, db_path)
        
        run_config = RunConfig(
            model_provider=model_provider or self._core_agent._model_provider
        )
        
        import asyncio
        result = asyncio.run(
            Runner.run(self._core_agent.agent, user_input, session=session, run_config=run_config)
        )
        
        if self._current_session_id in self._sessions:
            self._sessions[self._current_session_id].last_active = datetime.now()
        
        return result.final_output
    
    async def run_with_session_streamed(
        self,
        user_input: str,
        model_provider: Optional[BaseModelProvider] = None,
        on_token: Optional[callable] = None
    ):
        """Run agent with streaming output
        
        Args:
            user_input: User input
            model_provider: Optional model provider
            on_token: Callback for each token received
            
        Yields:
            Agent response tokens
        """
        from diona.project.agent import get_core_agent
        
        if self._current_session_id is None:
            self.create_session()
        
        if self._core_agent is None:
            self._core_agent = get_core_agent(
                model_provider=model_provider or self._model_provider
            )
        
        db_path = self._get_session_db_path(self._current_session_id)
        session = SQLiteSession(self._current_session_id, db_path)
        
        run_config = RunConfig(
            model_provider=model_provider or self._core_agent._model_provider
        )
        
        result = Runner.run_streamed(
            self._core_agent.agent,
            user_input,
            session=session,
            run_config=run_config
        )
        
        final_output = ""
        
        async for event in result.stream_events():
            if event.type == "raw_response_event":
                if hasattr(event, 'data') and hasattr(event.data, 'delta'):
                    token = event.data.delta
                    final_output += token
                    if on_token:
                        on_token(token)
                    yield token
            
            elif event.type == "run_item_stream_event":
                if hasattr(event, 'item'):
                    pass
        
        if self._current_session_id in self._sessions:
            self._sessions[self._current_session_id].last_active = datetime.now()
        
        # 不再 yield final_output，避免重复输出
        # yield final_output
    
    def clear_all_sessions(self) -> None:
        """Clear all sessions"""
        for session_id in list(self._sessions.keys()):
            self.delete_session(session_id)


def get_session_manager() -> AgentSessionManager:
    """Get singleton session manager
    
    Returns:
        AgentSessionManager instance
    """
    return AgentSessionManager()
