scripts/tools/quota_manager.py

import os
import psycopg2
import subprocess
from datetime import datetime, timezone

# 🔱 Imperial Gemini API 2026 Quota Limits (Gemini 2.5 / 3.1 Pro/Flash-Lite / etc.)
# 0 means Unlimited (무제한)
MODEL_LIMITS = {
    "gemini-3.1-flash-lite": {"rpd": 150000, "category": "텍스트 출력 모델", "display_name": "Gemini 3.1 Flash-Lite"},
    "gemini-3.1-pro": {"rpd": 250, "category": "텍스트 출력 모델", "display_name": "Gemini 3.1 Pro"},
    "gemini-3.0-flash": {"rpd": 10000, "category": "텍스트 출력 모델", "display_name": "Gemini 3.0 Flash"},
    "gemini-2.0-flash": {"rpd": 0, "category": "텍스트 출력 모델", "display_name": "Gemini 2.0 Flash"},
    "gemini-2.0-flash-lite": {"rpd": 0, "category": "텍스트 출력 모델", "display_name": "Gemini 2.0 Flash Lite"},
    "gemini-3.1-flash-tts-preview": {"rpd": 100, "category": "멀티모달 생성 모델", "display_name": "Gemini 3.1 Flash TTS"},
    "gemini-3-flash-live": {"rpd": 0, "category": "Live API", "display_name": "Gemini 3 Flash Live"},
    "imagen-4.0-generate-001": {"rpd": 70, "category": "멀티모달 생성 모델", "display_name": "Imagen 4.0 Pro"},
    "imagen-4.0-fast-generate-001": {"rpd": 70, "category": "멀티모달 생성 모델", "display_name": "Imagen 4.0 Fast"},
    "gemini-3.1-flash-image-preview": {"rpd": 1000, "category": "멀티모달 생성 모델", "display_name": "나노바나나 2"},
    "gemini-3-pro-image-preview": {"rpd": 250, "category": "멀티모달 생성 모델", "display_name": "나노바나나 Pro"},
    "gemini-embedding-2-preview": {"rpd": 0, "category": "기타 모델", "display_name": "Gemini Embedding 2"},
    "gemma-4-26b-it": {"rpd": 14400, "category": "기타 모델", "display_name": "Gemma 4 26B"},
    "gemma-4-31b-it": {"rpd": 14400, "category": "기타 모델", "display_name": "Gemma 4 31B"},
    "google-search-grounding-default": {"rpd": 5000, "category": "검색 그라우딩", "display_name": "Default"},
    "google-search-grounding-gemini2": {"rpd": 5000, "category": "검색 그라우딩", "display_name": "Gemini 2"},
    "google-search-grounding-gemini3": {"rpd": 1500, "category": "검색 그라우딩", "display_name": "Gemini 3"},
}

