app.py

import io, uuid, time, os
import flask
from flask_socketio import SocketIO, emit, join_room, leave_room
import eventlet
from core.mtdna_backend import *
from werkzeug.middleware.proxy_fix import ProxyFix
import mimetypes

# Many things running at once on HuggingFace and sockets so money_patch makes codes that looks synchronous run cooperatively
# == resolve incomplete responses for static files caused by socket didnt cooperate with Eventletl's loop
eventlet.monkey_patch()

print("CWD:", os.getcwd())
print("templates/ exists:", os.path.isdir("templates"), "->", os.listdir("templates") if os.path.isdir("templates") else "missing")
print("static/ exists:", os.path.isdir("static"), "->", os.listdir("static")[:10] if os.path.isdir("static") else "missing")

# accessions = []
isvip = True  # or True depending on the user

mimetypes.add_type("text/css", ".css") 

BASE_DIR = os.path.dirname(os.path.abspath(__file__))
app = flask.Flask(__name__, static_folder=os.path.join(BASE_DIR, "static"), static_url_path="/static", template_folder=os.path.join(BASE_DIR, "templates"))
app.config["DEBUG"] = True
app.config["SECRET_KEY"] = "dev-key"

# important behind HF proxy
app.wsgi_app = ProxyFix(app.wsgi_app, x_for=1, x_proto=1, x_host=1, x_prefix=1)

# # Socket.IO (logs on for debugging)
socketio = SocketIO(app, async_mode="eventlet", cors_allowed_origins="*")


# --- job registry for cancel flags ---
# Use a simple boolean flag in eventlet mode: True => cancel requested
CANCEL_FLAGS = {}  # {job_id: bool}
JOBS = {} # {job_id: {"accs": [...], "started": False}}

# (optional) no-cache while debugging
# @app.after_request
# def nocache(resp):
#     resp.headers["Cache-Control"] = "no-store"
#     return resp

print(app.url_map)

# request log to verify what the proxy is asking for
@app.before_request
def _log():
    if flask.request.path.startswith('/static/') or flask.request.path == '/favicon.ico' or flask.request.path.startswith('/socket.io'):
        return
    print("REQ:", flask.request.method, flask.request.path)

# Home
@app.route("/")
def home():
    return flask.render_template("Home.html", isvip=isvip)

# Submit route
@app.route("/submit", methods=["POST"])
def submit():
    raw_text = flask.request.form.get("raw_text", "").strip()
    file_upload = flask.request.files.get("file_upload")
    user_email = flask.request.form.get("user_email", "").strip()

    if file_upload and getattr(file_upload, "filename", ""):
        data = file_upload.read()
        buf = io.BytesIO(data); buf.name = file_upload.filename
        file_upload = buf

    accs, error = extract_accessions_from_input(file=file_upload, raw_text=raw_text)

    job_id = uuid.uuid4().hex[:8]
    CANCEL_FLAGS[job_id] = False

    # Obtain user's past usage
    user_hash = hash_user_id(user_email)
    user_usage, max_allowed = increment_usage(user_hash, 0) # get how much they have run and the maximun #queries they have
    remaining_trials = max(max_allowed - user_usage, 0) # remaining trials if everything goes well
    total_queries = max(0, min(len(accs), max_allowed - user_usage)) # limited the number of queries of users so that won't have to run all.
    
    # the list of IDs that will be run within allowance
    accs = accs[:total_queries]

    # Save var to the global environment
    JOBS[job_id] = {"accs": accs, 
                    "user": {
                        "user_hash": user_hash, 
                        "user_usage": user_usage, 
                        "max_allowed": max_allowed,
                        "total_queries": total_queries,
                        "remaining_trials": remaining_trials
                        }, 
                    "started": False} 

    return flask.redirect(flask.url_for("output", job_id=job_id), code=303)

@app.route("/output/")
def output_root():
    return flask.redirect(flask.url_for("home"))

# Output page (must accept job_id!)
@app.route("/output/<job_id>")
def output(job_id):
    started_at = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime())
    total_queries = JOBS[job_id]['user']['total_queries']
    return flask.render_template(
        "Output.html",
        job_id=job_id,
        started_at=started_at,
        total_queries=total_queries,
        isvip=isvip,
        ws_url="",  # leave empty for mock/demo mode; set later to real WS URL
    )

