"""AdapterPool — Sprint 16 runtime 核心：共享 base model 上的 LoRA 熱插拔。

【這個模組在做什麼？】
  Gemma 4 E4B 整個 ML service 只載入一份（共享 base model），但每位使用者可能
  有自己的 stylist LoRA adapter。peft 的 `set_adapter()` 會直接改動模型狀態 —
  不同併發請求如果沒序列化，後一個會看到前一個的 adapter，造成跨使用者資料污染。

  AdapterPool 用 asyncio.Lock 把「load → set → yield → reset」整段序列化，
  並搭配 LRU 管理已載入的 adapter 數量（避免 HF Spaces Free tier VRAM 爆掉）。

【關鍵設計決策】

  1. **default adapter 永遠 pin 在 pool**
     主流程（stylist 給新手 + 非訂閱者）都走 default；如果被驅逐會造成所有
     無 adapter 請求每次都要 reload，CPU 100% 拉爆。

  2. **asyncio.Lock 序列化 set_adapter**
     peft 的 set_adapter 是 mutation，不 serialize 會 race。這個 lock 粒度
     粗一點（包整個 generate 區塊）比讓 peft 內部 lock 再 fine-grain 更安全。
     trade-off：單 ML service process 併發容量 = 1（在 set_adapter 階段）。
     對 Wardrobe OS 的流量規模（< 10 req/s）完全夠用。

  3. **capacity 包含 default**
     更精確地說：capacity 是 non-default adapter 數量的上限。default 不算進
     LRU queue，永遠 pinned。所以 capacity=4 實際 loaded 總數最多是 5。

  4. **HF Spaces Free tier 預設 capacity=4**
     review v8 P1-7 的選項 (c)：免費方案 + capacity=4 + cold-start penalty
     acceptable。MVP 階段 marketplace subscribers 現實上 < 10 個，capacity=4
     夠用；> 100 時再升級到 (a) 或 (b)。

【併發模型】
  asyncio.Lock → FastAPI async 請求自然串起來。不是 threading.Lock（ML service
  是 async，沒 GIL 問題；threading.Lock 會整個 event loop 掛掉）。

【對應測試】
  apps/ml-service/tests/test_adapter_pool.py（用 FakePeftModel mock，不載實模型）
"""

from __future__ import annotations

import asyncio
import logging
from collections import OrderedDict
from contextlib import asynccontextmanager
from typing import Any, AsyncIterator

logger = logging.getLogger(__name__)


class AdapterPool:
    """LRU pool for hot-swapping PEFT LoRA adapters on a shared base model."""

    def __init__(
        self,
        model: Any,
        *,
        default_name: str,
        capacity: int = 4,
    ) -> None:
        """初始化 pool。

        Args:
            model: peft.PeftModel 或相容物件（需要 load_adapter / set_adapter /
                delete_adapter 三個方法，以及 loaded_adapters 的存在性檢查能力）。
                Production 會是實際的 PeftModel；測試會是 FakePeftModel。
            default_name: 已經 pin 在 base model 上的 default adapter name。
                AdapterPool 假設呼叫者在 __init__ 之前已經把 default 載入好。
            capacity: non-default adapter 同時 loaded 的上限。default 不算進去。
        """
        if capacity < 1:
            raise ValueError(f"capacity must be >= 1, got {capacity}")

        self._model = model
        self._default = default_name
        self._capacity = capacity

        # LRU queue 只裝 non-default adapter。最近使用的放在尾端，最舊的在頭端。
        # 用 OrderedDict 是因為 move_to_end 是 O(1)，加上 popitem(last=False) 也是 O(1)。
        # value 欄不重要（我們只關心 key 順序），存 hf_repo_id 方便 debug。
        self._lru: OrderedDict[str, str] = OrderedDict()

        # asyncio.Lock — 整個 use() context 內 set_adapter + yield 都 hold 著
        self._lock = asyncio.Lock()

    # ------------------------------------------------------------------
    # 公開 API
    # ------------------------------------------------------------------

    def is_loaded(self, name: str) -> bool:
        """檢查 adapter 是否在 pool 內（default 永遠是 True）。"""
        if name == self._default:
            return True
        return name in self._lru

    @asynccontextmanager
    async def use(self, name: str, hf_repo_id: str) -> AsyncIterator[Any]:
        """切到指定 adapter 執行 generation，離開時切回 default。

        【使用方式】
            async with pool.use("user-uuid", "hf://fengchen31/wardrobe-os-stylist-<uuid>") as model:
                output = model.generate(..., use_cache=False)  # ← Sprint 16 Task 2.0 incident

        【契約】
          - 進入時：adapter 被載入（如果還沒）並 set 為 active
          - 離開時：無論是否有 exception，adapter 都切回 default
          - 整段序列化：同時呼叫 use() 的請求會 FIFO 排隊

        Args:
            name: adapter name（通常是使用者 UUID，或 "default"）。
            hf_repo_id: HF Hub 路徑，load_adapter 時用。name=="default" 時會被忽略。

        Yields:
            底層的 model 物件（方便 caller 直接拿來 generate）。
        """
        async with self._lock:
            # 載入 + 切換
            await self._ensure_loaded(name, hf_repo_id)
            # peft 的 set_adapter 是同步 torch mutation — 放到 worker thread 以免
            # 凍結 event loop（Gemma 11 capability 共用同一個 loop）
            await asyncio.to_thread(self._model.set_adapter, name)
            logger.debug("adapter_pool: set active → %s", name)

            try:
                yield self._model
            finally:
                # 一律切回 default — 防跨使用者污染，即使 caller 拋 exception
                await asyncio.to_thread(self._model.set_adapter, self._default)
                logger.debug("adapter_pool: reset → %s", self._default)

    # ------------------------------------------------------------------
    # 內部
    # ------------------------------------------------------------------

    async def _ensure_loaded(self, name: str, hf_repo_id: str) -> None:
        """確保 adapter 已在 pool 內；該驅逐就驅逐，該載入就載入。"""
        if name == self._default:
            # default 永遠 loaded，no-op
            return

        if name in self._lru:
            # hit — 把它推到 LRU 最新端
            self._lru.move_to_end(name, last=True)
            return

        # miss — 需要載入，先看是否要驅逐最舊的
        if len(self._lru) >= self._capacity:
            oldest_name, oldest_repo = self._lru.popitem(last=False)
            logger.info(
                "adapter_pool: evict %s (hf=%s) to make room for %s",
                oldest_name,
                oldest_repo,
                name,
            )
            try:
                await asyncio.to_thread(self._model.delete_adapter, oldest_name)
            except Exception as exc:
                # peft 的 delete_adapter 理論上不會失敗，但保險起見不讓它炸掉
                # 整個 pool。記 warning 繼續 — LRU tracking 已經移除了。
                logger.warning(
                    "adapter_pool: delete_adapter(%s) failed: %s — LRU still evicted",
                    oldest_name,
                    exc,
                )

        # 真正 load — HF Hub download + torch weight load，30s+ on cold miss
        # 必須 to_thread，否則整個 event loop 被這段凍住
        logger.info("adapter_pool: load_adapter %s from %s", name, hf_repo_id)
        await asyncio.to_thread(self._model.load_adapter, hf_repo_id, adapter_name=name)
        self._lru[name] = hf_repo_id