default/telegram-agy-bot/usage_monitor.py

import os
import re
import time
import asyncio
import logging
from datetime import datetime, timezone
from typing import Dict, Any, List, Optional, Tuple

logger = logging.getLogger("AGYUsageMonitor")

def make_quota_bar(percent_remaining: float, length: int = 10) -> str:
    """
    Generates an emoji progress bar for remaining quota:
    >= 60%: 🟩 Green (Safe/High)
    20% - 59%: 🟨 Yellow (Medium)
    < 20%: 🟥 Red (Low)
    """
    percent = max(0.0, min(100.0, float(percent_remaining)))
    filled = int(round((percent / 100.0) * length))
    filled = max(0, min(length, filled))

    if percent >= 60:
        fill_char = "🟩"
    elif percent >= 20:
        fill_char = "🟨"
    else:
        fill_char = "🟥"

    empty_char = "⬜"
    return f"{fill_char * filled}{empty_char * (length - filled)}"

def parse_iso_time_remaining(iso_str: str, is_fa: bool = True) -> str:
    """Calculates relative time remaining until reset timestamp."""
    if not iso_str:
        return "نامشخص" if is_fa else "Unknown"
    try:
        clean_iso = iso_str.strip().replace("Z", "+00:00")
        target_dt = datetime.fromisoformat(clean_iso)
        now_dt = datetime.now(timezone.utc)
        diff_sec = int((target_dt - now_dt).total_seconds())

        if diff_sec <= 0:
            return "هم‌اکنون در حال تمدید" if is_fa else "Resetting now"

        days = diff_sec // 86400
        hours = (diff_sec % 86400) // 3600
        mins = (diff_sec % 3600) // 60

        if is_fa:
            parts = []
            if days > 0:
                parts.append(f"{days} روز")
            if hours > 0:
                parts.append(f"{hours} ساعت")
            if mins > 0 or not parts:
                parts.append(f"{mins} دقیقه")
            return f"{' و '.join(parts)} دیگر"
        else:
            parts = []
            if days > 0:
                parts.append(f"{days}d")
            if hours > 0:
                parts.append(f"{hours}h")
            parts.append(f"{mins}m")
            return f"in {' '.join(parts)}"
    except Exception as e:
        return iso_str

def parse_agy_usage_output(usage_raw: str, credits_raw: str) -> Dict[str, Any]:
    """Parses raw stdout from 'agy --print /usage' and 'agy --print /credits'."""
    data: Dict[str, Any] = {
        "gemini_5h": None,
        "gemini_5h_reset": "",
        "gemini_weekly": None,
        "gemini_weekly_reset": "",
        "claude_5h": None,
        "claude_5h_reset": "",
        "claude_weekly": None,
        "claude_weekly_reset": "",
        "credits": "0",
        "raw_text": usage_raw + "\n" + credits_raw,
    }

    lines = (usage_raw + "\n" + credits_raw).splitlines()
    for line in lines:
        line_clean = line.strip()
        if not line_clean:
            continue

        # Check Gemini 5H
        if re.search(r"Gemini.*Five Hour Limit Remaining", line_clean, re.I):
            m = re.search(r"(\d+)%\s+(\S+)", line_clean)
            if m:
                data["gemini_5h"] = int(m.group(1))
                data["gemini_5h_reset"] = m.group(2)

        # Check Gemini Weekly
        elif re.search(r"Gemini.*Weekly Limit Remaining", line_clean, re.I):
            m = re.search(r"(\d+)%\s+(\S+)", line_clean)
            if m:
                data["gemini_weekly"] = int(m.group(1))
                data["gemini_weekly_reset"] = m.group(2)

        # Check Claude/GPT 5H
        elif re.search(r"(Claude|GPT).*Five Hour Limit Remaining", line_clean, re.I):
            m = re.search(r"(\d+)%\s+(\S+)", line_clean)
            if m:
                data["claude_5h"] = int(m.group(1))
                data["claude_5h_reset"] = m.group(2)

        # Check Claude/GPT Weekly
        elif re.search(r"(Claude|GPT).*Weekly Limit Remaining", line_clean, re.I):
            m = re.search(r"(\d+)%\s+(\S+)", line_clean)
            if m:
                data["claude_weekly"] = int(m.group(1))
                data["claude_weekly_reset"] = m.group(2)

        # Check credits
        elif re.search(r"Remaining credits", line_clean, re.I):
            m = re.search(r"Remaining credits\s+(\d+)", line_clean, re.I)
            if m:
                data["credits"] = m.group(1)

    return data

