"""
CollectExperienceUsecase
협상 중 Experience 데이터를 수집하여 저장
"""

from typing import Optional
from datetime import datetime
from ..domain.model.state import State
from ..domain.model.experience import Experience
from ..domain.repository.experience_repository import ExperienceRepository


class CollectExperienceUsecase:
    """
    협상 과정에서 발생하는 Experience를 수집하여 저장

    오프라인 학습을 위한 데이터 수집 담당
    """

    def __init__(
        self,
        experience_repository: ExperienceRepository,
        filename: Optional[str] = None,
    ):
        """
        Args:
            experience_repository: Experience 저장소
            filename: 저장할 파일명 (None이면 기본 파일 사용)
        """
        self.experience_repository = experience_repository
        self.filename = filename or "experiences.jsonl"
        self.current_episode_id: Optional[str] = None
        self.current_step: int = 0

    def start_episode(self, episode_id: Optional[str] = None):
        """
        새로운 에피소드 시작

        Args:
            episode_id: 에피소드 ID (None이면 자동 생성)
        """
        if episode_id is None:
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S_%f")
            episode_id = f"ep_{timestamp}"

        self.current_episode_id = episode_id
        self.current_step = 0

    def collect(
        self,
        state: State,
        action_id: int,
        reward: float,
        next_state: Optional[State],
        done: bool,
    ) -> Experience:
        """
        단일 Experience 수집 및 저장

        Args:
            state: 현재 상태
            action_id: 선택한 액션 ID
            reward: 받은 보상
            next_state: 다음 상태 (종료 시 None)
            done: 에피소드 종료 여부

        Returns:
            수집된 Experience
        """
        if self.current_episode_id is None:
            self.start_episode()

        # Experience 생성
        experience = Experience(
            state=state,
            action_id=action_id,
            reward=reward,
            next_state=next_state,
            done=done,
            episode_id=self.current_episode_id,
            step=self.current_step,
            timestamp=datetime.now().isoformat(),
        )

        # 저장
        self.experience_repository.save(experience, self.filename)

        # 스텝 증가
        self.current_step += 1

        return experience

    def end_episode(self):
        """에피소드 종료"""
        self.current_episode_id = None
        self.current_step = 0

    def get_episode_count(self) -> int:
        """수집된 에피소드 개수 (근사치)"""
        # 파일에서 unique episode_id 개수 계산
        experiences = self.experience_repository.load_all(self.filename)
        episode_ids = set(exp.episode_id for exp in experiences if exp.episode_id)
        return len(episode_ids)

    def get_total_count(self) -> int:
        """수집된 총 Experience 개수"""
        return self.experience_repository.count(self.filename)

    def get_collection_info(self) -> dict:
        """
        수집 정보 조회

        Returns:
            {
                'filename': 'experiences.jsonl',
                'total_experiences': 1000,
                'episodes': 50,
                'current_episode': 'ep_001' or None,
                'current_step': 5
            }
        """
        return {
            "filename": self.filename,
            "total_experiences": self.get_total_count(),
            "episodes": self.get_episode_count(),
            "current_episode": self.current_episode_id,
            "current_step": self.current_step,
            "file_info": self.experience_repository.get_file_info(self.filename),
        }

    def create_new_collection_file(self, prefix: str = "exp") -> str:
        """
        새로운 수집 파일 생성 및 전환

        Args:
            prefix: 파일명 접두사

        Returns:
            생성된 파일명
        """
        new_filename = self.experience_repository.create_new_file(prefix)
        self.filename = new_filename
        return new_filename
