import warnings

warnings.filterwarnings("ignore", category=DeprecationWarning)
warnings.filterwarnings("ignore", category=FutureWarning)
from flask import (
    Flask,
    render_template,
    request,
    send_from_directory,
    make_response,
    redirect,
    url_for,
)
import json
import pandas as pd
import shutil
from bson.objectid import ObjectId
import zipfile
from flask_mongoengine import MongoEngine
from datetime import datetime, timedelta, timezone
# from SiteContext import SiteContext
import os
import glob
from rq import Queue
from env_info import environment_var, UNZIP_FOLDER, STATIC_FOLDER
from site_context_worker import conn
from site_context_compute import (
    unzip_folder,
    insert_data_mongodb, 
    run_single_molecules, 
    get_shape_feature_result,
    fileHandler,
    run_demo
)


q = Queue(connection=conn, default_timeout=300000)

app = Flask(__name__)

app.environment = environment_var

if app.environment == "DEV":
    app.config["MONGODB_SETTINGS"] = {"host": "mongodb://localhost:27017/SiteContext"}
    app_route_front = "/"
    print("In dev")
else:
    app.config['MONGODB_SETTINGS'] = {'host':"mongodb://"}
    app_route_front = "/SiteContext"
    print("In prod")

# upload and download folder location
OUTPUT_FOLDER = os.path.join(app.root_path, "output")
UPLOAD_FOLDER = os.path.join(app.root_path, "run_uploads")
app.config["UPLOAD_FOLDER"] = UPLOAD_FOLDER
app.config["SECRET_KEY"] = "arnold"

db = MongoEngine(app)


class Project(db.Document):

    
    meta = {
        "collection": "SiteContext_projects",
        "indexes": [
            # Each doc expires at its own `expires_at` timestamp
            {"fields": ["expires_at"], "expireAfterSeconds": 0}
        ],
    }
    file = db.StringField()
    job_id = db.StringField()
    job_type = db.StringField()
    job_status = db.StringField()
    feature_type = db.StringField() # string of list
    set1_molecules_name = db.StringField()  
    set2_molecules_name = db.StringField()
    expires_at = db.DateTimeField(required=True)


@app.route("/")
@app.route("/index")
@app.route(app_route_front)
@app.route(app_route_front +"/")
def hello_world():
    return render_template("index.html")

@app.route(app_route_front + "/run-demo-SiteContext")
def run_demo_Sitecontext():
    return render_template("run_SiteContext_demo.html")


@app.route(app_route_front + "/run-demo-SiteContext", methods=["POST"])
def run_demo_SiteContext_withPost():
    print("IN demo siteContext")    
    feature_option = request.form.get('feature_option', 'shape')
    pre_selected_pair = request.form.get('preselect_option')
    molecule1 = request.form.get('molecule1')
    molecule2 = request.form.get('molecule2')

    if feature_option == 'shape' or feature_option == '':
        feature_type_list = ['shape']
    elif feature_option == 'shape+feature':
        feature_type_list = ['shape+feature']
    elif feature_option == 'both':
        feature_type_list = ['shape', 'shape+feature']
    else:
        return json.dumps({"error": 'feature input is not right.'})

    if pre_selected_pair != '':
        pdb_set1_single, pdb_set_2_single = pre_selected_pair.split("_")
        pdb_set1_names = [pdb_set1_single]
        pdb_set2_names = [pdb_set_2_single]
    elif molecule1 != '' and molecule2 != '':
        pdb_set1_names = [molecule1]
        pdb_set2_names = [molecule2]
    else:
        return json.dumps({"error": 'Option need to re-select.'})

    project = Project(
        file="",
        job_id="-1",
        job_status="-1",
        job_type="DEMO",
        feature_type = json.dumps(feature_type_list), 
        set1_molecules_name=json.dumps(pdb_set1_names),
        set2_molecules_name=json.dumps(pdb_set2_names),
        expires_at=datetime.now(timezone.utc) + timedelta(hours=1),
    ).save()
    print("Project id: " + str(project.id))
    # ---------------------------------------
    output_folder = os.path.join(OUTPUT_FOLDER, str(project["id"]))
    if not os.path.exists(output_folder):
        os.makedirs(output_folder)

    # This part need to be changed
    # save data into mongodb
    
    run_demo(output_folder, project)
    
    return json.dumps({'project_id': str(project.id)})


@app.route(app_route_front + "/run-single-SiteContext")
def run_single_Sitecontext():
    return render_template("run_SiteContext_single.html")


