import json
import uuid
import os


class TokenUtilizationSystem:
    def __init__(self, db_path="user_database.json"):
        self.db_path = db_path
        self.default_limit = 1_000_000
        self.users = self._load_db()

    def _load_db(self):
        if os.path.exists(self.db_path):
            with open(self.db_path, 'r', encoding='utf-8') as f:
                return json.load(f)
        return {}

    def _save_db(self):
        safe = {}
        for uid, data in self.users.items():
            safe[uid] = {
                "name": data["name"],
                "input_tokens_used": data["input_tokens_used"],
                "output_tokens_used": data["output_tokens_used"],
                "token_balance": data["token_balance"]
            }
        with open(self.db_path, 'w', encoding='utf-8') as f:
            json.dump(safe, f, indent=4)

    def register(self, name):
        """
        Register a new user or return existing user's ID.

        Args:
            name (str): User's name
        Returns:
            user_id (str)
        """
        for uid, data in self.users.items():
            if data['name'].lower() == name.lower():
                return uid

        user_id = f"USR-{uuid.uuid4().hex[:6].upper()}"
        self.users[user_id] = {
            "name": name,
            "input_tokens_used": 0,
            "output_tokens_used": 0,
            "token_balance": self.default_limit
        }
        self._save_db()
        return user_id

    def track(self, user_id, input_tokens, output_tokens):
        """
        Track token usage after a chat response.

        Args:
            user_id (str): ID returned from register()
            input_tokens (int): prompt_tokens from your LLM response
            output_tokens (int): completion_tokens from your LLM response
        Returns:
            dict with updated usage, or error string if limit exceeded
        """
        user = self.users.get(user_id)
        if not user:
            return "Error: User not found."

        if user['token_balance'] <= 0:
            return "Token limit of 1,000,000 reached. Please upgrade to Premium."

        user['input_tokens_used'] += input_tokens
        user['output_tokens_used'] += output_tokens
        user['token_balance'] = max(0, user['token_balance'] - (input_tokens + output_tokens))
        self._save_db()

        return {
            "input_tokens_used": user['input_tokens_used'],
            "output_tokens_used": user['output_tokens_used'],
            "remaining_balance": user['token_balance']
        }

    def get_usage(self, user_id):
        """
        Get token usage stats for a user.

        Args:
            user_id (str): ID returned from register()
        """
        user = self.users.get(user_id)
        if not user:
            return "Error: User not found."
        return {
            "name": user["name"],
            "input_tokens_used": user["input_tokens_used"],
            "output_tokens_used": user["output_tokens_used"],
            "token_balance": user["token_balance"]
        }

    def has_balance(self, user_id):
        """
        Check if user still has token balance remaining.

        Args:
            user_id (str): ID returned from register()
        Returns:
            bool
        """
        user = self.users.get(user_id)
        return bool(user and user['token_balance'] > 0)
