from flask import Flask, render_template, jsonify, request
import json
import random
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch
import os

# Set the TRANSFORMERS_CACHE environment variable
os.environ['TRANSFORMERS_CACHE'] = '/app/cache'

app = Flask(__name__)

# Load initial coins
with open('coins.json', 'r') as f:
    coins = json.load(f)

# Initialize game state
game_state = {
    'balance': 0,
    'flips_left': 1000,
    'current_coin': 0
}

# Load leaderboard
try:
    with open('leaderboard.json', 'r') as f:
        leaderboard = json.load(f)
except FileNotFoundError:
    leaderboard = []

# Load AI model
tokenizer = AutoTokenizer.from_pretrained("TinyLlama/TinyLlama-1.1B-Chat-v1.0")
model = AutoModelForCausalLM.from_pretrained("TinyLlama/TinyLlama-1.1B-Chat-v1.0")

@app.route('/')
def index():
    return render_template('index.html', coins=coins, game_state=game_state, leaderboard=leaderboard)

@app.route('/flip', methods=['POST'])
def flip_coin():
    if game_state['flips_left'] > 0:
        coin = coins[game_state['current_coin']]
        is_heads = random.random() < coin['winrate']
        if is_heads:
            game_state['balance'] += coin['value']
            # Apply coin bonus
            game_state['balance'] += coin['bonus']
        game_state['flips_left'] -= 1
        
        return jsonify({
            'result': 'H' if is_heads else 'T',
            'balance': game_state['balance'],
            'flips_left': game_state['flips_left']
        })
    return jsonify({'error': 'No flips left'})

@app.route('/buy', methods=['POST'])
def buy_coin():
    index = int(request.json['index'])
    if index < len(coins) and game_state['balance'] >= coins[index]['price'] and index != game_state['current_coin']:
        game_state['balance'] -= coins[index]['price']
        game_state['current_coin'] = index
        return jsonify({'success': True, 'balance': game_state['balance'], 'current_coin': index})
    return jsonify({'success': False})

@app.route('/generate_coin', methods=['POST'])
def generate_coin():
    if game_state['balance'] < 4:
        return jsonify({'success': False, 'error': 'Not enough balance to generate a coin'})

    prompt = """Generate a new coin for a game in JSON format with the following properties:
    name (a creative name for the coin),
    color (as a hex code),
    price (between 10 and 100),
    value (between 0.1 and 1.0),
    winrate (between 0.5 and 0.8),
    bonus (a flat $ value between 0.01 and 0.1)
    
    Example:
    {
        "name": "Lucky Clover",
        "color": "#2ecc71",
        "price": 15,
        "value": 0.2,
        "winrate": 0.65,
        "bonus": 0.05
    }
    """

    try:
        inputs = tokenizer(prompt, return_tensors="pt")
        with torch.no_grad():
            outputs = model.generate(**inputs, max_length=300)
        response = tokenizer.decode(outputs[0], skip_special_tokens=True)

        new_coin = json.loads(response)
        if validate_coin(new_coin):
            coins.append(new_coin)
            with open('coins.json', 'w') as f:
                json.dump(coins, f)
            game_state['balance'] -= 4
            return jsonify({'success': True, 'coin': new_coin, 'balance': game_state['balance']})
        else:
            return jsonify({'success': False, 'error': 'Invalid coin format'})
    except Exception as e:
        # Refund the player if coin generation fails
        return jsonify({'success': False, 'error': str(e), 'balance': game_state['balance']})

@app.route('/game_over', methods=['POST'])
def game_over():
    initials = request.json['initials']
    score = game_state['balance']
    leaderboard.append({'initials': initials, 'score': score})
    leaderboard.sort(key=lambda x: x['score'], reverse=True)
    leaderboard = leaderboard[:10]  # Keep only top 10
    with open('leaderboard.json', 'w') as f:
        json.dump(leaderboard, f)
    return jsonify({'success': True})

@app.route('/reset_game', methods=['POST'])
def reset_game():
    game_state['balance'] = 0
    game_state['flips_left'] = 1000
    game_state['current_coin'] = 0
    return jsonify({'success': True})

def validate_coin(coin):
    required_keys = ['name', 'color', 'price', 'value', 'winrate', 'bonus']
    return all(key in coin for key in required_keys) and \
           isinstance(coin['name'], str) and \
           isinstance(coin['color'], str) and coin['color'].startswith('#') and \
           isinstance(coin['price'], (int, float)) and 10 <= coin['price'] <= 100 and \
           isinstance(coin['value'], (int, float)) and 0.1 <= coin['value'] <= 1.0 and \
           isinstance(coin['winrate'], (int, float)) and 0.5 <= coin['winrate'] <= 0.8 and \
           isinstance(coin['bonus'], (int, float)) and 0.01 <= coin['bonus'] <= 0.1

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=7860)