#!/usr/bin/env python3
"""
토큰 사용량 추적 시스템
일일 토큰 예산 관리
"""

import json
from datetime import datetime
from pathlib import Path
from typing import Dict, Optional


class TokenTracker:
    """토큰 사용량 추적기"""

    def __init__(self, memory_path: str = None):
        """초기화"""
        if memory_path is None:
            memory_path = Path(__file__).parent
        self.memory_path = Path(memory_path)
        self.tracker_file = self.memory_path / "token-usage.json"
        self.daily_limit = 100000  # 일일 한도

        # 오늘 날짜
        self.today = datetime.now().strftime("%Y-%m-%d")

        # 데이터 로드
        self.data = self.load_data()

    def load_data(self) -> Dict:
        """데이터 로드"""
        if self.tracker_file.exists():
            with open(self.tracker_file, "r", encoding="utf-8") as f:
                return json.load(f)
        else:
            return {"daily": {}, "history": []}

    def save_data(self):
        """데이터 저장"""
        with open(self.tracker_file, "w", encoding="utf-8") as f:
            json.dump(self.data, f, ensure_ascii=False, indent=2)

    def track(self, tokens: int, category: str = "general") -> Dict:
        """
        토큰 사용량 기록

        Args:
            tokens: 사용한 토큰 수
            category: 카테고리 (search, generate, boris, delegation, etc.)

        Returns:
            상태 정보
        """
        # 오늘 데이터 초기화
        if self.today not in self.data["daily"]:
            self.data["daily"][self.today] = {"total": 0, "categories": {}}

        # 사용량 추가
        self.data["daily"][self.today]["total"] += tokens

        # 카테고리별 기록
        if category not in self.data["daily"][self.today]["categories"]:
            self.data["daily"][self.today]["categories"][category] = 0
        self.data["daily"][self.today]["categories"][category] += tokens

        # 저장
        self.save_data()

        # 상태 반환
        return self.status()

    def status(self) -> Dict:
        """
        현재 상태 반환

        Returns:
            상태 정보
        """
        if self.today not in self.data["daily"]:
            return {"date": self.today, "used": 0, "remaining": self.daily_limit, "percentage": 0.0, "categories": {}}

        used = self.data["daily"][self.today]["total"]
        remaining = self.daily_limit - used
        percentage = (used / self.daily_limit) * 100

        return {
            "date": self.today,
            "used": used,
            "remaining": remaining,
            "percentage": round(percentage, 2),
            "categories": self.data["daily"][self.today]["categories"],
        }

    def report(self) -> str:
        """
        사용량 보고서 생성

        Returns:
            보고서 문자열
        """
        status = self.status()

        report = []
        report.append("=" * 50)
        report.append("📊 토큰 사용량 보고서")
        report.append("=" * 50)
        report.append(f"📅 날짜: {status['date']}")
        report.append(f"✅ 사용: {status['used']:,} tokens")
        report.append(f"⏳ 남음: {status['remaining']:,} tokens")
        report.append(f"📈 비율: {status['percentage']}%")
        report.append("")

        if status["categories"]:
            report.append("📋 카테고리별 사용량:")
            for cat, count in status["categories"].items():
                report.append(f"  - {cat}: {count:,} tokens")

        report.append("=" * 50)

        # 경고
        if status["percentage"] > 80:
            report.append("⚠️ 주의: 일일 예산의 80% 이상 사용")

        if status["percentage"] > 95:
            report.append("🚨 경고: 일일 예산 거의 소진")

        return "\n".join(report)

    def history(self, days: int = 7) -> str:
        """
        지난 N일간 사용량 이력

        Args:
            days: 조회할 일수

        Returns:
            이력 문자열
        """
        report = []
        report.append("=" * 50)
        report.append(f"📊 지난 {days}일간 토큰 사용량")
        report.append("=" * 50)

        # 지난 N일간 데이터 수집
        dates = sorted(self.data["daily"].keys(), reverse=True)[:days]

        for date in dates:
            used = self.data["daily"][date]["total"]
            percentage = (used / self.daily_limit) * 100
            report.append(f"{date}: {used:,} tokens ({percentage:.1f}%)")

        report.append("=" * 50)

        return "\n".join(report)


# CLI 인터페이스
if __name__ == "__main__":
    import sys

    tracker = TokenTracker()

    if len(sys.argv) < 2:
        # 기본: 상태 보고
        print(tracker.report())
    else:
        command = sys.argv[1]

        if command == "track":
            # 토큰 기록
            if len(sys.argv) < 3:
                print("사용법: python3 token_tracker.py track <토큰수> [카테고리]")
                sys.exit(1)

            tokens = int(sys.argv[2])
            category = sys.argv[3] if len(sys.argv) > 3 else "general"

            status = tracker.track(tokens, category)
            print(f"✅ 기록됨: {tokens:,} tokens ({category})")
            print(f"📊 현재: {status['used']:,}/{status['remaining']:,} ({status['percentage']}%)")

        elif command == "status":
            # 상태 확인
            print(tracker.report())

        elif command == "history":
            # 이력 확인
            days = int(sys.argv[2]) if len(sys.argv) > 2 else 7
            print(tracker.history(days))

        else:
            print("알 수 없는 명령:", command)
            print("사용법:")
            print("  python3 token_tracker.py status       # 현재 상태")
            print("  python3 token_tracker.py track <n>    # 토큰 기록")
            print("  python3 token_tracker.py history [n]  # 이력 확인")