"""로컬 학습 리셋 — learning 스키마를 비운다 (카드 카탈로그·협상 데이터는 보존).

카탈로그(카드) 구성을 바꾼 뒤 Q-table 을 처음부터 다시 학습시키고 싶을 때 사용한다.
learning.* (버전/Q값/방문수/경험로그/세션/카드매핑) 만 삭제 → 다음 협상부터 fresh 재학습.
card.*·negotiation.*·quotation.* 등 실제 데이터는 건드리지 않는다.

사용 (agent 디렉토리에서):
  APP_ENV=local python -m tools.reset_learning                 # 전체 회사 학습 리셋
  APP_ENV=local python -m tools.reset_learning --company <id>  # 특정 회사만
  APP_ENV=local python -m tools.reset_learning --yes           # 확인 프롬프트 생략
"""

import argparse
import asyncio

import asyncpg

from config.server_configs import main_db_config

# 삭제 대상(learning 스키마). 전부 company_id 스코프라 --company 로 회사별 리셋 가능.
_LEARNING_TABLES = [
    "q_values", "visit_counts", "experience_logs", "chat_sessions",
    "tenant_action_cards", "q_table_versions",
]


async def _reset(company_id: str | None) -> None:
    cfg = main_db_config
    conn = await asyncpg.connect(
        host=cfg.write_host, port=cfg.write_port,
        user=cfg.write_id, password=cfg.write_pw, database=cfg.name,
    )
    try:
        total = 0
        for t in _LEARNING_TABLES:
            if company_id:
                res = await conn.execute(f"DELETE FROM learning.{t} WHERE company_id = $1", company_id)
            else:
                res = await conn.execute(f"DELETE FROM learning.{t}")
            n = int(res.split()[-1]) if res else 0
            total += n
            print(f"  learning.{t}: {n} 행 삭제")
        scope = f"회사 {company_id}" if company_id else "전체 회사"
        print(f"== 학습 리셋 완료: 총 {total} 행 삭제 (scope={scope}) ==")
    finally:
        await conn.close()


def main() -> None:
    p = argparse.ArgumentParser(description="learning 스키마 리셋 (카드/협상 데이터는 보존)")
    p.add_argument("--company", help="특정 company_id 만 리셋. 생략 시 전체")
    p.add_argument("--yes", action="store_true", help="확인 프롬프트 생략")
    args = p.parse_args()

    scope = f"회사 {args.company}" if args.company else "전체 회사"
    if not args.yes:
        ans = input(f"[{main_db_config.name}] learning 스키마({scope})를 비웁니다. 계속? [y/N] ")
        if ans.strip().lower() != "y":
            print("취소됨.")
            return
    asyncio.run(_reset(args.company))


if __name__ == "__main__":
    main()