"""init_base — 공유 베이스 정책(_base) 시드 (P5, 계획서 D).

시뮬레이터로 UCB Q-Table 을 학습시켜 learning 스키마의 _base(scope=base) 버전에 저장한다.
신규 테넌트는 cold-start 시 이 베이스를 warm-start 복제해 어느 정도 학습된 상태로 시작한다.

실행: cd agent && APP_ENV=local python -m tools.init_base [--action-space 9] [--episodes 500]
차원(state×action)은 신규 테넌트와 호환돼야 복제된다(데모 테넌트=162×9).
"""

import argparse
import asyncio

from common.database.db_session_manager import DB_SESSION_MNG
from common.database.model.models import BASE_COMPANY_ID
from common.logger import LOG
from eval_harness.buyer import HeuristicBuyer, Scenario, best_actions, make_card_effectiveness
from eval_harness.registry import build_policy
from eval_harness.simulator import run_episode
from negotiation.qtable.infra.repository.learning_repository import LearningRepository
from tenancy.config_loader import TenantConfigLoader


async def seed_base(action_space: int = 9, episodes: int = 500, seed: int = 123,
                    anchor: float = 8000, target: float = 10000) -> dict:
    cfg = TenantConfigLoader().load(BASE_COMPANY_ID)   # _base: state 162
    S = cfg.state.state_space_size
    A = action_space
    lr, gamma = cfg.policy.learning_rate, cfg.policy.gamma

    # 시뮬레이터로 베이스 학습
    policy = build_policy("qtable_ucb", cfg.state, A, cfg.policy, seed=seed)
    eff = make_card_effectiveness(A, seed=seed)
    buyer = HeuristicBuyer(eff, seed=seed)
    scenario = Scenario(anchor_price=anchor, target_price=target)
    for i in range(episodes):
        buyer.reseed(seed * 100_000 + i)
        run_episode(policy, buyer, scenario, cfg.state, cfg.reward, A, learn=True)

    # _base 버전(scope=base)에 저장 (재시드 시 기존 셀 비우고 갱신)
    repo = LearningRepository(BASE_COMPANY_ID)
    vid = await repo.get_or_create_active_version(
        state_space_size=S, action_space_size=A, learning_rate=lr, discount_factor=gamma,
        scope=1, version_name="base_v000")
    await repo.reset_learning()  # 이전 셀 정리(버전은 유지)
    cells = policy.qtable.nonzero_cells()
    for st, a, q, c in cells:
        await repo.upsert_cell(vid, st, a, q, c)

    return {"version_id": str(vid), "state_space": S, "action_space": A,
            "episodes": episodes, "cells": len(cells), "good_cards": best_actions(eff)}


async def _main(args):
    info = await seed_base(action_space=args.action_space, episodes=args.episodes)
    print("=" * 56)
    print("  베이스 정책 시드 완료 (_base, scope=base)")
    print(f"  버전: {info['version_id']}  차원: {info['state_space']}x{info['action_space']}")
    print(f"  학습 에피소드: {info['episodes']}  저장 셀: {info['cells']}")
    print(f"  (숨은) 좋은 카드: {info['good_cards']}")
    print("  → 이제 신규 테넌트 첫 협상 시 warm-start 로 이 베이스를 복제합니다.")
    print("=" * 56)
    await DB_SESSION_MNG.dispose_all()


def main():
    ap = argparse.ArgumentParser(description="공유 베이스 정책 시드 (P5)")
    ap.add_argument("--action-space", type=int, default=9)
    ap.add_argument("--episodes", type=int, default=500)
    args = ap.parse_args()
    asyncio.run(_main(args))


if __name__ == "__main__":
    main()