app.py

import os
import json
import requests
import gradio as gr
from datetime import datetime
from dotenv import load_dotenv

# Load environment variables from .env file
load_dotenv()

# API Configuration - Hugging Face Spaces stores secrets in environment variables
API_KEY = os.getenv('StableCogKey', '')
if not API_KEY:
    # For local testing fallback
    API_KEY = "StableCogKey"

API_HOST = 'https://api.stablecog.com'
API_ENDPOINT = '/v1/credits'
API_URL = f'{API_HOST}{API_ENDPOINT}'

headers = {
    'Authorization': f'Bearer {API_KEY}',
    'Content-Type': 'application/json'
}

def check_credits():
    """Check StableCog credits and return formatted results"""
    try:
        response = requests.get(API_URL, headers=headers, timeout=10)
        
        if response.status_code == 200:
            res_json = response.json()
            
            # Debug: Print raw response to see structure
            print(f"Raw API Response: {json.dumps(res_json, indent=2)}")
            
            # Parse the actual response structure
            total_remaining_credits = res_json.get('total_remaining_credits', 0)
            credits_list = res_json.get('credits', [])
            
            # Calculate total initial credits from all credit entries
            total_initial_credits = 0
            total_used_credits = 0
            credit_details = []
            
            for credit in credits_list:
                credit_type = credit.get('type', {})
                initial_amount = credit_type.get('amount', 0)
                remaining_amount = credit.get('remaining_amount', 0)
                credit_name = credit_type.get('name', 'Unknown')
                
                total_initial_credits += initial_amount
                used_credits = initial_amount - remaining_amount
                total_used_credits += used_credits
                
                # Store credit details for display
                credit_details.append({
                    'name': credit_name,
                    'initial': initial_amount,
                    'remaining': remaining_amount,
                    'used': used_credits,
                    'expires': credit.get('expires_at', 'Never'),
                    'description': credit_type.get('description', '')
                })
            
            # Use total_remaining_credits from API or calculate it
            if total_remaining_credits == 0 and credits_list:
                # Calculate total remaining from all credit entries
                total_remaining_credits = sum(credit.get('remaining_amount', 0) for credit in credits_list)
            
            # Calculate total credits (max of initial total or used + remaining)
            total_credits = max(total_initial_credits, total_remaining_credits + total_used_credits)
            
            # Calculate percentage used
            if total_credits > 0:
                percentage_used = (total_used_credits / total_credits * 100)
            else:
                percentage_used = 0
            
            # Create formatted output
            timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S UTC")
            
            result = {
                "success": True,
                "total_credits": total_credits,
                "total_remaining_credits": total_remaining_credits,
                "total_used_credits": total_used_credits,
                "percentage_used": round(percentage_used, 2),
                "credit_details": credit_details,
                "total_credit_types": len(credits_list),
                "timestamp": timestamp,
                "raw_data": json.dumps(res_json, indent=2)
            }
            
            print(f"Parsed result: Total={total_credits}, Remaining={total_remaining_credits}, Used={total_used_credits}")
            return result
            
        else:
            return {
                "success": False,
                "error": f"API Error: {response.status_code}",
                "message": response.text if response.text else "No response text",
                "status_code": response.status_code
            }
            
    except requests.exceptions.Timeout:
        return {
            "success": False,
            "error": "Timeout Error",
            "message": "The request timed out. Please try again."
        }
    except requests.exceptions.ConnectionError:
        return {
            "success": False,
            "error": "Connection Error",
            "message": "Could not connect to the API. Check your internet connection."
        }
    except Exception as e:
        return {
            "success": False,
            "error": f"Unexpected Error: {type(e).__name__}",
            "message": str(e)
        }