# 🔱 GitHub Copilot Model List with Premium Multipliers
# multiplier=0 means FREE (no premium requests consumed)
COPILOT_MONTHLY_BUDGET = 300  # GitHub Copilot Pro plan default
COPILOT_MODEL_INFO = {
    "claude-sonnet-4.6":           {"display_name": "Claude Sonnet 4.6",                   "multiplier": 1.0,  "category": "Anthropic"},
    "claude-sonnet-4.5":           {"display_name": "Claude Sonnet 4.5",                   "multiplier": 1.0,  "category": "Anthropic"},
    "claude-opus-4.6":             {"display_name": "Claude Opus 4.6",                     "multiplier": 3.0,  "category": "Anthropic"},
    "claude-opus-4.6-fast":        {"display_name": "Claude Opus 4.6 (fast mode)",         "multiplier": 30.0, "category": "Anthropic"},
    "claude-haiku-4.5":            {"display_name": "Claude Haiku 4.5",                    "multiplier": 0.33, "category": "Anthropic"},
    "gpt-4.1":                     {"display_name": "GPT-4.1",                             "multiplier": 0.0,  "category": "OpenAI"},
    "gpt-4o":                      {"display_name": "GPT-4o",                              "multiplier": 0.0,  "category": "OpenAI"},
    "gpt-4.1-mini":                {"display_name": "GPT-4.1 mini",                        "multiplier": 0.0,  "category": "OpenAI"},
    "gpt-4.1-nano":                {"display_name": "GPT-4.1 nano",                        "multiplier": 0.0,  "category": "OpenAI"},
    "gpt-5.4":                     {"display_name": "GPT-5.4",                             "multiplier": 1.0,  "category": "OpenAI"},
    "gpt-5.4-mini":                {"display_name": "GPT-5.4 mini",                        "multiplier": 0.33, "category": "OpenAI"},
    "gpt-5.1":                     {"display_name": "GPT-5.1",                             "multiplier": 1.0,  "category": "OpenAI"},
    "gpt-5.2":                     {"display_name": "GPT-5.2",                             "multiplier": 1.0,  "category": "OpenAI"},
    "gpt-5.3-codex":               {"display_name": "GPT-5.3-Codex",                       "multiplier": 1.0,  "category": "OpenAI"},
    "gpt-5.2-codex":               {"display_name": "GPT-5.2-Codex",                       "multiplier": 1.0,  "category": "OpenAI"},
    "gpt-5.1-codex":               {"display_name": "GPT-5.1-Codex",                       "multiplier": 1.0,  "category": "OpenAI"},
    "gpt-5.1-codex-max":           {"display_name": "GPT-5.1-Codex-Max",                   "multiplier": 1.0,  "category": "OpenAI"},
    "gpt-5.1-codex-mini":          {"display_name": "GPT-5.1-Codex-Mini",                  "multiplier": 0.33, "category": "OpenAI"},
    "gpt-5-mini":                  {"display_name": "GPT-5 mini",                          "multiplier": 0.0,  "category": "OpenAI"},
    "o1":                          {"display_name": "o1",                                  "multiplier": 1.0,  "category": "OpenAI"},
    "o1-mini":                     {"display_name": "o1-mini",                             "multiplier": 0.0,  "category": "OpenAI"},
    "o3-mini":                     {"display_name": "o3-mini",                             "multiplier": 0.0,  "category": "OpenAI"},
    "gemini-2.0-pro":              {"display_name": "Gemini 2.0 Pro",                      "multiplier": 1.0,  "category": "Google"},
    "gemini-3-flash":              {"display_name": "Gemini 3 Flash",                      "multiplier": 0.33, "category": "Google"},
    "gemini-3-pro":                {"display_name": "Gemini 3 Pro",                        "multiplier": 1.0,  "category": "Google"},
    "gemini-3.1-pro":              {"display_name": "Gemini 3.1 Pro",                      "multiplier": 1.0,  "category": "Google"},
    "grok-code-fast-1":            {"display_name": "Grok Code Fast 1",                    "multiplier": 0.25, "category": "xAI"},
    "raptor-mini":                 {"display_name": "Raptor mini",                         "multiplier": 0.0,  "category": "기타"},
}


class CopilotUsageManager:
    """GitHub Copilot API 사용량을 월간 기준으로 추적합니다."""
    DB_PREFIX = "copilot:"

    @staticmethod
    def get_db_conn():
        if os.environ.get("LOCAL_AGENT_MODE", "false").lower() == "true":
            return None
        return psycopg2.connect(_resolve_db_url(), connect_timeout=5)

    @staticmethod
    def _month_start():
        now = datetime.now(timezone.utc)
        return now.replace(day=1).date()

    @classmethod
    def track_usage(cls, model_id):
        db_key = cls.DB_PREFIX + model_id
        now_utc = datetime.now(timezone.utc)
        month_start = cls._month_start()
        conn = None
        cur = None
        try:
            conn = cls.get_db_conn()
            if not conn:
                return True
            cur = conn.cursor()
            cur.execute(
                "SELECT last_reset_date FROM taemingames.api_usage_tracker WHERE model_id = %s FOR UPDATE;",
                (db_key,)
            )
            row = cur.fetchone()
            if not row:
                cur.execute(
                    "INSERT INTO taemingames.api_usage_tracker (model_id, last_reset_date, daily_requests, daily_tokens) VALUES (%s, %s, 1, 0) ON CONFLICT DO NOTHING;",
                    (db_key, month_start)
                )
                cur.execute(
                    "UPDATE taemingames.api_usage_tracker SET last_request_at = %s WHERE model_id = %s;",
                    (now_utc, db_key)
                )
            else:
                last_reset = row[0]
                if last_reset < month_start:
                    cur.execute(
                        "UPDATE taemingames.api_usage_tracker SET daily_requests=1, last_request_at=%s, last_reset_date=%s WHERE model_id=%s;",
                        (now_utc, month_start, db_key)
                    )
                else:
                    cur.execute(
                        "UPDATE taemingames.api_usage_tracker SET daily_requests=daily_requests+1, last_request_at=%s WHERE model_id=%s;",
                        (now_utc, db_key)
                    )
            conn.commit()
            return True
        except Exception as e:
            print(f"❌ [CopilotUsageManager] track failed: {e}")
            return False
        finally:
            if cur:
                try: cur.close()
                except: pass
            if conn:
                try: conn.close()
                except: pass

    @classmethod
    def get_all_usage(cls, monthly_budget=None):
        if monthly_budget is None:
            monthly_budget = COPILOT_MONTHLY_BUDGET
        results = []
        conn = None
        cur = None
        try:
            prefix = cls.DB_PREFIX
            conn = cls.get_db_conn()
            if not conn:
                return {"models": [], "monthly_budget": monthly_budget, "total_premium_used": 0}
            cur = conn.cursor()
            cur.execute(
                "SELECT model_id, daily_requests, last_request_at FROM taemingames.api_usage_tracker WHERE model_id LIKE %s;",
                (prefix + "%",)
            )
            rows = {r[0].replace(prefix, ""): (r[1], r[2]) for r in cur.fetchall()}

            total_premium = 0.0
            for mid, info in COPILOT_MODEL_INFO.items():
                calls, last = rows.get(mid, (0, None))
                premium = round(calls * info["multiplier"], 2)
                total_premium += premium
                results.append({
                    "model_id": mid,
                    "display_name": info["display_name"],
                    "category": info["category"],
                    "multiplier": info["multiplier"],
                    "calls_this_month": calls,
                    "premium_used": premium,
                    "last_request_at": last.isoformat() if last else None,
                })
            return {
                "models": results,
                "monthly_budget": monthly_budget,
                "total_premium_used": round(total_premium, 2),
            }
        except Exception as e:
            print(f"❌ [CopilotUsageManager] get_all_usage failed: {e}")
            return {"models": [], "monthly_budget": monthly_budget, "total_premium_used": 0}
        finally:
            if cur:
                try: cur.close()
                except: pass
            if conn:
                try: conn.close()
                except: pass

