import utils
import time
import datetime

import pandas as pd
import bittensor as bt
from typing import Dict, List, Any, Optional, Tuple
from flask import Flask, request, jsonify


app = Flask(__name__)

# Global variables (saves time on loading data)

state_vars = utils.test_load_state_vars()
metagraph = state_vars["metagraph"]
model_data = state_vars["model_data"]
vali_runs = state_vars["vali_runs"]
scores = state_vars["scores"]
validator_df = state_vars["validator_df"]
benchmarks = state_vars.get("benchmarks", None)
benchmark_timestamp = state_vars.get("benchmark_timestamp", None)

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

@app.route('/reload', methods=['GET'])
def reload():
    """
    Reload the state variables
    """
    global metagraph, model_data, vali_runs, scores, validator_df, benchmarks, benchmark_timestamp
    state_vars = utils.load_state_vars()
    metagraph = state_vars["metagraph"]
    model_data = state_vars["model_data"]
    vali_runs = state_vars["vali_runs"]
    scores = state_vars["scores"]
    validator_df = state_vars["validator_df"]
    benchmarks = state_vars.get("benchmarks", None)
    benchmark_timestamp = state_vars.get("benchmark_timestamp", None)

    return jsonify({"message": "State variables reloaded"})

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

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

    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)
    """
    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)
    """
    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)
    """
    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)
    """
    return jsonify(
        utils.make_validator_dataframe(validator_df, model_data).to_dict(orient='records')
        )


if __name__ == '__main__':

    app.run(port=5000, debug=True)