def update_display():
    """Update the UI with credit information"""
    result = check_credits()
    
    if result["success"]:
        # Create a visual progress bar with color coding
        percentage = result["percentage_used"]
        if percentage < 50:
            bar_color = "#4CAF50"  # Green
            status = "🟢 Good"
            status_color = "#4CAF50"
        elif percentage < 80:
            bar_color = "#FF9800"  # Orange
            status = "🟡 Moderate"
            status_color = "#FF9800"
        else:
            bar_color = "#F44336"  # Red
            status = "🔴 Low"
            status_color = "#F44336"
        
        # Build credit details HTML
        credit_details_html = ""
        for i, credit in enumerate(result["credit_details"]):
            if credit['remaining'] > 0 or credit['initial'] > 0:
                credit_percentage = (credit['used'] / credit['initial'] * 100) if credit['initial'] > 0 else 0
                credit_details_html += f"""
                <div style='margin-bottom: 15px; padding: 12px; background: rgba(255,255,255,0.05); border-radius: 8px;'>
                    <div style='display: flex; justify-content: space-between; margin-bottom: 5px;'>
                        <span style='font-weight: bold;'>{credit['name']}</span>
                        <span>{credit['remaining']} / {credit['initial']} ⭐</span>
                    </div>
                    <div style='font-size: 12px; opacity: 0.8; margin-bottom: 8px;'>{credit['description']}</div>
                    <div style='height: 6px; background: rgba(255,255,255,0.1); border-radius: 3px; overflow: hidden;'>
                        <div style='height: 100%; width: {min(credit_percentage, 100)}%; background: {bar_color}; border-radius: 3px;'></div>
                    </div>
                </div>
                """
        
        if credit_details_html:
            credit_details_section = f"""
            <div style='margin-top: 20px;'>
                <h3 style='margin-bottom: 15px; font-size: 18px;'>📊 Credit Breakdown</h3>
                {credit_details_html}
            </div>
            """
        else:
            credit_details_section = ""
        
        html_content = f"""
        <div style='font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", Roboto, sans-serif; padding: 25px; background: linear-gradient(135deg, #667eea 0%, #764ba2 100%); border-radius: 15px; color: white; box-shadow: 0 10px 30px rgba(0,0,0,0.2);'>
            <h2 style='text-align: center; margin-bottom: 30px; font-size: 28px; font-weight: 600;'>🎨 StableCog Credit Status</h2>
            
            <div style='background: rgba(255,255,255,0.1); backdrop-filter: blur(10px); padding: 25px; border-radius: 12px; margin-bottom: 25px; border: 1px solid rgba(255,255,255,0.2);'>
                <div style='display: flex; justify-content: space-between; margin-bottom: 15px; padding-bottom: 10px; border-bottom: 1px solid rgba(255,255,255,0.1);'>
                    <span style='font-size: 16px; opacity: 0.9;'>Total Credits:</span>
                    <span style='font-weight: bold; font-size: 24px;'>{result['total_credits']} ⭐</span>
                </div>
                
                <div style='display: flex; justify-content: space-between; margin-bottom: 15px; padding-bottom: 10px; border-bottom: 1px solid rgba(255,255,255,0.1);'>
                    <span style='font-size: 16px; opacity: 0.9;'>Remaining Credits:</span>
                    <span style='font-weight: bold; font-size: 24px; color: #90EE90;'>{result['total_remaining_credits']} ⭐</span>
                </div>
                
                <div style='display: flex; justify-content: space-between; margin-bottom: 20px;'>
                    <span style='font-size: 16px; opacity: 0.9;'>Used Credits:</span>
                    <span style='font-weight: bold; font-size: 20px;'>{result['total_used_credits']} ⭐</span>
                </div>
                
                <div style='margin-bottom: 20px;'>
                    <div style='display: flex; justify-content: space-between; margin-bottom: 8px;'>
                        <span style='font-size: 16px; opacity: 0.9;'>Overall Usage:</span>
                        <span style='font-weight: bold;'>{result['percentage_used']}%</span>
                    </div>
                    <div style='height: 22px; background: rgba(255,255,255,0.15); border-radius: 11px; overflow: hidden; position: relative;'>
                        <div style='height: 100%; width: {min(result['percentage_used'], 100)}%; background: {bar_color}; 
                                border-radius: 11px; transition: width 0.5s ease-in-out; box-shadow: 0 0 10px {bar_color}80;'></div>
                        <div style='position: absolute; right: 10px; top: 50%; transform: translateY(-50%); color: white; font-size: 12px; font-weight: bold; text-shadow: 0 1px 2px rgba(0,0,0,0.5);'>
                            {result['percentage_used']}%
                        </div>
                    </div>
                </div>
                
                <div style='display: flex; justify-content: space-between; align-items: center; margin-top: 20px; padding-top: 20px; border-top: 1px solid rgba(255,255,255,0.2);'>
                    <span style='font-size: 16px; opacity: 0.9;'>Status:</span>
                    <span style='font-weight: bold; font-size: 18px; color: {status_color}; padding: 5px 15px; background: rgba(255,255,255,0.1); border-radius: 20px;'>
                        {status}
                    </span>
                </div>
                
                {credit_details_section}
            </div>
            
            <div style='text-align: center; font-size: 14px; opacity: 0.7; margin-top: 10px;'>
                ⏰ Last checked: {result['timestamp']} | 📋 Credit types: {result['total_credit_types']}
            </div>
        </div>
        """
        
        return html_content, result['raw_data'], result['total_remaining_credits'], result['percentage_used']
    
    else:
        # Error display
        html_content = f"""
        <div style='font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", Roboto, sans-serif; padding: 25px; background: linear-gradient(135deg, #f093fb 0%, #f5576c 100%); border-radius: 15px; color: white; text-align: center; box-shadow: 0 10px 30px rgba(0,0,0,0.2);'>
            <h2 style='margin-bottom: 20px; font-size: 26px;'>⚠️ API Connection Error</h2>
            
            <div style='background: rgba(255,255,255,0.15); backdrop-filter: blur(10px); padding: 20px; border-radius: 10px; margin-bottom: 20px; border: 1px solid rgba(255,255,255,0.2);'>
                <p style='margin: 0 0 10px 0; font-size: 18px; font-weight: bold;'>{result.get('error', 'Unknown error')}</p>
                <p style='margin: 0; font-size: 14px; opacity: 0.9;'>{result.get('message', '')}</p>
                {'<p style="margin: 10px 0 0 0; font-size: 14px;">Status Code: ' + str(result.get('status_code', '')) + '</p>' if result.get('status_code') else ''}
            </div>
            
            <div style='margin-top: 25px; padding: 15px; background: rgba(255,255,255,0.1); border-radius: 10px;'>
                <h3 style='margin-top: 0;'>🔧 Troubleshooting Tips:</h3>
                <ul style='text-align: left; margin: 10px 0; padding-left: 20px;'>
                    <li>Check if your API key is set in Hugging Face Secrets</li>
                    <li>Verify the API key has proper permissions</li>
                    <li>Ensure StableCog API is currently available</li>
                    <li>Check your internet connection</li>
                </ul>
            </div>
        </div>
        """
        
        return html_content, f"Error: {result.get('error', 'Unknown error')}\n\nDetails: {result.get('message', '')}", 0, 0

