"""QTablePolicyStore — learning 스키마에서 테넌트 UCBQTablePolicy 를 로드/영속화 (H1).

요청마다 활성 버전의 q_values/visit_counts 를 numpy QTable 로 적재해 정책을 조립한다(상태는 DB 가
단일 소스). 갱신은 touched 셀만 write-through → 재시작/요청 간 학습이 보존된다.

PoC 단순화: 온라인 single-step 갱신(다음상태 부트스트랩은 outcome 종료 시 생략). 시퀀스 보상링크는
P5/H 트랙에서 transition_id 기반으로 정교화한다.
"""

import math
from typing import Tuple

from common.logger import LOG
from negotiation.policies.qtable_policy import UCBQTablePolicy
from negotiation.qtable.domain.model.q_table import QTable
from negotiation.qtable.infra.repository.learning_repository import LearningRepository
from tenancy.registry import TenantEngine


class QTablePolicyStore:
    @staticmethod
    async def load(engine: TenantEngine) -> Tuple[UCBQTablePolicy, object, LearningRepository]:
        repo = LearningRepository(engine.company_id)
        S = engine.state_space_size
        A = engine.action_space_size
        pol_cfg = engine.config.policy
        lr = pol_cfg.learning_rate
        gamma = pol_cfg.gamma

        # 현재 카탈로그 스냅샷(action_id → 카드번호). 버전에 저장해 카드번호 기반 마이그레이션에 쓴다.
        card_list = [engine.mapper.get_card_id(i) for i in range(A)]

        # cold-start 3단 (계획서 D):
        #  ① 활성 버전 있으면 그대로  ② 없고 inherits_base 면 base warm-start 복제(차원 호환 시)
        #  ③ 차원 불일치/base 없음 → 휴리스틱 빈 버전
        err, active = await repo.read(lambda s: repo.get_active_version(s))
        if active is not None:
            dim_changed = active.action_space_size != A or active.state_space_size != S
            # 동일 차원이라도 카탈로그 내용(카드 구성)이 바뀌면 마이그레이션(카드번호 리맵).
            #   레거시 버전(action_cards=None)은 옛 구성을 몰라 내용변경 감지 불가 → 차원만 본다.
            cards_changed = active.action_cards is not None and list(active.action_cards) != card_list
            if dim_changed or cards_changed:
                # 카드번호 기반 학습 보존 마이그레이션: 같은 카드의 Q값을 새 action_id 로 이동.
                migrated = await repo.migrate_active_version_dim(
                    old_version=active, state_space_size=S, action_space_size=A,
                    learning_rate=lr, discount_factor=gamma, version_name=f"v_migrated_a{A}",
                    action_cards=card_list)
                if migrated is not None:
                    LOG.i(f"[QTablePolicyStore] 카탈로그 변경 마이그레이션 company={engine.company_id} "
                          f"action {active.action_space_size}→{A} (카드번호 리맵, 학습 보존)")
                version_id = migrated or active.version_id
            else:
                version_id = active.version_id
                if active.action_cards is None:   # 레거시 버전 backfill → 향후 카탈로그 변경 감지 가능
                    await repo.set_version_action_cards(version_id, card_list)
        else:
            version_id = None
            if engine.config.inherits_base:
                version_id = await repo.warm_start_from_base(
                    state_space_size=S, action_space_size=A, learning_rate=lr, discount_factor=gamma,
                    action_cards=card_list)
            if version_id is None:
                version_id = await repo.get_or_create_active_version(
                    state_space_size=S, action_space_size=A, learning_rate=lr, discount_factor=gamma,
                    action_cards=card_list)

        qtable = QTable(S, A, learning_rate=lr, discount_factor=gamma)
        qrows, vrows = await repo.load_cells(version_id)
        for st, a, q in qrows:
            if 0 <= st < S and 0 <= a < A:
                qtable.q[st, a] = q
        for st, a, c in vrows:
            if 0 <= st < S and 0 <= a < A:
                qtable.visits[st, a] = c

        params = pol_cfg.params or {}
        policy = UCBQTablePolicy(
            qtable,
            exploration_constant=params.get("exploration_constant", math.sqrt(2.0)),
            epsilon=params.get("propensity_epsilon", 0.1),
        )
        return policy, version_id, repo

    @staticmethod
    async def persist_cell(repo: LearningRepository, version_id, policy: UCBQTablePolicy,
                           state_index: int, action_id: int):
        """(state, action) 셀 write-through (select 의 visit 증가 + update 의 Q 변화 반영)."""
        q_value = float(policy.qtable.q[state_index, action_id])
        count = int(policy.qtable.visits[state_index, action_id])
        await repo.upsert_cell(version_id, state_index, action_id, q_value, count)