import pathlib

from gevent.pywsgi import WSGIServer
from tabulate import tabulate

from flask import Flask, request, render_template
from wtforms import Form, StringField, SubmitField
from wtforms.validators import DataRequired

from backend import AIBackend

app = Flask("xkcd_retriever", template_folder=pathlib.Path(__file__).parent / "templates")

print("Initializing backend")
backend = AIBackend()

# backend.warmup() # Only needed when there is no data in the pgvector database
backend._ready = True
print("AI backend initialized ...")

class BasicForm(Form):
    ids = StringField("ID",validators=[DataRequired()])
    submit = SubmitField("Submit")

@app.route("/",methods =['POST','GET'])
def main():
    form = BasicForm()
    return render_template("index.html", form = form)

@app.route("/search",methods =['POST','GET'])
def results():
    topic = request.form.get('ids')

    result = backend.query(f"Give an example for {topic}")

    headers = ["ID", "Title", "Link"]
    print(tabulate(backend.format_result(result), headers=headers, tablefmt="github"))

    res = backend.format_for_api(result)
    return render_template("results.html", results=res, topic=topic)

@app.route("/find")
def query_backend():
    topic = request.args.get('topic')
    result = backend.query(f"Give an example for {topic}")

    headers = ["ID", "Title", "Link"]
    print(tabulate(backend.format_result(result), headers=headers, tablefmt="github"))

    return backend.format_for_api(result)

if __name__ == "__main__":
    http_server = WSGIServer(("0.0.0.0", 8000), app)
    http_server.serve_forever()