src/submission_tracker.py

"""
사용자 제출 추적 모듈
로그인한 사용자의 user_id를 기반으로 하루 3번 제한 기능을 제공합니다.
제출 정보는 별도의 HuggingFace repository에서 관리됩니다.
"""

import os
import json
import pandas as pd
import tempfile
from datetime import datetime
from typing import Dict, List, Optional, Tuple
from huggingface_hub import hf_hub_download, login, HfApi
import pytz

from src.utils import file_lock, get_current_date_str, get_current_datetime_str

# 한국 시간대 설정
KOREA_TZ = pytz.timezone('Asia/Seoul')


class SubmissionTracker:
    """사용자 제출 추적 클래스 - HuggingFace Repository 기반"""
    
    def __init__(
        self, 
        filename: str = "user_submissions.json",
    ):
        """
        Args:
            filename: 제출 기록 파일명
        """
        # 환경변수에서 설정 가져오기
        self.repo_id = os.getenv("SUBMISSION_TRACKER_REPO_ID")
        self.admin_token = os.getenv("HF_TOKEN")
        self.filename = filename
        
        if not self.repo_id:
            raise ValueError(
                "SUBMISSION_TRACKER_REPO_ID 환경변수가 설정되지 않았습니다. "
            )
        
        if not self.admin_token:
            raise ValueError(
                "HuggingFace Admin 토큰이 필요합니다. "
                "HF_TOKEN 환경변수를 설정하세요."
            )
        
        # HuggingFace API 초기화
        self.api = HfApi()
        try:
            # 관리자 토큰으로 로그인 (dataset read/write 용)
            login(token=self.admin_token)
        except Exception as e:
            print(f"❌ HuggingFace 로그인 실패: {e}")
            raise
        
        # 제출 기록 로드
        self.submissions: Dict = self.load_submissions()
    
    def load_submissions(self) -> Dict:
        """HuggingFace repository에서 제출 기록 로드"""
        try:
            # 임시 디렉토리에 파일 다운로드
            with tempfile.TemporaryDirectory() as temp_dir:
                file_path = hf_hub_download(
                    repo_id=self.repo_id,
                    filename=self.filename,
                    local_dir=temp_dir,
                    repo_type="dataset",
                    token=self.admin_token,
                )
                
                # JSON 파일 로드
                with open(file_path, "r", encoding="utf-8") as f:
                    submissions = json.load(f)
                
                return submissions
                
        except Exception as e:
            print(f"⚠️ 제출 기록 로드 실패 (새로 시작): {e}")
            return {}

    def get_today_submissions(self, user_id: str) -> List[Dict]:
        """오늘 사용자의 제출 기록 가져오기"""
        if not user_id:
            return []
        today = get_current_date_str()
        user_submissions = self.submissions.get(user_id, {})
        return user_submissions.get(today, [])
    
    def can_submit(
        self,
        user_id: str,
        submissions_data: Optional[Dict] = None,
    ) -> Tuple[bool, str, int]:
        """
        사용자가 제출할 수 있는지 확인.

        Args:
            user_id: 로그인한 사용자의 고유 ID (HF 계정 ID 등)
            submissions_data: 검사에 사용할 제출 데이터(테스트/락 내부 재검사용).
                              None이면 self.submissions 사용.
        """
        if not user_id:
            raise ValueError("❌ HuggingFace 로그인 상태에서만 제출 가능합니다. 로그인 정보를 확인해주세요.")
        
        data = submissions_data if submissions_data is not None else self.submissions
        today = get_current_date_str()
        today_submissions = data.get(user_id, {}).get(today, [])
        successful_count = len([s for s in today_submissions if s.get("success", False)])

        if successful_count >= 3:
            raise Exception("❌ 오늘 제출 한도를 초과했습니다. 내일 다시 시도해주세요.")

        remaining = 3 - successful_count
        return True, f"✅ 제출 가능합니다. (오늘 {successful_count}/3회 사용, {remaining}회 남음)", remaining
    
    def record_submission(
        self,
        user_id: str,
        submitter_name: str,
        file_name: str,
        success: bool,
        error_message: Optional[str] = None,
        submit_model: Optional[str] = None,
        submit_description: Optional[str] = None,
    ) -> bool:
        """
        제출 기록 추가 (파일 잠금으로 보호)

        Args:
            user_id: 로그인한 사용자의 고유 ID (HF 계정 ID 등)
        """
        if not user_id:
            raise ValueError("❌ HuggingFace 로그인 상태에서만 제출 가능합니다. 로그인 정보를 확인해주세요.")
        
        # 잠금 파일 경로 생성
        lock_file_path = tempfile.gettempdir() + f'/{self.repo_id.replace("/", "_")}.lock'
        
        # 파일 잠금으로 전체 과정을 atomic하게 보호
        with file_lock(lock_file_path):
            try:
                # 최신 데이터를 다시 로드 (다른 프로세스에서 업데이트했을 수 있음)
                latest_submissions = self.load_submissions()
                
                # Lock 내부에서 최신 데이터 기준으로 제출 가능 여부 재확인
                try:
                    self.can_submit(
                        user_id=user_id,
                        submissions_data=latest_submissions,
                    )
                except Exception as e:
                    # 제출 제한 초과 시
                    print(f"제출 제한 초과: {e}")
                    # 메모리만 최신으로 맞추고 저장하지 않음
                    self.submissions = latest_submissions
                    return False
                
                # 새로운 제출 기록 추가
                current_datetime = get_current_datetime_str()
                
                if user_id not in latest_submissions:
                    latest_submissions[user_id] = {}
                
                today = get_current_date_str()
                if today not in latest_submissions[user_id]:
                    latest_submissions[user_id][today] = []
                
                submission_record = {
                    "timestamp": current_datetime,
                    "submitter_name": submitter_name,
                    "file_name": file_name,
                    "success": success,
                    "error_message": error_message,
                    "submit_model": submit_model,
                    "submit_description": submit_description,
                }
                
                latest_submissions[user_id][today].append(submission_record)
                
                # 메모리 업데이트
                self.submissions = latest_submissions
                
                # 저장
                return self._save_submissions_internal(latest_submissions)
                
            except Exception as e:
                print(f"❌ 제출 기록 추가 실패: {e}")
                return False
    
    def _save_submissions_internal(self, submissions_data: Dict) -> bool:
        """내부 저장 함수 (lock은 이미 획득된 상태)"""
        try:
            # 임시 파일에 JSON 데이터 저장
            with tempfile.NamedTemporaryFile(
                mode="w",
                encoding="utf-8",
                suffix=".json",
                delete=False,
            ) as temp_file:
                json.dump(submissions_data, temp_file, ensure_ascii=False, indent=2)
                temp_file_path = temp_file.name
            
            # HuggingFace repository에 파일 업로드
            self.api.upload_file(
                path_or_fileobj=temp_file_path,
                path_in_repo=self.filename,
                repo_id=self.repo_id,
                repo_type="dataset",
                token=self.admin_token,
                commit_message=(
                    "Update submission records - "
                    f"{datetime.now(KOREA_TZ).strftime('%Y-%m-%d %H:%M:%S')}"
                ),
            )
            
            # 임시 파일 삭제
            os.unlink(temp_file_path)
            
            return True
            
        except Exception as e:
            print(f"❌ 제출 기록 저장 실패: {e}")
            return False
    
    def get_user_submission_history(self, user_id: str, days: int = 7) -> Dict:
        """사용자의 최근 제출 기록 가져오기"""
        if not user_id or user_id not in self.submissions:
            return {}
        
        user_submissions = self.submissions[user_id]
        today = datetime.now(KOREA_TZ).date()
        
        history: Dict[str, List[Dict]] = {}
        for i in range(days):
            check_date = today - pd.Timedelta(days=i)
            date_str = check_date.strftime("%Y-%m-%d")
            
            if date_str in user_submissions:
                history[date_str] = user_submissions[date_str]
        
        return history
    
    def get_submission_stats(self, user_id: str) -> Dict:
        """사용자 제출 통계 가져오기"""
        if not user_id:
            return {}
        
        today_submissions = self.get_today_submissions(user_id)
        successful_today_count = len([s for s in today_submissions if s.get("success", False)])
        history = self.get_user_submission_history(user_id, 7)
        
        # 통계 계산
        total_submissions = sum(len(day_submissions) for day_submissions in history.values())
        successful_submissions = sum(
            len([s for s in day_submissions if s.get("success", False)])
            for day_submissions in history.values()
        )
        failed_submissions = total_submissions - successful_submissions
        
        return {
            "today_count": len(today_submissions),
            "today_remaining": max(0, 3 - successful_today_count),
            "week_total": total_submissions,
            "week_successful": successful_submissions,
            "week_failed": failed_submissions,
            "history": history,
        }
    
    def cleanup_old_records(self, days_to_keep: int = 30) -> int:
        """오래된 제출 기록 정리 (파일 잠금 사용)"""
        # 잠금 파일 경로 생성
        lock_file_path = tempfile.gettempdir() + f'/{self.repo_id.replace("/", "_")}.lock'
        
        # 파일 잠금으로 전체 과정을 atomic하게 보호
        with file_lock(lock_file_path):
            try:
                # 최신 데이터를 다시 로드
                latest_submissions = self.load_submissions()
                
                cutoff_date = datetime.now(KOREA_TZ) - pd.Timedelta(days=days_to_keep)
                cutoff_str = cutoff_date.strftime("%Y-%m-%d")
                
                cleaned_count = 0
                for uid in list(latest_submissions.keys()):
                    user_submissions = latest_submissions[uid]
                    for date_str in list(user_submissions.keys()):
                        if date_str < cutoff_str:
                            del user_submissions[date_str]
                            cleaned_count += 1
                    
                    # 빈 사용자 기록 제거
                    if not user_submissions:
                        del latest_submissions[uid]
                
                # 메모리 업데이트
                self.submissions = latest_submissions
                
                if cleaned_count > 0:
                    if self._save_submissions_internal(latest_submissions):
                        print(f"🧹 {cleaned_count}개의 오래된 제출 기록을 정리했습니다.")
                    else:
                        print(f"⚠️ {cleaned_count}개의 오래된 제출 기록을 정리했지만 저장에 실패했습니다.")
                
                return cleaned_count
                
            except Exception as e:
                print(f"❌ 오래된 기록 정리 실패: {e}")
                return 0


def get_submission_tracker() -> Optional[SubmissionTracker]:
    """SubmissionTracker 인스턴스 반환"""
    try:
        return SubmissionTracker()
    except Exception as e:
        print(f"❌ SubmissionTracker 초기화 실패: {e}")
        return None