import glob
from io import BytesIO
from quart import Blueprint, request, send_file
from quart.utils import run_sync
from matplotlib import pyplot as plt

from CONSTANT import DATA_BASEPATH
from modules.tidi.plot import TidiPlotv2
from modules.tidi.staged.plot import tidi_render


tidi_module = Blueprint("Tidi", __name__)


@tidi_module.route('/metadata')
def get_all_years():
    res = glob.glob(f"{DATA_BASEPATH.tidi}/**/**.txt", recursive=True)
    # search for the folder name that is year

    return {
        "path": res
    }


@tidi_module.route('/render/wave')
async def render_wave():
    mode = request.args.get('mode')
    year = request.args.get('year')
    k = request.args.get('k')
    T = request.args.get('T')
    height = request.args.get('height')
    lat_range = request.args.get('lat_range')
    # like `0 ~ 5`
    year = int(year)
    k = int(k)
    T = int(T)
    height = float(height)
    lat_range = tuple(map(int, lat_range.split('~')))

    await run_sync(tidi_render)(mode, year, k, height, lat_range, T)
    buffer = BytesIO()
    plt.savefig(buffer, format="png")
    buffer.seek(0)

    return await send_file(buffer, mimetype="image/png")


@tidi_module.route('/render/month_stats_v1')
async def render_stats_v1():
    year = request.args.get('year')
    year = int(year)

    plotter = await run_sync(TidiPlotv2)(year)
    await run_sync(plotter.plot_v1)()
    buffer = BytesIO()
    plt.savefig(buffer, format="png")
    buffer.seek(0)
    return await send_file(buffer, mimetype="image/png")


@tidi_module.route('/render/month_stats_v2')
async def render_stats_v2():
    year = request.args.get('year')
    year = int(year)

    plotter = await run_sync(TidiPlotv2)(year)
    await run_sync(plotter.plot_month)()
    buffer = BytesIO()
    plt.savefig(buffer, format="png")
    buffer.seek(0)
    return await send_file(buffer, mimetype="image/png")