"""DQNServingStore — action-as-feature DQN 서빙 (선택 전용, 학습 없음).

tools/export_dqn_serving.py 가 만든 dqn_serving.npz(ScoreNet 가중치 + 카드특징 392차원)를
numpy 로 추론한다 — 서빙 컨테이너에 PyTorch 불필요.

역할 분담(계획서 H 트랙으로 가기 전 파일럿):
- 카드 '선택'만 DQN(greedy). Q-table 학습/영속/experience_logs 로깅 경로는 기존 그대로 유지
  (Q-learning 은 오프폴리시라 DQN 이 고른 행동으로 갱신해도 유효, 로그는 DQN 오프라인 재학습 재료).
- 폴백: 플래그 꺼짐 / 번들 없음 / 가용 카드 전부 특징 미보유(신규 카드) → None 반환,
  호출부(ChatService)가 기존 UCB Q-table 선택으로 진행한다.

활성화: 환경변수 DQN_SERVING=1 (docker-compose agent environment).
신규 카드 주의: 번들에 없는 카드번호는 후보에서 제외된다 — 카드 추가 시
tools/build_card_embeddings.py → tools/export_dqn_serving.py 재실행 후 재배포 필요.
"""

import os
from typing import List, Optional

import numpy as np

from common.logger import LOG
from negotiation.policies.base import ActionDecision, PolicyContext
from negotiation.qtable.domain.service.feature_builder import (
    build_state_features, build_tenant_features)

_HERE = os.path.dirname(os.path.abspath(__file__))
BUNDLE_PATH = os.path.join(_HERE, "..", "..", "artifacts", "dqn_serving.npz")


class _Bundle:
    def __init__(self, z):
        self.W0, self.b0 = z["W0"], z["b0"]
        self.W1, self.b1 = z["W1"], z["b1"]
        self.W2, self.b2 = z["W2"], z["b2"]
        self.state_dim = int(z["state_dim"])
        self.card_feats = {str(n): z["card_feats"][i]
                           for i, n in enumerate(z["card_numbers"])}

    def scores(self, state_feat: np.ndarray, card_feats: np.ndarray) -> np.ndarray:
        """가용 카드 K개 일괄 채점: [K, state+card] → [K]."""
        k = card_feats.shape[0]
        x = np.concatenate([np.repeat(state_feat[None, :], k, axis=0), card_feats], axis=1)
        h = np.maximum(x @ self.W0.T + self.b0, 0.0)
        h = np.maximum(h @ self.W1.T + self.b1, 0.0)
        return (h @ self.W2.T + self.b2).squeeze(-1)


class DQNServingPolicy:
    """UCBQTablePolicy.select 와 동일한 PolicyContext → ActionDecision 계약(선택 전용)."""

    name = "feature_dqn"
    _EPS = 0.1   # propensity 근사용 ε (UCB 정책과 동일 관례 — OPE 지지 확보용, 선택은 greedy)

    def __init__(self, bundle: _Bundle, engine):   # engine: tenancy.registry.TenantEngine
        self._bundle = bundle
        self._mapper = engine.mapper
        self._tenant_feat = build_tenant_features(engine.config.reward)

    def _available(self, ctx: PolicyContext) -> List[int]:
        # UCBQTablePolicy._available 과 동일 규칙 (마스크 → used 제외 → 소진 시 전체 허용)
        if ctx.available_mask is not None:
            avail = [a for a in range(ctx.action_space_size) if ctx.available_mask[a]]
        else:
            used = ctx.episode.used_action_ids if ctx.episode else set()
            avail = [a for a in range(ctx.action_space_size) if a not in used]
        return avail or list(range(ctx.action_space_size))

    def select(self, ctx: PolicyContext) -> Optional[ActionDecision]:
        """카드특징이 있는 가용 카드가 없으면 None → 호출부가 Q-table 로 폴백."""
        candidates = []   # (action_id, card_feat)
        for a in self._available(ctx):
            num = self._mapper.get_card_id(a)
            feat = self._bundle.card_feats.get(num) if num else None
            if feat is not None:
                candidates.append((a, feat))
        if not candidates:
            return None
        state_feat = np.concatenate([build_state_features(ctx.snapshot), self._tenant_feat])
        if state_feat.shape[0] != self._bundle.state_dim:
            LOG.e_no_callstack(
                f"[DQNServing] state_dim 불일치: {state_feat.shape[0]} != {self._bundle.state_dim}")
            return None
        sc = self._bundle.scores(state_feat, np.stack([f for _, f in candidates]))
        i = int(sc.argmax())
        n = len(candidates)
        return ActionDecision(
            action_id=candidates[i][0],
            propensity=(1.0 - self._EPS) + self._EPS / n,
            q_value=float(sc[i]),
            ucb_score=float(sc[i]),
            available_actions=[a for a, _ in candidates],
        )


class DQNServingStore:
    """번들 lazy 로드 + 캐시. 비활성/부재 시 None (호출부 Q-table 폴백)."""

    _bundle: Optional[_Bundle] = None
    _load_failed = False

    @classmethod
    def enabled(cls) -> bool:
        return os.getenv("DQN_SERVING", "0").lower() in ("1", "true", "yes")

    @classmethod
    def policy_for(cls, engine) -> Optional[DQNServingPolicy]:
        if not cls.enabled() or cls._load_failed:
            return None
        if cls._bundle is None:
            try:
                cls._bundle = _Bundle(np.load(BUNDLE_PATH, allow_pickle=False))
                LOG.i(f"[DQNServing] 번들 로드 완료: 카드 {len(cls._bundle.card_feats)}장")
            except Exception as ex:
                cls._load_failed = True   # 요청마다 재시도하지 않음
                LOG.e_no_callstack(f"[DQNServing] 번들 로드 실패 → Q-table 폴백: {ex}")
                return None
        return DQNServingPolicy(cls._bundle, engine)