from app import db
from app.models import Annual_Data, Glacier, User
from flask_sqlalchemy import BaseQuery
from flask import flash, redirect, url_for


def query_annual_data(form) -> BaseQuery:
    annual_data = db.session.query(Annual_Data).filter_by(year=form.year.data)
    if form.name.data:
        glacier_name = (
            db.session.query(Annual_Data)
            .join(Glacier, Glacier.id == Annual_Data.id)
            .filter_by(name=form.name.data)
            .group_by(Annual_Data.year)
        )
        if glacier_name.first() is None:
            flash("Sorry, no results found")
            return redirect(url_for("table_selection"))
        return glacier_name
    return annual_data


def query_user(form) -> BaseQuery:
    user = User.query.filter_by(username=form.username.data).first()
    return user


def query_plot_data(form) -> BaseQuery:
    query = (
        db.session.query(Annual_Data)
        .join(Glacier, Glacier.id == Annual_Data.id)
        .filter_by(name=form.name.data)
        .group_by(Annual_Data.year)
    )
    if query.first() is None:
        flash("Sorry, no results found")
        return redirect(url_for("plot_selection"))
    return query