import json
import time
import os
import redis
from datetime import datetime, timedelta

class BudgetManager:
    """
    Tracks and gates AI costs to stay within a specific budget.
    Default: $10 per 24-hour cycle.
    """
    def __init__(self, daily_limit: float = 10.0):
        self.daily_limit = daily_limit
        self.namespace = "genesis:budget"
        self.redis_client = None
        
        # Load Redis config
        config_path = r"E:\genesis-system\genesis_config.json"
        if os.path.exists(config_path):
            with open(config_path, 'r') as f:
                config = json.load(f).get("redis")
                if config:
                    self.redis_client = redis.Redis(
                        host=config["host"],
                        port=config["port"],
                        password=config["password"],
                        ssl=config.get("ssl", False),
                        decode_responses=True
                    )

    def log_cost(self, amount: float, model: str):
        """Log the cost of an AI operation."""
        if not self.redis_client: return
        
        now = datetime.now()
        day_key = f"{self.namespace}:{now.strftime('%Y-%m-%d')}"
        
        # Increment daily total
        self.redis_client.incrbyfloat(day_key, amount)
        self.redis_client.expire(day_key, 86400 * 2) # Keep for 2 days

        # Log detailed entry
        log_entry = {
            "timestamp": now.isoformat(),
            "amount": amount,
            "model": model
        }
        self.redis_client.lpush(f"{self.namespace}:logs", json.dumps(log_entry))
        self.redis_client.ltrim(f"{self.namespace}:logs", 0, 1000)

    def is_within_budget(self) -> bool:
        """Check if we are still within the daily budget."""
        if not self.redis_client: return True # Fail open if Redis is down
        
        now = datetime.now()
        day_key = f"{self.namespace}:{now.strftime('%Y-%m-%d')}"
        current_spend = float(self.redis_client.get(day_key) or 0.0)
        
        return current_spend < self.daily_limit

    def get_current_spend(self) -> float:
        """Return the total spent today."""
        if not self.redis_client: return 0.0
        now = datetime.now()
        day_key = f"{self.namespace}:{now.strftime('%Y-%m-%d')}"
        return float(self.redis_client.get(day_key) or 0.0)

if __name__ == "__main__":
    bm = BudgetManager()
    print(f"Current Spend: ${bm.get_current_spend()}")
    print(f"Within Budget ($10): {bm.is_within_budget()}")