# Functions that communicates between web and server - run via socketio
def run_job(job_id, accessions):
    total_queries = JOBS[job_id]['user']['total_queries']
    user_hash = JOBS[job_id]['user']['user_hash'] # to update allowance in case cancelled
    room = job_id


    def send_log(msg): socketio.emit("log", {"msg": msg}, room=room)
    def send_row(row): socketio.emit("row", row, room=room)

    
    try:
        socketio.emit("status", {"state": "started", "total": len(accessions)}, room=room)
        start_time = time.perf_counter()
        send_log(f"Job {job_id} started. {total_queries} accession(s).")

        outs = []
        for i, acc in enumerate(accessions, 1):
            if CANCEL_FLAGS.get(job_id):
                send_log("Cancellation requested. Stopping…")
                socketio.emit("status", {"state": "cancelled"}, room=room)
                increment_usage(user_hash, i - 1)
                return

            t0 = time.perf_counter()
            out = summarize_results(acc)  # may be dict / [] / None
            dt = time.perf_counter() - t0

            # ---- normalise 'out' to a dict we can emit safely ----
            if not out:
                out = {}
            elif isinstance(out, list):
                # If a list slipped through, try to coerce sensibly
                if out and isinstance(out[0], dict):
                    out = out[0]
                elif out and isinstance(out[0], list):
                    # very defensive: list-of-lists -> map by expected order
                    keys = [
                        "Sample ID", "Predicted Country", "Country Explanation",
                        "Predicted Sample Type", "Sample Type Explanation",
                        "Sources", "Time cost"
                    ]
                    row0 = out[0]
                    out = {k: (row0[idx] if idx < len(row0) else "") for idx, k in enumerate(keys)}
                else:
                    out = {}
            elif not isinstance(out, dict):
                out = {}

            # ---- map backend keys (Title Case) to frontend keys (snake_case) ----
            sample_id                 = out.get("Sample ID") or str(acc)
            predicted_country         = out.get("Predicted Country", "unknown")
            country_explanation       = out.get("Country Explanation", "unknown")
            predicted_sample_type     = out.get("Predicted Sample Type", "unknown")
            sample_type_explanation   = out.get("Sample Type Explanation", "unknown")
            sources                   = out.get("Sources", "No Links")
            time_cost                 = out.get("Time cost") or f"{dt:.2f}s"

            send = {
                "idx": i,
                "sample_id": sample_id,
                "predicted_country": predicted_country,
                "country_explanation": country_explanation,
                "predicted_sample_type": predicted_sample_type,
                "sample_type_explanation": sample_type_explanation,
                "sources": sources,
                "time_cost": time_cost,
            }

            socketio.emit("row", send, room=room)
            socketio.sleep(0)  # <- correct spelling; yield so the emit flushes
            send_log(f"Processed {acc} in {dt:.2f}s")

        total_dt = time.perf_counter() - start_time
        
        # Update user allowance
        increment_usage(user_hash, total_queries)
        # Calculate remaining_trials to display for user
        remaining_trials = JOBS[job_id]['user']['remaining_trials']

        socketio.emit("status", {"state": "finished", "elapsed": f"{total_dt:.2f}s"}, room=room)
        send_log(f"Job finished successfully. Number of trials left is")
    except Exception as e:
        send_log(f"ERROR: {e}")
        socketio.emit("status", {"state": "error", "message": str(e)}, room=room)
    finally:
        CANCEL_FLAGS.pop(job_id, None)
        JOBS.pop(job_id, None)  # <— tidy queued job

# ---- Socket.IO events ----
@socketio.on("connect")
def on_connect():
    emit("connected", {"ok": True})

@socketio.on("join")
def on_join(data):
    job_id = data.get("job_id")
    if job_id:
        join_room(job_id)
        emit("joined", {"room": job_id})

        # Start the job once the client is in the room
        job = JOBS.get(job_id)
        if job and not job["started"]:
            job["started"] = True
            total = len(job["accs"])
            # Send an initial queued/total status so the UI can set progress denominator
            socketio.emit("status", {"state": "queued", "total": total}, room=job_id)
            socketio.start_background_task(run_job, job_id, job["accs"])

@socketio.on("leave")
def on_leave(data):
    job_id = data.get("job_id")
    if job_id:
        leave_room(job_id)

@socketio.on("cancel")
def on_cancel(data):
    job_id = data.get("job_id")
    if job_id in CANCEL_FLAGS:
        CANCEL_FLAGS[job_id] = True  # flip the flag
        emit("status", {"state": "cancelling"}, room=job_id)

@app.route("/about")
def about():
    return flask.render_template("About.html", isvip=isvip)

@app.route("/pricing")
def pricing():
    return flask.render_template("Pricing.html", isvip=isvip)

@app.route("/contact")
def contact():
    return flask.render_template("Contact.html", isvip=isvip)

if __name__ == "__main__":
    port = int(os.environ.get("PORT", 7860))  # HF Spaces injects PORT
    socketio.run(app, host="0.0.0.0", port=port)