"""Offline Q-learning agent backed by a UCB policy."""

from __future__ import annotations

import os
from typing import Optional

import numpy as np

from ..model.visit_table import VisitTable
from ..model.q_table import QTable
from .policy import UCBPolicy


class QLearningAgent:
    def __init__(
        self,
        agent_params,
        state_size: int,
        action_size: int,
        visit_table: Optional[VisitTable] = None,
    ) -> None:
        """
        Args:
            agent_params: 에이전트 파라미터 (learning_rate, discount_factor 등)
            state_size: 상태 공간 크기
            action_size: 액션 공간 크기
            visit_table: 방문 기록 테이블 (None이면 새로 생성)
        """
        self.state_size = state_size
        self.action_size = action_size
        
        # Q-Table 객체 생성 (Composition)
        self.q_table = QTable(
            state_space_size=state_size,
            action_space_size=action_size,
            learning_rate=agent_params["learning_rate"],
            discount_factor=agent_params["discount_factor"]
        )

        self.visit_table = visit_table or VisitTable(state_size, action_size)
        self.policy = UCBPolicy(
            visit_table=self.visit_table,
            exploration_constant=agent_params.get("exploration_constant", np.sqrt(2.0)),
        )

    def get_action(self, state: int, action_mask=None):
        """
        현재 상태에서 액션 선택 (UCB 정책 사용)

        Args:
            state: 현재 상태 인덱스
            action_mask: 가능한 액션 마스크 (Optional)

        Returns:
            선택된 액션 ID
        """
        # QTable 객체에서 Q-value 조회
        q_values = self.q_table.get_q_values_for_state(state)
        mask = None if action_mask is None else np.asarray(action_mask, dtype=bool)
        return self.policy.select_action(state, q_values, available_mask=mask)

    def learn(self, batch):
        """
        배치 데이터를 사용하여 Q-Table 업데이트

        Args:
            batch: 학습 데이터 배치 (observations, actions, rewards, next_observations, terminals)
        """
        for state, action, reward, next_state, terminated in zip(
            batch["observations"],
            batch["actions"],
            batch["rewards"],
            batch["next_observations"],
            batch["terminals"],
        ):
            # QTable 객체의 update 메서드 사용
            # terminated가 True이면 next_state는 의미가 없거나 None이어야 함
            # QTable.update는 next_state_index가 None이면 종료 상태로 처리
            
            next_s = next_state if not terminated else None
            
            self.q_table.update(
                state_index=state,
                action_id=action,
                reward=reward,
                next_state_index=next_s,
                done=terminated
            )
            
            self.visit_table.increment(state, action)

    def save_model(self, path):
        """
        Q-Table을 파일로 저장

        Args:
            path: 저장할 파일 경로
        """
        # QTable 내부의 numpy array를 저장 (기존 호환성 유지)
        np.save(path, self.q_table.q_values)
        print(f"Q-Table saved to {path}")

    def load_q_table(self, file_path):
        """
        파일에서 Q-Table 로드

        Args:
            file_path: 로드할 파일 경로
        """
        if os.path.exists(file_path):
            # QTable 객체의 q_values 속성에 직접 할당
            self.q_table.q_values = np.load(file_path)
            print(f"Q-Table loaded from {file_path}")
        else:
            print(f"Error: No Q-Table found at {file_path}")

    def reset_episode(self):
        """에피소드 초기화 (정책 상태 등 리셋)"""
        self.policy.reset_episode()