@app.route(app_route_front + "/run-single-SiteContext", methods=["POST"])
def run_single_SiteContext_withPost():
    print("IN single siteContext")
    file_name = request.form.get("filename")
    feature_option = request.form.get('feature_option', 'shape')

    
    if feature_option == 'shape' or feature_option == '':
        feature_type_list = ['shape']
    elif feature_option == 'shape+feature':
        feature_type_list = ['shape+feature']
    elif feature_option == 'both':
        feature_type_list = ['shape', 'shape+feature']
    else:
        return json.dumps({"error": 'feature input is not right.'})
    
    file_path = os.path.join(UPLOAD_FOLDER, file_name)
    unzipped_folder_name = unzip_folder(file_path)
    input_folder = os.path.join(UNZIP_FOLDER, file_name).replace(".zip", "")
    os.rename(os.path.join(UNZIP_FOLDER, unzipped_folder_name), input_folder)
    # check pdb files in side uncompared folder
    sub_folders = os.listdir(input_folder)
    print(f'sub_folders: {sub_folders}')
    
    if "set1" not in sub_folders or "set2" not in sub_folders:
        return json.dumps({"error": "The 'set1' or 'set2' folder is missing."})
    for sub_folder in os.listdir(input_folder):
        if sub_folder[0] == ".":
            continue
        if sub_folder == "set1":
            pdb_files = [f for f in glob.glob(input_folder + "/" + sub_folder + "/*.pdb")]
            if pdb_files == []:
                return json.dumps({"error": "No files found in the 'set1' folder."})
            pdb_set1_names = [os.path.splitext(os.path.basename(pdb_files[0]))[0]]
        elif sub_folder == "set2":
            pdb_files = [f for f in glob.glob(input_folder + "/" + sub_folder + "/*.pdb")]
            if pdb_files == []:
                return json.dumps({"error": "No files found in the 'set2' folder."})
            pdb_set2_names = [os.path.splitext(os.path.basename(pdb_files[0]))[0]]
        else:
            return json.dumps({"error": "The folder structure does not comply with the required format."})
    project = Project(
        file=file_name,
        job_id="-1",
        job_status="-1",
        job_type="SINGLE",
        feature_type = json.dumps(feature_type_list), 
        set1_molecules_name=json.dumps(pdb_set1_names),
        set2_molecules_name=json.dumps(pdb_set2_names),
        expires_at=datetime.now(timezone.utc) + timedelta(hours=1)
    ).save()
    # ---------------------------------------
    output_folder = os.path.join(OUTPUT_FOLDER, str(project["id"]))

    if not os.path.exists(output_folder):
        os.makedirs(output_folder)

    # save data into mongodb
    response = insert_data_mongodb(input_folder, project)
    if response != []:
        return {'error': json.dumps(response)}
    run_single_molecules(output_folder, project)

    return json.dumps({'project_id': str(project.id)})


@app.route(app_route_front + "/run-batch-SiteContext")
def run_batch_SiteContext():
    return render_template(
        "run_SiteContext_batch.html",
        info="Upload Reddit posts of multiple users to investigate their opioid addiction and recovery journey.",
    )


@app.route(app_route_front + "/run-batch-SiteContext", methods=["POST"])
def run_batch_SiteContext_withPost():
    file_name = request.form.get("filename")
    feature_option = request.form.get('feature_option', 'shape')
    
    if feature_option == 'shape' or feature_option == '':
        feature_type_list = ['shape']
    elif feature_option == 'shape+feature':
        feature_type_list = ['shape+feature']
    elif feature_option == 'both':
        feature_type_list = ['shape', 'shape+feature']
    else:
        return json.dumps({"error": 'feature input is not right.'})
    
    file_path = os.path.join(UPLOAD_FOLDER, file_name)
    unzipped_folder_name = unzip_folder(file_path)
    input_folder = os.path.join(UNZIP_FOLDER, file_name).replace(".zip", "")
    os.rename(os.path.join(UNZIP_FOLDER, unzipped_folder_name), input_folder)
    # check pdb files in side uncompared folder
    sub_folders = os.listdir(input_folder)
    if "set1" not in sub_folders or "set2" not in sub_folders:
        return json.dumps({"error": "The 'set1' or 'set2' folder is missing."})
    for sub_folder in os.listdir(input_folder):
        if sub_folder[0] == ".":
            continue
        if sub_folder == "set1":
            pdb_files = [f for f in glob.glob(input_folder + "/" + sub_folder + "/*.pdb")]
            if pdb_files == []:
                return json.dumps({"error": "No files found in the 'set1' folder."})
            pdb_set1_names = [os.path.splitext(os.path.basename(pdb_file))[0] for pdb_file in pdb_files[:10]]
        elif sub_folder == "set2":
            pdb_files = [f for f in glob.glob(input_folder + "/" + sub_folder + "/*.pdb")]
            if pdb_files == []:
                return json.dumps({"error": "No files found in the 'set2' folder."})
            pdb_set2_names = [os.path.splitext(os.path.basename(pdb_file))[0] for pdb_file in pdb_files[:10]]
        else:
            return json.dumps({"error": "The folder structure does not comply with the required format."})

    project = Project(
        file=file_name,
        job_id="-1",
        job_status="-1",
        job_type="BATCH",
        feature_type = json.dumps(feature_type_list), 
        set1_molecules_name=json.dumps(pdb_set1_names),
        set2_molecules_name=json.dumps(pdb_set2_names),
        expires_at=datetime.now(timezone.utc) + timedelta(hours=48)
    ).save()


    job = q.enqueue_call(
        func=fileHandler, args=(project.id,), result_ttl=5000, timeout=300000
    )
    job_id = str(job.get_id())

    Project.objects(id=project.id).update_one(set__job_id=job_id)

    return render_template(
        "run_SiteContext_batch.html",
        submission_successful=True,
        project_id=project.id,
        job_id=job_id,
        info="Upload multiple molecules to learn their bidding site information.",
    )

