from __future__ import annotations

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

from agents import Agent, Runner, SQLiteSession, RunConfig
from agents.memory import SessionSettings

from diona.project.agent.model import BaseModelProvider
from diona.project.agent.sub_agents.registry import get_sub_agent_registry

MAX_TURNS = 20
MAX_HISTORY_ITEMS = 30

SESSION_SETTINGS = SessionSettings(limit=None)


@dataclass
class SessionInfo:
    session_id: str
    name: str
    created_at: datetime
    last_active: datetime


class SessionManager:
    _instance: Optional[SessionManager] = None

    def __new__(cls) -> SessionManager:
        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):
        from diona.project.storage import get_storage_instance
        return get_storage_instance()

    def _get_sessions_dir(self) -> str:
        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:
        return os.path.join(self._sessions_dir, f"{session_id}.db")

    def _load_sessions(self) -> None:
        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:
        self._model_provider = model_provider

    def create_session(self, name: Optional[str] = None) -> str:
        session_id = str(uuid.uuid4())

        db_path = self._get_session_db_path(session_id)
        session = SQLiteSession(session_id, db_path, session_settings=SESSION_SETTINGS)

        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]:
        return self._sessions.get(session_id)

    def list_sessions(self) -> list[SessionInfo]:
        return list(self._sessions.values())

    def delete_session(self, session_id: str) -> bool:
        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:
        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]:
        return self._current_session_id

    async def get_session_items(self, session_id: Optional[str] = None, limit: Optional[int] = None) -> list:
        sid = session_id or self._current_session_id
        if sid is None:
            return []

        db_path = self._get_session_db_path(sid)
        session = SQLiteSession(sid, db_path, session_settings=SESSION_SETTINGS)

        return await session.get_items(limit=limit)

    def get_session_items_sync(self, session_id: Optional[str] = None, limit: Optional[int] = None) -> list:
        import asyncio

        sid = session_id or self._current_session_id
        if sid is None:
            return []

        db_path = self._get_session_db_path(sid)
        session = SQLiteSession(sid, db_path, session_settings=SESSION_SETTINGS)

        try:
            loop = asyncio.get_running_loop()
            try:
                import nest_asyncio
                nest_asyncio.apply()
            except ImportError:
                pass
            return loop.run_until_complete(session.get_items(limit=limit))
        except RuntimeError:
            return asyncio.run(session.get_items(limit=limit))

    def run_with_session(self, user_input: str, model_provider: Optional[BaseModelProvider] = None) -> str:
        from diona.project.agent.core import get_core_agent
        import asyncio

        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, session_settings=SESSION_SETTINGS)

        run_config = RunConfig(
            model_provider=model_provider or self._core_agent._model_provider
        )

        async def _run():
            result = await Runner.run(
                self._core_agent.inner_agent,
                user_input,
                session=session,
                run_config=run_config,
                max_turns=MAX_TURNS
            )
            return result.final_output

        try:
            loop = asyncio.get_running_loop()
            try:
                import nest_asyncio
                nest_asyncio.apply()
            except ImportError:
                pass
            final_result = loop.run_until_complete(_run())
        except RuntimeError:
            final_result = asyncio.run(_run())

        if self._current_session_id in self._sessions:
            self._sessions[self._current_session_id].last_active = datetime.now()

        return final_result

    async def run_with_session_streamed(
        self,
        user_input: str,
        model_provider: Optional[BaseModelProvider] = None,
        on_token: Optional[callable] = None
    ):
        from diona.project.agent.core 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, session_settings=SESSION_SETTINGS)

        run_config = RunConfig(
            model_provider=model_provider or self._core_agent._model_provider
        )

        result = Runner.run_streamed(
            self._core_agent.inner_agent,
            user_input,
            session=session,
            run_config=run_config,
            max_turns=MAX_TURNS
        )

        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()

    def clear_all_sessions(self) -> None:
        for session_id in list(self._sessions.keys()):
            self.delete_session(session_id)

    def run_without_session(self, user_input: str, model_provider: Optional[BaseModelProvider] = None) -> str:
        from diona.project.agent.core import get_core_agent
        import asyncio

        if self._core_agent is None:
            self._core_agent = get_core_agent(
                model_provider=model_provider or self._model_provider
            )

        run_config = RunConfig(
            model_provider=model_provider or self._core_agent._model_provider
        )

        async def _run():
            result = await Runner.run(
                self._core_agent.inner_agent,
                user_input,
                run_config=run_config,
                max_turns=MAX_TURNS
            )
            return result.final_output

        try:
            loop = asyncio.get_running_loop()
            try:
                import nest_asyncio
                nest_asyncio.apply()
            except ImportError:
                pass
            return loop.run_until_complete(_run())
        except RuntimeError:
            return asyncio.run(_run())

    async def run_without_session_streamed(
        self,
        user_input: str,
        model_provider: Optional[BaseModelProvider] = None,
        on_token: Optional[callable] = None
    ):
        from diona.project.agent.core import get_core_agent

        if self._core_agent is None:
            self._core_agent = get_core_agent(
                model_provider=model_provider or self._model_provider
            )

        run_config = RunConfig(
            model_provider=model_provider or self._core_agent._model_provider
        )

        result = Runner.run_streamed(
            self._core_agent.inner_agent,
            user_input,
            run_config=run_config,
            max_turns=MAX_TURNS
        )

        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


_manager: Optional[SessionManager] = None


def get_session_manager() -> SessionManager:
    global _manager
    if _manager is None:
        _manager = SessionManager()
    return _manager