def get_recommendation(remaining_credits, percentage_used):
    """Provide recommendations based on credit status"""
    if remaining_credits == 0:
        return "💸 **No credits remaining.** Please add credits to continue using StableCog services."
    elif percentage_used >= 90:
        return "🛑 **Critically low credits!** Consider purchasing more credits before starting new projects."
    elif percentage_used >= 75:
        return "⚠️ **Credits are running low.** You can still do some work, but plan ahead for larger projects."
    elif remaining_credits < 10:
        return "📝 **Limited credits available.** Good for small tasks, testing, or single images."
    elif remaining_credits < 50:
        return "✨ **Credits available!** Suitable for several medium-sized projects or batch processing."
    else:
        return "🚀 **Plenty of credits!** Ready for extensive image generation work and experimentation."

# Create Gradio interface with Gradio 5 theme
theme = gr.themes.Soft(
    primary_hue="purple",
    secondary_hue="indigo",
    font=[gr.themes.GoogleFont("Inter"), "ui-sans-serif", "system-ui", "sans-serif"]
)

with gr.Blocks(
    theme=theme, 
    title="StableCog Credit Monitor",
    css="""
    footer {display: none !important;}
    .gradio-container {max-width: 1200px !important;}
    .stat-box input {font-weight: bold !important; font-size: 18px !important;}
    .recommendation-box textarea {font-size: 16px !important; line-height: 1.5 !important;}
    """
) as demo:
    gr.Markdown("""
    # 🎯 StableCog Credit Dashboard
    
    *Monitor your API credits and plan your image generation projects efficiently.*
    """)
    
    # Status row
    with gr.Row():
        with gr.Column(scale=2):
            html_output = gr.HTML(label="Credit Status")
        with gr.Column(scale=1):
            raw_output = gr.Code(
                label="📋 Raw API Response", 
                language="json", 
                interactive=False,
                lines=15
            )
    
    # Stats row
    with gr.Row():
        with gr.Column():
            credits_display = gr.Number(
                label="Remaining Credits", 
                interactive=False,
                elem_classes="stat-box"
            )
        with gr.Column():
            usage_display = gr.Number(
                label="Usage Percentage", 
                interactive=False,
                elem_classes="stat-box"
            )
    
    # Recommendation row
    with gr.Row():
        recommendation_box = gr.Textbox(
            label="🎯 AI Recommendation", 
            interactive=False,
            lines=3,
            elem_classes="recommendation-box"
        )
    
    # Control row
    with gr.Row():
        check_btn = gr.Button(
            "🔄 Check Credits Now", 
            variant="primary", 
            scale=1,
            size="lg"
        )
        gr.Button(
            "📖 View Documentation",
            variant="secondary",
            scale=0,
            size="lg",
            link="https://stablecog.com/docs/api"
        )
    
    # Auto-check on load
    demo.load(
        fn=update_display, 
        inputs=None, 
        outputs=[html_output, raw_output, credits_display, usage_display]
    ).then(
        fn=get_recommendation,
        inputs=[credits_display, usage_display],
        outputs=recommendation_box
    )
    
    # Connect button and update recommendation
    def update_all():
        html, raw, credits, usage = update_display()
        recommendation = get_recommendation(credits, usage)
        return html, raw, credits, usage, recommendation
    
    check_btn.click(
        fn=update_all, 
        inputs=None, 
        outputs=[html_output, raw_output, credits_display, usage_display, recommendation_box]
    )
    
    # Instructions and info
    with gr.Accordion("📚 How to Use & Setup", open=False):
        gr.Markdown("""
        ### Setting Up on Hugging Face Spaces:
        
        1. **Add your API key as a Secret:**
           - Go to your Space's Settings → Secrets
           - Add a new secret with:
             - Key: `STABLECOG_API_KEY`
             - Value: `your_actual_api_key_here`
        
        2. **Understanding Your Credits:**
           - **Total Credits**: Sum of all initial credits received
           - **Remaining Credits**: Currently available credits (sum of all credit types)
           - **Used Credits**: Credits spent on image generation
           - **Credit Types**: Different credit sources (Free, Refund, etc.)
        
        3. **Credit Breakdown:**
           - **Free Credits**: Base credits provided to users
           - **Refund Credits**: Credits returned for failed generations
           - Each credit type may have different amounts and expiration
        
        4. **Recommendations are based on:**
           - Remaining credit count
           - Usage percentage
           - Common workload patterns
        """)
    
    gr.Markdown("""
    ---
    *Built with ❤️ for StableCog users | [Report Issues](https://github.com/stability-ai/stablecog/issues)*
    """)

if __name__ == "__main__":
    # For Hugging Face Spaces - Gradio 5 syntax
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False,
        show_error=True,
        debug=False
    )