@app.route(app_route_front + '/upload-data', methods=['POST'])
def upload():
    file = request.files['file']
    save_path = os.path.join(app.config['UPLOAD_FOLDER'], file.filename)
    current_chunk = int(request.form['dzchunkindex'])
    # If the file already exists it's ok if we are appending to it,
    # but not if it's new file that would overwrite the existing one
    if os.path.exists(save_path) and current_chunk == 0:
        # 400 and 500s will tell dropzone that an error occurred and show an error
        return make_response(('File already exists', 400))
    try:
        with open(save_path, 'ab') as f:
            f.seek(int(request.form['dzchunkbyteoffset']))
            f.write(file.stream.read())
    except OSError:
        # log.exception will include the traceback so we can see what's wrong
        print ('Could not write to file')
        return make_response(("Not sure why,"
                              " but we couldn't write the file to disk", 500))

    total_chunks = int(request.form['dztotalchunkcount'])

    if current_chunk + 1 == total_chunks:
        # This was the last chunk, the file should be complete and the size we expect
        if os.path.getsize(save_path) != int(request.form['dztotalfilesize']):
            print ("File {file.filename} was completed, " +
                      "but has a size mismatch." +
                      "Was {os.path.getsize(save_path)} but we" +
                      " expected {request.form['dztotalfilesize']} ")
            return make_response(('Size mismatch', 500))

    else:
        print('Chunk {current_chunk + 1} of {total_chunks} ' + 'for file {file.filename} complete')

    extension = ".zip" if file.filename.endswith('.zip') else '.csv'

    new_name = str(datetime.now(timezone.utc).time()).replace("/", "-") + extension
    renamed_path = os.path.join(app.config['UPLOAD_FOLDER'], new_name)
    os.rename(save_path, renamed_path)
    return make_response((new_name, 200))


@app.route(app_route_front + "/results")
def get_results():
    return render_template("download.html")

@app.route(app_route_front + '/results/<project_id>')
def get_visualization(project_id):
    print("getting visualzination")
    root_output_folder = os.path.join(OUTPUT_FOLDER, project_id)
    shape, shape_radius, sinkhorn_score_shape, final_RMSD_shape, shape_feature, shape_feature_radius, sinkhorn_score_shape_feature, final_RMSD_shape_feature = get_shape_feature_result(root_output_folder)
    return render_template("visualization_SiteContext.html", project_id=project_id,
                            shape=shape, shape_radius=shape_radius, sinkhorn_score_shape=sinkhorn_score_shape, final_RMSD_shape=final_RMSD_shape, 
                            shape_feature=shape_feature, shape_feature_radius=shape_feature_radius, sinkhorn_score_shape_feature=sinkhorn_score_shape_feature, final_RMSD_shape_feature=final_RMSD_shape_feature)

@app.route(app_route_front + "/results", methods=["POST"])
def post_results():
    print("-------In post_results function.-------")
    project_id = request.form.get("project_id")
    if not project_id:
        return render_template("download.html")
    try:
        project = Project.objects.get(id=project_id)
        print(project)
        if not project:
            return render_template("download.html", project_id=project_id, error=True)

        if project.job_status != "1" or not os.path.exists(
            os.path.join(OUTPUT_FOLDER, project_id)
        ):
            return render_template("download.html", project_id=project_id)

        if not os.path.exists(os.path.join(OUTPUT_FOLDER, project_id) + ".zip"):
            shutil.make_archive(
                os.path.join(OUTPUT_FOLDER, project_id),
                "zip",
                os.path.join(OUTPUT_FOLDER, project_id),
            )

        return send_from_directory(OUTPUT_FOLDER, project_id + ".zip")

    except:
        return render_template("download.html", project_id=project_id, error=True)


@app.route(app_route_front + '/output/<project_id>/<filename>')
def visualization_file(project_id, filename):
    print("in visualization_file")
    return send_from_directory(os.path.join(OUTPUT_FOLDER, project_id), 
                               filename, max_age=0) # cache_timeout is replaced by max_age


@app.route(app_route_front + "/visual-sample-binding")
def run_sample_binding():
    return render_template("visualization_Sample_Binding.html")


@app.route(app_route_front + "/tutorial")
def get_tutorial_quick():
    return render_template("tutorial.html")


@app.route(app_route_front + "/about")
def get_citation():
    return render_template("about.html")


@app.route(app_route_front + "/funding")
def get_funding():
    return render_template("funding.html")


@app.route(app_route_front + "/contact")
def get_contact():
    return render_template("contact.html")


if __name__ == "__main__":
    if app.environment == "DEV":
        app.run(debug=True, threaded=True, host="localhost", port=4500)
    else:
        app.run(debug=False, threaded=True, host="localhost", port=8000)