def _sanitize_db_url(url: str) -> str:
    from urllib.parse import urlparse, urlunparse, urlencode, parse_qs
    url = url.strip().strip('"').strip("'").strip()
    parsed = urlparse(url)
    qs = parse_qs(parsed.query, keep_blank_values=True)
    qs.pop("schema", None)
    qs.pop("search_path", None)
    qs.pop("connection_limit", None)
    clean_query = urlencode(qs, doseq=True)
    return urlunparse(parsed._replace(query=clean_query))

def _resolve_db_url():
    url = os.environ.get("DATABASE_URL")
    if url: return _sanitize_db_url(url)
    try:
        result = subprocess.run(
            'bw get password "Imperial: LOCAL_DATABASE_URL" --raw',
            capture_output=True, text=True, shell=True, timeout=10,
            encoding='utf-8', errors='replace'
        )
        val = result.stdout.strip()
        if val and "postgresql://" in val:
            val = _sanitize_db_url(val)
            os.environ["DATABASE_URL"] = val
            return val
    except: pass
    return "postgresql://localhost:5432/postgres"

class QuotaExceededError(Exception):
    """API 쿼터(RPD)가 초과되었을 때 발생하는 제국 표준 예외입니다."""
    def __init__(self, model_id, limit):
        self.model_id = model_id
        self.limit = limit
        super().__init__(f"🔱 [황실 보안 경보] '{model_id}' 모델의 일일 사용 한도({limit} RPD)를 초과했습니다. 내일 다시 초기화됩니다.")

