import asyncio
import base64
import logging
from pathlib import Path
from typing import Any

from app.bots.assistant import AssistantRuntime
from app.config import Settings
from app.storage.hf_dataset import HFDataStore


class SessionManager:
    def __init__(self, settings: Settings, store: HFDataStore) -> None:
        self.settings = settings
        self.store = store
        self.base_dir = Path("sessions")
        self.base_dir.mkdir(exist_ok=True)
        self._assistants: dict[str, AssistantRuntime] = {}

    def running_ids(self) -> list[str]:
        return list(self._assistants.keys())

    def get_runtime(self, assistant_id: str) -> AssistantRuntime | None:
        return self._assistants.get(assistant_id)

    async def start_assistant(self, assistant_id: str, assistant_data: dict[str, Any]) -> None:
        if assistant_id in self._assistants:
            return

        session_b64 = assistant_data.get("session_b64", "")
        session_path = self.base_dir / f"{assistant_id}.session"
        if session_b64:
            try:
                session_path.write_bytes(base64.b64decode(session_b64.encode("utf-8")))
            except Exception as exc:
                raise ValueError(f"Corrupted session data for assistant {assistant_id}") from exc

        runtime = AssistantRuntime(self.settings, self.store, assistant_id, str(session_path))
        await runtime.start()
        self._assistants[assistant_id] = runtime

    async def stop_assistant(self, assistant_id: str) -> None:
        runtime = self._assistants.pop(assistant_id, None)
        if runtime:
            await runtime.stop()
        (self.base_dir / f"{assistant_id}.session").unlink(missing_ok=True)

    async def load_all(self) -> None:
        assistants = self.store.get_data().get("assistants", {})
        for assistant_id, data in assistants.items():
            try:
                await self.start_assistant(assistant_id, data)
            except Exception:
                logging.exception("Failed to start assistant %s", assistant_id)

    async def shutdown(self) -> None:
        for assistant_id in list(self._assistants.keys()):
            await self.stop_assistant(assistant_id)