"""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

        # 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:
            if active.action_space_size != A or active.state_space_size != S:
                # 카탈로그 카드 수(또는 상태 차원) 변경 → 학습 보존 마이그레이션.
                #   겹치는 셀 복사(append/truncate 안전) + 새 카드 fresh. 실패 시 기존 버전 유지(load 가 reshape).
                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}")
                if migrated is not None:
                    LOG.i(f"[QTablePolicyStore] 차원 변경 마이그레이션 company={engine.company_id} "
                          f"{active.action_space_size}→{A} action (학습 보존)")
                version_id = migrated or active.version_id
            else:
                version_id = active.version_id
        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)
            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)

        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)