class ImperialQuotaManager:
    """Gemini API 사용량을 추적하고 한도를 강제하는 전담 관리자입니다."""
    
    @staticmethod
    def get_db_conn():
        if os.environ.get("LOCAL_AGENT_MODE", "false").lower() == "true":
            return None
        return psycopg2.connect(_resolve_db_url(), connect_timeout=5)

    @staticmethod
    def _get_utc_now():
        return datetime.now(timezone.utc)

    @classmethod
    def validate_quota(cls, model_id):
        """API 호출 전 RPD 한도를 체크하고 초과 시 QuotaExceededError를 발생시킵니다."""
        limit_data = MODEL_LIMITS.get(model_id)
        if not limit_data:
            return  # 미등록 모델은 통과
            
        rpd_limit = int(limit_data.get("rpd", 0))
        if rpd_limit <= 0:
            return  # 무제한(0 or -1) 통과
            
        conn = None
        cur = None
        try:
            conn = cls.get_db_conn()
            if not conn:
                return
            cur = conn.cursor()
            cur.execute(
                "SELECT daily_requests, last_reset_date FROM taemingames.api_usage_tracker WHERE model_id = %s;",
                (model_id,)
            )
            row = cur.fetchone()
            if not row:
                return  # 기록 없으면 첫 호출이므로 통과
                
            used_requests, last_reset = row
            today_utc = cls._get_utc_now().date()
            
            # 날짜가 오늘이면 실제 사용량 체크
            if last_reset == today_utc:
                if used_requests >= rpd_limit:
                    raise QuotaExceededError(model_id, rpd_limit)
        except QuotaExceededError:
            raise
        except Exception as e:
            print(f"[Quota Guard] Validation failed: {e}")
            return  # DB 장애 시 가동 유지를 위해 통과
        finally:
            if cur:
                try: cur.close()
                except: pass
            if conn:
                try: conn.close()
                except: pass

    @classmethod
    def track_usage(cls, model_id, tokens=0):
        """API 호출 시 사용량을 기록하고 일일 리셋을 처리합니다."""
        now_utc = cls._get_utc_now()
        today_utc = now_utc.date()
        
        conn = None
        cur = None
        try:
            conn = cls.get_db_conn()
            if not conn:
                return True
            cur = conn.cursor()
            
            # 1. 현재 상태 조회 및 리셋 체크
            cur.execute(
                "SELECT last_reset_date FROM taemingames.api_usage_tracker WHERE model_id = %s FOR UPDATE;",
                (model_id,)
            )
            row = cur.fetchone()
            
            if not row:
                # 모델이 없으면 생성
                cur.execute(
                    "INSERT INTO taemingames.api_usage_tracker (model_id, last_reset_date, daily_requests, daily_tokens) VALUES (%s, %s, 1, 0) ON CONFLICT DO NOTHING;",
                    (model_id, today_utc)
                )
                last_reset = today_utc
            else:
                last_reset = row[0]

            # 2. 날짜가 바뀌었으면 리셋 (구글 리셋 주기: 00:00 UTC)
            if last_reset != today_utc:
                cur.execute(
                    """UPDATE taemingames.api_usage_tracker 
                       SET daily_requests = 1, daily_tokens = 0, last_request_at = %s, last_reset_date = %s 
                       WHERE model_id = %s;""",
                    (now_utc, today_utc, model_id)
                )
            else:
                # 일반 누적 업데이트
                cur.execute(
                    """UPDATE taemingames.api_usage_tracker 
                       SET daily_requests = daily_requests + 1, daily_tokens = 0, last_request_at = %s 
                       WHERE model_id = %s;""",
                    (now_utc, model_id)
                )
            
            conn.commit()
            return True
        except Exception as e:
            print(f"❌ [QuotaManager] Failed to track usage for {model_id}: {e}")
            return False
        finally:
            if cur:
                try: cur.close()
                except: pass
            if conn:
                try: conn.close()
                except: pass

    @classmethod
    def get_all_usage_with_limits(cls):
        """모든 모델의 현재 사용량과 설정된 한도를 반환합니다. (사용량 0인 모델 포함)"""
        results = []
        conn = None
        cur = None
        try:
            conn = cls.get_db_conn()
            if not conn:
                # Return dummy limits when bypassing DB
                for mid, limit in MODEL_LIMITS.items():
                    results.append({
                        "model_id": mid,
                        "rpm": 0,
                        "tpm": 0,
                        "rpd": 0,
                        "last_request_at": None,
                        "limit": limit
                    })
                return results
                
            cur = conn.cursor()
            cur.execute("SELECT model_id, daily_requests, daily_tokens, last_request_at FROM taemingames.api_usage_tracker;")
            rows_data = cur.fetchall()
            rows_dict = {r[0]: (r[1], r[2], r[3]) for r in rows_data}

            # MODEL_LIMITS에 정의된 모든 모델을 순회하며 데이터가 없으면 0으로 채움
            for mid, limit in MODEL_LIMITS.items():
                req, tok, last = rows_dict.get(mid, (0, 0, None))
                results.append({
                    "model_id": mid,
                    "rpm": 0,       # 실시간 RPM은 DB에 누적하지 않으므로 0으로 표시 (한도는 limit에 포함)
                    "tpm": tok,     # Daily Tokens (UI에서 TPM으로 표시될 수 있으나 실제론 일일 누적)
                    "rpd": req,     # Daily Requests
                    "last_request_at": last.isoformat() if last else None,
                    "limit": limit
                })
            return results
        except Exception as e:
            print(f"[QuotaManager] ERROR: Failed to get usage: {e}")
            return []
        finally:
            if cur:
                try: cur.close()
                except: pass
            if conn:
                try: conn.close()
                except: pass

if __name__ == "__main__":
    # Test
    ImperialQuotaManager.track_usage("gemini-3.1-flash-lite", tokens=100)
    print(ImperialQuotaManager.get_all_usage_with_limits())