from flask import Flask, render_template, request, jsonify
from chatbot.orchestrator import run_pipeline
from datetime import datetime, timezone
import uuid
import traceback

app = Flask(__name__)
user_sessions = {}

#BETA_EXPIRATION = datetime(2026, 12, 31, 20, 59, tzinfo=timezone.utc)

#@app.before_request
#def check_beta_expiration():
    #if datetime.now(timezone.utc) > BETA_EXPIRATION:
        #return render_template("expired.html")

@app.route("/")
def home():
    return render_template("chat.html")

@app.route("/message", methods=["POST"])
def message():
    data = request.json
    query = data.get("message", "").strip()
    user_id = data.get("user_id", str(uuid.uuid4()))

    if user_id not in user_sessions:
        user_sessions[user_id] = {
            "targets": [],
            "selected_target": None,
            "mode_selected": False
        }

    session = user_sessions[user_id]

    # Step 1: Target selection by number
    if query.isdigit() and session["targets"]:
        index = int(query) - 1
        if 0 <= index < len(session["targets"]):
            selected = session["targets"][index]
            session["selected_target"] = selected["id"]
            return jsonify({
                "response": f"You selected ***{selected['name']}*** from ***{selected.get('organism', 'unknown organism')}***. Would you prefer a **strict** or **lenient** prediction?",
                "options": ["Strict", "Lenient"]
            })
        else:
            return jsonify({"response": "Invalid number. Please choose a valid index from the table."})

    # Step 2: Handle ChEMBL ID or SMILES or keyword
    response = run_pipeline(query, user_id)

    # If multiple targets found, save them in session
    if response.get("targets"):
        session["targets"] = response["targets"]

    # If a single target is selected, remember it
    if response.get("target"):
        session["selected_target"] = response["target"]["id"]

    return jsonify(response)

@app.route("/apply_filter", methods=["POST"])
def apply_filter():
    data = request.json
    user_id = data.get("user_id", str(uuid.uuid4()))
    mode = data.get("choice", "").lower()

    session = user_sessions.get(user_id, {})
    target_id = session.get("selected_target")

    if not target_id or not target_id.startswith("CHEMBL"):
        return jsonify({"response": "No target selected. Please start by providing a target name or ID."})

    strict = mode == "strict"
    lipinski_threshold = 4 if strict else 3
    exploratory = not strict

    try:
        response = run_pipeline(
            query_or_target_id=target_id,
            user_id=user_id,
            apply_lipinski=True,
            lipinski_threshold=lipinski_threshold,
            apply_admet=True,
            exploratory=exploratory
        )
        return jsonify(response)
    except Exception as e:
        print("❌ Error during prediction:\n")
        traceback.print_exc()
        return jsonify({
            "response": "An error occurred while making predictions. Please try again or contact support."
        })
        
application = app

if __name__ == "__main__":
    app.run(debug=True)