import atexit
import datetime

from flask import Flask, request, jsonify
from apscheduler.schedulers.background import BackgroundScheduler

import utils

app = Flask(__name__)

# Global variables (saves time on loading data)
state_vars = None
reload_timestamp = datetime.datetime.now().strftime('%D %T')


def load_data(test=False):
    """
    Reload the state variables
    """
    global state_vars, reload_timestamp
    if test:
        state_vars = utils.test_load_state_vars()    
    else:  
        state_vars = utils.load_state_vars()
        
    reload_timestamp = datetime.datetime.now().strftime('%D %T')

    print(f'Reloaded data at {reload_timestamp}')


def start_scheduler():
    scheduler = BackgroundScheduler()
    scheduler.add_job(func=load_data, trigger="interval", seconds=60*30)
    scheduler.start()

    # Shut down the scheduler when exiting the app
    atexit.register(lambda: scheduler.shutdown())


@app.route('/', methods=['GET'])
def home():
    return "Welcome to the Bittensor Pretraining Leaderboard API!"


@app.route('/updated', methods=['GET'])
def updated():
    return reload_timestamp


@app.route('/benchmark', methods=['GET'])
def benchmark():
    """
    Get the benchmarks and the timestamp

    Returns:
    - benchmarks: List of dicts (from pandas DataFrame)
    - benchmark_timestamp: String
    """

    benchmarks = state_vars.get("benchmarks", None)
    benchmark_timestamp = state_vars.get("benchmark_timestamp", None)

    return jsonify(
        {
            "benchmarks": benchmarks.to_dict(orient='records'),
            "benchmark_timestamp": benchmark_timestamp.strftime('%Y-%m-%d %H:%M:%S')
        }
    )

@app.route('/metagraph', methods=['GET'])
def metagraph():
    """
    Get the metagraph data
    Returns:
    - metagraph_data: List of dicts (from pandas DataFrame)
    """
    
    metagraph = state_vars["metagraph"]
     
    return jsonify(
        utils.make_metagraph_dataframe(metagraph).to_dict(orient='records')
    )

@app.route('/leaderboard', methods=['GET'])
def leaderboard():
    """
    Get the leaderboard data
    Returns:
    - leaderboard_data: List of dicts (from pandas DataFrame)
    """

    model_data = state_vars["model_data"]
    scores = state_vars["scores"]
    show_stale = request.args.get('show_stale')
    return jsonify(
        utils.leaderboard_data(model_data, scores, show_stale=show_stale)
        )


@app.route('/loss', methods=['GET'])
def loss():
    """
    Get the losses over time
    Returns:
    - losses_over_time: List of dicts (from pandas DataFrame)
    """
    vali_runs = state_vars["vali_runs"]

    return jsonify(
        utils.get_losses_over_time(vali_runs).to_dict(orient='records')
        )


@app.route('/validator', methods=['GET'])
def validator():
    """
    Get the validator data
    Returns:
    - validator_data: List of dicts (from pandas DataFrame)
    """
    model_data = state_vars["model_data"]
    validator_df = state_vars["validator_df"]       

    return jsonify(
        utils.make_validator_dataframe(validator_df, model_data).to_dict(orient='records')
        )


if __name__ == '__main__':
    
    load_data()
    start_scheduler()

    app.run(host='0.0.0.0', port=5000, debug=True)