async def fetch_and_render_usage_report(is_fa: bool = True, user_id: Optional[int] = None) -> str:
    """Executes agy /usage and /credits and formats with visual progress bars."""
    from auth_manager import auth_manager
    if user_id and not auth_manager.has_custom_account(user_id):
        if is_fa:
            return (
                "⚠️ <b>حساب کاربری هوش مصنوعی متصل نیست!</b>\n\n"
                "در این ربات حساب‌های اشتراکی غیرفعال هستند و هر کاربر باید حساب گوگل اختصاصی خود را متصل کند.\n\n"
                "👉 برای ورود و استعلام سهمیه اختصاصی، لطفاً دستور <code>/login</code> را ارسال کنید یا از منوی <code>/account</code> وارد شوید."
            )
        else:
            return (
                "⚠️ <b>AGY Account Not Connected!</b>\n\n"
                "Shared accounts are disabled. Please log in with your own Google account using <code>/login</code> or <code>/account</code>."
            )

    env = auth_manager.get_user_env(user_id) if user_id else os.environ.copy()
    usage_raw = ""
    credits_raw = ""
    try:
        proc1 = await asyncio.create_subprocess_exec(
            "agy", "--print", "/usage",
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
            env=env,
        )
        out1, _ = await proc1.communicate()
        usage_raw = out1.decode("utf-8", errors="replace").strip()

        proc2 = await asyncio.create_subprocess_exec(
            "agy", "--print", "/credits",
            stdout=asyncio.subprocess.PIPE,
            stderr=asyncio.subprocess.PIPE,
            env=env,
        )
        out2, _ = await proc2.communicate()
        credits_raw = out2.decode("utf-8", errors="replace").strip()
    except Exception as e:
        logger.error(f"Failed to query AGY quota: {e}")
        usage_raw = f"Error: {e}"

    parsed = parse_agy_usage_output(usage_raw, credits_raw)

    g_5h = parsed["gemini_5h"]
    g_wk = parsed["gemini_weekly"]
    c_5h = parsed["claude_5h"]
    c_wk = parsed["claude_weekly"]
    credits_val = parsed["credits"]

    if g_5h is None and g_wk is None and c_5h is None:
        # Fallback to pre-formatted raw output if parsing failed
        if is_fa:
            return (
                "📊 <b>گزارش لحظه‌ای سهمیه و وضعیت مصرف AGY</b>\n\n"
                f"<pre>{usage_raw}\n\n{credits_raw}</pre>\n\n"
                "💡 <i>سهمیه‌های ۵ ساعته و هفتگی به صورت خودکار تمدید می‌شوند.</i>"
            )
        else:
            return (
                "📊 <b>Real-time AGY Quota & Usage Report</b>\n\n"
                f"<pre>{usage_raw}\n\n{credits_raw}</pre>\n\n"
                "💡 <i>Quotas are reset on 5-hour and 7-day windows.</i>"
            )

    if is_fa:
        lines = [
            "📈 <b>گزارش وضعیت سهمیه و مصرف هوش مصنوعی (AGY)</b>",
            "",
            "✨ <b>۱. مدل‌های Gemini (جمینای):</b>",
        ]
        if g_5h is not None:
            bar_5h = make_quota_bar(g_5h)
            reset_5h = parse_iso_time_remaining(parsed["gemini_5h_reset"], is_fa=True)
            lines.extend([
                f"  ⏳ <b>سهمیه ۵ ساعته:</b>",
                f"  {bar_5h} <b>{g_5h}% باقیمانده</b>",
                f"  • 🔄 <i>زمان تمدید:</i> <code>{reset_5h}</code>",
            ])
        if g_wk is not None:
            bar_wk = make_quota_bar(g_wk)
            reset_wk = parse_iso_time_remaining(parsed["gemini_weekly_reset"], is_fa=True)
            lines.extend([
                f"  📅 <b>سهمیه هفتگی:</b>",
                f"  {bar_wk} <b>{g_wk}% باقیمانده</b>",
                f"  • 🔄 <i>زمان تمدید:</i> <code>{reset_wk}</code>",
            ])

        lines.extend([
            "",
            "🌟 <b>۲. مدل‌های Claude و GPT:</b>",
        ]
        )
        if c_5h is not None:
            bar_c5 = make_quota_bar(c_5h)
            reset_c5 = parse_iso_time_remaining(parsed["claude_5h_reset"], is_fa=True)
            lines.extend([
                f"  ⏳ <b>سهمیه ۵ ساعته:</b>",
                f"  {bar_c5} <b>{c_5h}% باقیمانده</b>",
                f"  • 🔄 <i>زمان تمدید:</i> <code>{reset_c5}</code>",
            ])
        if c_wk is not None:
            bar_cwk = make_quota_bar(c_wk)
            reset_cwk = parse_iso_time_remaining(parsed["claude_weekly_reset"], is_fa=True)
            lines.extend([
                f"  📅 <b>سهمیه هفتگی:</b>",
                f"  {bar_cwk} <b>{c_wk}% باقیمانده</b>",
                f"  • 🔄 <i>زمان تمدید:</i> <code>{reset_cwk}</code>",
            ])

        lines.extend([
            "",
            f"🪙 <b>اعتبار مدل‌ها (Model Credits):</b> <code>{credits_val}</code>",
            "",
            "💡 <i>سهمیه‌های مصرف به صورت خودکار در بازه‌های زمانی فوق بازنشانی می‌شوند.</i>"
        ])
    else:
        lines = [
            "📈 <b>Real-Time AGY Quota & Usage Report</b>",
            "",
            "✨ <b>1. Gemini Models:</b>",
        ]
        if g_5h is not None:
            bar_5h = make_quota_bar(g_5h)
            reset_5h = parse_iso_time_remaining(parsed["gemini_5h_reset"], is_fa=False)
            lines.extend([
                f"  ⏳ <b>5-Hour Limit:</b>",
                f"  {bar_5h} <b>{g_5h}% Remaining</b>",
                f"  • 🔄 <i>Reset:</i> <code>{reset_5h}</code>",
            ])
        if g_wk is not None:
            bar_wk = make_quota_bar(g_wk)
            reset_wk = parse_iso_time_remaining(parsed["gemini_weekly_reset"], is_fa=False)
            lines.extend([
                f"  📅 <b>Weekly Limit:</b>",
                f"  {bar_wk} <b>{g_wk}% Remaining</b>",
                f"  • 🔄 <i>Reset:</i> <code>{reset_wk}</code>",
            ])

        lines.extend([
            "",
            "🌟 <b>2. Claude & GPT Models:</b>",
        ])
        if c_5h is not None:
            bar_c5 = make_quota_bar(c_5h)
            reset_c5 = parse_iso_time_remaining(parsed["claude_5h_reset"], is_fa=False)
            lines.extend([
                f"  ⏳ <b>5-Hour Limit:</b>",
                f"  {bar_c5} <b>{c_5h}% Remaining</b>",
                f"  • 🔄 <i>Reset:</i> <code>{reset_c5}</code>",
            ])
        if c_wk is not None:
            bar_cwk = make_quota_bar(c_wk)
            reset_cwk = parse_iso_time_remaining(parsed["claude_weekly_reset"], is_fa=False)
            lines.extend([
                f"  📅 <b>Weekly Limit:</b>",
                f"  {bar_cwk} <b>{c_wk}% Remaining</b>",
                f"  • 🔄 <i>Reset:</i> <code>{reset_cwk}</code>",
            ])

        lines.extend([
            "",
            f"🪙 <b>Model Credits:</b> <code>{credits_val}</code>",
            "",
            "💡 <i>Quotas automatically reset periodically in the windows shown above.</i>"
        ])

    return "\n".join(lines)