"""정책 컨텍스트/결정 + 요청 스코프 episode 상태 (계획서 B 동시성, F 추상화 출발점).

핵심(계획서 리스크): Chat_server 의 UCBPolicy._episode_actions 는 **인스턴스 멤버**라,
테넌트 엔진을 여러 요청이 공유하면 세션 간 오염이 발생한다. 여기서는 episode 상태를
**요청 스코프 객체(EpisodeState)** 로 외부화하고, 정책은 stateless 하게 이를 주입받는다.

NegotiationPolicy(상위 추상)의 train/predict_action_dist 등 전체 인터페이스는 H0 에서 확장한다.
P4 는 동시성에 필요한 PolicyContext/ActionDecision/EpisodeState 만 정의한다.
"""

from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from typing import List, Optional, Set

import numpy as np

from negotiation.qtable.domain.model.snapshot import NegotiationSnapshot


@dataclass
class EpisodeState:
    """한 협상 세션(요청 흐름)의 가변 상태. 절대 정책/엔진 인스턴스에 두지 않는다.

    used_action_ids: 이 에피소드에서 이미 제시한 action(중복 방지 마스킹용).
    """

    used_action_ids: Set[int] = field(default_factory=set)

    def mark_used(self, action_id: int):
        self.used_action_ids.add(action_id)


@dataclass
class PolicyContext:
    """정책이 행동을 고르기 위해 받는 입력. state_index(이산)와 snapshot(연속) 모두 포함
    → Q-Table/LinUCB/Offline RL 이 같은 컨텍스트를 공유한다(계획서 F).
    """

    state_index: int
    snapshot: NegotiationSnapshot
    action_space_size: int
    episode: EpisodeState
    available_mask: Optional[np.ndarray] = None  # None 이면 used_action_ids 로 산출


@dataclass
class ActionDecision:
    """정책의 출력. propensity(행동확률)는 OPE 의 전제(계획서 H0)."""

    action_id: int
    propensity: float
    card_id: Optional[str] = None
    q_value: Optional[float] = None
    ucb_score: Optional[float] = None
    available_actions: Optional[List[int]] = None


@dataclass
class Transition:
    """학습용 (s, a, r, s', done) + OPE 메타. experience_logs 한 row 에 대응."""

    state_index: int
    action_id: int
    reward: float
    next_state_index: Optional[int] = None
    done: bool = False
    propensity: Optional[float] = None


class NegotiationPolicy(ABC):
    """상위 정책 추상 (계획서 F). Q-Table/LinUCB/Offline RL 이 모두 이 인터페이스를 구현한다.

    Chat_server 의 Policy.select_action(state_index, q_values: np.ndarray, ...) 는 Q값 배열에
    결합돼 LinUCB/CQL 을 감쌀 수 없었다. 여기서는 PolicyContext(state_index+snapshot)를 받아
    ActionDecision(propensity 포함)을 돌려주는 알고리즘-중립 인터페이스로 둔다.
    """

    name: str = "base"

    @abstractmethod
    def select(self, ctx: PolicyContext) -> ActionDecision:
        """행동 선택. propensity(선택확률)를 반드시 채운다(OPE 전제)."""

    @abstractmethod
    def update(self, transition: Transition) -> None:
        """온라인 1-스텝 갱신 (Q-learning 등). 배치 학습은 train()."""

    def train(self, transitions: List[Transition]) -> None:
        """오프라인 배치 학습 (기본: update 반복). 알고리즘별 override."""
        for t in transitions:
            self.update(t)

    def predict_action_dist(self, ctx: PolicyContext) -> np.ndarray:
        """상태에서의 행동 분포 (OPE/시뮬레이터용). 기본: 선택 액션에 1.0."""
        dist = np.zeros(ctx.action_space_size)
        dist[self.select(ctx).action_id] = 1.0
        return dist

    def warm_start(self, other: "NegotiationPolicy") -> None:
        """베이스 정책으로부터 초기화 (계획서 D). 기본: 미지원."""
        raise NotImplementedError

    def snapshot(self) -> dict:
        """직렬화 가능한 파라미터 스냅샷(저장용). 알고리즘별 구현."""
        raise NotImplementedError

    def load_snapshot(self, data: dict) -> None:
        raise NotImplementedError