# retoor <retoor@molodetz.nl>

from __future__ import annotations

import logging
from typing import Any, Callable

from devplacepy.database import db

from .agentic import LessonStore
from .behavior import BehaviorStore
from .config import Settings
from .cost.tracker import Pricing
from .llm import LLMClient
from .session import DeviiSession
from .store import ConversationStore, TurnAudit, UsageLedger
from .tasks.store import TaskStore, memory_db
from .virtual_tools import VirtualToolStore

logger = logging.getLogger("devii.hub")

SettingsBuilder = Callable[[str, str, str, bool], "tuple[Settings, Pricing]"]


def _user_key_resolver(owner_id: str) -> Callable[[], str]:
    def resolve() -> str:
        if "users" not in db.tables:
            return ""
        row = db["users"].find_one(uid=owner_id, deleted_at=None)
        return (row or {}).get("api_key") or ""

    return resolve


class DeviiHub:
    def __init__(self, settings_builder: SettingsBuilder) -> None:
        self._build = settings_builder
        self._stores: dict[str, Any] = {
            "conversations": ConversationStore(),
            "ledger": UsageLedger(),
            "turns": TurnAudit(),
        }
        self._sessions: dict[tuple[str, str, str], DeviiSession] = {}

    @property
    def ledger(self) -> UsageLedger:
        return self._stores["ledger"]

    def get_or_create(
        self,
        owner_kind: str,
        owner_id: str,
        username: str,
        api_key: str,
        base_url: str,
        is_admin: bool = False,
        is_primary_admin: bool = False,
        channel: str = "main",
    ) -> DeviiSession:
        key = (owner_kind, owner_id, channel)
        session = self._sessions.get(key)
        if session is not None:
            return session
        settings, pricing = self._build(api_key, base_url, owner_kind, is_admin)
        key_resolver = (
            _user_key_resolver(owner_id) if owner_kind == "user" else None
        )
        llm = LLMClient(settings, key_resolver=key_resolver)
        owned_db = (
            db
            if (owner_kind == "user" and channel in ("main", "telegram"))
            else memory_db()
        )
        task_store = TaskStore(owned_db, owner_kind, owner_id)
        lessons = LessonStore(owned_db, owner_kind, owner_id)
        virtual_tool_store = VirtualToolStore(owned_db, owner_kind, owner_id)
        behavior_store = BehaviorStore(owned_db, owner_kind, owner_id)
        session = DeviiSession(
            owner_kind,
            owner_id,
            username,
            settings,
            llm,
            lessons,
            pricing,
            task_store,
            virtual_tool_store,
            behavior_store,
            self._stores,
            is_admin=is_admin,
            is_primary_admin=is_primary_admin,
            channel=channel,
            key_resolver=key_resolver,
        )
        if owner_kind == "user":
            saved = self._stores["conversations"].load(owner_kind, owner_id, channel)
            if saved:
                session.restore_history(saved)
        self._sessions[key] = session
        logger.info(
            "Created session %s/%s [%s] (total %d)",
            owner_kind,
            owner_id,
            channel,
            len(self._sessions),
        )
        return session

    def find(
        self, owner_kind: str, owner_id: str, channel: str = "main"
    ) -> DeviiSession | None:
        return self._sessions.get((owner_kind, owner_id, channel))

    def active_sessions(self) -> int:
        return len(self._sessions)

    def connection_count(self) -> int:
        return sum(s.connection_count for s in self._sessions.values())

    async def gc_idle(self) -> int:
        removed = 0
        for key, session in list(self._sessions.items()):
            if session.connection_count == 0 and not session.is_busy():
                await session.aclose()
                del self._sessions[key]
                removed += 1
        if removed:
            logger.info("GC removed %d idle session(s)", removed)
        return removed

    async def aclose(self) -> None:
        for session in list(self._sessions.values()):
            await session.aclose()
        self._sessions.clear()
