# app.py — Hugging Face Space entrypoint
# Simple name gate (First/Last only) → Chat + thumbs feedback + optional Admin CSV export.

import os
os.environ.setdefault("OMP_NUM_THREADS", "1")
os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")

import sqlite3, csv, time
from datetime import datetime
from typing import Optional, Tuple

import gradio as gr
from First_Pass import ask  # your RAG/LLM core

TITLE = "Askstein — CT Rigidity / FE Q&A"
DESC  = (
    "Enter your first and last name to start chatting. After each answer, leave a 👍 or 👎 — feedback is saved."
)

# ====== Storage (sessions + feedback) ========================================

# Prefer Space persistent storage if mounted
DEFAULT_DATA_DIR = "/data" if os.path.isdir("/data") and os.access("/data", os.W_OK) else "./data"
DATA_DIR = os.path.abspath(os.getenv("DATA_DIR", DEFAULT_DATA_DIR))
DB_PATH  = os.path.join(DATA_DIR, "askstein.db")
os.makedirs(DATA_DIR, exist_ok=True)

def _db():
    conn = sqlite3.connect(DB_PATH, check_same_thread=False)
    conn.execute("PRAGMA journal_mode=WAL;")
    conn.execute("PRAGMA synchronous=NORMAL;")
    return conn

def init_db():
    conn = _db()
    cur = conn.cursor()
    # session: a lightweight record for this browser session
    cur.execute("""
    CREATE TABLE IF NOT EXISTS sessions (
        id INTEGER PRIMARY KEY AUTOINCREMENT,
        first_name TEXT NOT NULL,
        last_name  TEXT NOT NULL,
        created_at TEXT NOT NULL
    )
    """)
    # feedback tied to a session (also stores names redundantly for easier export)
    cur.execute("""
    CREATE TABLE IF NOT EXISTS feedback (
        id INTEGER PRIMARY KEY AUTOINCREMENT,
        session_id INTEGER,
        first_name TEXT NOT NULL,
        last_name  TEXT NOT NULL,
        question TEXT NOT NULL,
        answer_preview TEXT NOT NULL,
        rating INTEGER NOT NULL,   -- +1 = thumbs up, -1 = thumbs down
        created_at TEXT NOT NULL,
        FOREIGN KEY (session_id) REFERENCES sessions(id)
    )
    """)
    conn.commit()
    conn.close()

def create_session(first: str, last: str) -> int:
    conn = _db()
    cur = conn.cursor()
    cur.execute(
        "INSERT INTO sessions (first_name, last_name, created_at) VALUES (?, ?, ?)",
        (first, last, datetime.utcnow().isoformat())
    )
    sid = cur.lastrowid
    conn.commit()
    conn.close()
    return int(sid)

def save_feedback(session_id: Optional[int], first: str, last: str, question: str, answer: str, rating: int) -> Tuple[bool, str]:
    preview = (answer or "").replace("\n", " ").strip()
    if len(preview) > 350:
        preview = preview[:350] + "…"
    try:
        conn = _db()
        conn.execute(
            "INSERT INTO feedback (session_id, first_name, last_name, question, answer_preview, rating, created_at) "
            "VALUES (?, ?, ?, ?, ?, ?, ?)",
            (session_id, first, last, question, preview, int(rating), datetime.utcnow().isoformat())
        )
        conn.commit()
        return True, "Thanks for your feedback!"
    except Exception as e:
        return False, f"Could not save feedback: {e}"
    finally:
        conn.close()

def export_feedback_csv_for_session(session_id: int) -> str:
    ts = int(time.time())
    out_path = os.path.join(DATA_DIR, f"feedback_session_{session_id}_{ts}.csv")
    conn = _db()
    cur = conn.cursor()
    cur.execute("""
        SELECT id, created_at, rating, question, answer_preview, first_name, last_name
        FROM feedback WHERE session_id = ?
        ORDER BY id DESC
    """, (session_id,))
    rows = cur.fetchall()
    conn.close()
    with open(out_path, "w", newline="", encoding="utf-8") as f:
        w = csv.writer(f)
        w.writerow(["id","created_at","rating","question","answer_preview","first_name","last_name"])
        for r in rows:
            w.writerow(r)
    return out_path

def export_all_feedback_csv() -> str:
    ts = int(time.time())
    out_path = os.path.join(DATA_DIR, f"feedback_all_{ts}.csv")
    conn = _db()
    cur = conn.cursor()
    cur.execute("""
        SELECT f.id, f.created_at, f.rating, f.question, f.answer_preview,
               f.first_name, f.last_name, f.session_id
        FROM feedback f
        ORDER BY f.id DESC
    """)
    rows = cur.fetchall()
    conn.close()
    with open(out_path, "w", newline="", encoding="utf-8") as f:
        w = csv.writer(f)
        w.writerow(["id","created_at","rating","question","answer_preview","first_name","last_name","session_id"])
        for r in rows:
            w.writerow(r)
    return out_path

def export_sessions_csv() -> str:
    ts = int(time.time())
    out_path = os.path.join(DATA_DIR, f"sessions_{ts}.csv")
    conn = _db()
    cur = conn.cursor()
    cur.execute("SELECT id, first_name, last_name, created_at FROM sessions ORDER BY id DESC")
    rows = cur.fetchall()
    conn.close()
    with open(out_path, "w", newline="", encoding="utf-8") as f:
        w = csv.writer(f)
        w.writerow(["id","first_name","last_name","created_at"])
        for r in rows:
            w.writerow(r)
    return out_path

# Initialize DB on import
init_db()

ADMIN_TOKEN = os.getenv("ADMIN_TOKEN", "").strip()

# ====== UI ===================================================================

with gr.Blocks(title=TITLE) as demo:
    gr.Markdown(f"## {TITLE}\n{DESC}")

    # Session state
    st_session_id = gr.State(value=None)   # int session id
    st_first_name = gr.State(value="")     # first name
    st_last_name  = gr.State(value="")     # last name
    st_last_q     = gr.State(value="")
    st_last_a     = gr.State(value="")
    st_can_fb     = gr.State(value=False)
    st_admin_ok   = gr.State(value=False)

    # ---- Name Gate ----
    gate_view = gr.Column(visible=True)
    with gate_view:
        with gr.Row():
            first_tb = gr.Textbox(label="First name", placeholder="Ada")
            last_tb  = gr.Textbox(label="Last name",  placeholder="Lovelace")
        start_btn = gr.Button("Start chatting", variant="primary")
        gate_msg  = gr.Markdown("")

    # ---- Chat View ----
    chat_view = gr.Column(visible=False)
    with chat_view:
        welcome_md = gr.Markdown("### Chat")
        with gr.Row():
            inp = gr.Textbox(
                label="Your question",
                placeholder="e.g., How is bending rigidity (EI) computed from a cortical cross-section?",
                lines=3,
            )
        out = gr.Textbox(label="Askstein", lines=12)
        ask_btn = gr.Button("Ask", variant="primary")

        with gr.Row():
            fb_up   = gr.Button("👍 Helpful")
            fb_down = gr.Button("👎 Not helpful")
            fb_status = gr.Markdown("")

        with gr.Row():
            my_export_btn  = gr.Button("Download this session’s feedback (CSV)")
            my_export_file = gr.File(label="Session feedback CSV", visible=False)

    # ---- Admin (optional) ----
    admin_view = gr.Column(visible=True)
    with admin_view:
        gr.Markdown("### Admin (enter token to unlock)")
        admin_token_in = gr.Textbox(label="Admin token", type="password", placeholder="Set ADMIN_TOKEN secret to use")
        admin_unlock   = gr.Button("Unlock Admin")
        admin_status   = gr.Markdown("")
        with gr.Group(visible=False) as admin_controls:
            all_export_btn   = gr.Button("Export ALL feedback (CSV)")
            sessions_export_btn = gr.Button("Export sessions (CSV)")
            all_export_file  = gr.File(label="All feedback CSV", visible=False)
            sessions_export_file = gr.File(label="Sessions CSV", visible=False)

    # ---- Examples ----
    examples = gr.Examples(
        examples=[
            "Define axial rigidity (EA) and how it is estimated from CT-derived cortical masks.",
            "How does torsional rigidity (GJ) relate to polar moment of area in FE pre-processing?",
            "What are typical pitfalls when mapping Hounsfield units to elastic modulus?",
            "What boundary conditions are common in long-bone FE bending simulations?",
        ],
        inputs=inp,
    )

    # ===== Handlers =====

    # Start chatting (create session)
    def on_start(first, last):
        first = (first or "").strip()
        last  = (last  or "").strip()
        if not first or not last or len(first) > 80 or len(last) > 80:
            return gr.Markdown.update(value="Please enter valid first and last names."), None, "", "", gr.update(), gr.update()
        sid = create_session(first, last)
        welcome = f"Welcome, **{first} {last}** — session **#{sid}**."
        return (
            gr.Markdown.update(value=""),
            sid, first, last,
            gr.update(visible=False),   # hide gate
            (gr.update(visible=True), gr.Markdown.update(value=welcome))  # show chat + set welcome
        )

    # Compound output needs to match targets; unpack the group:
    def _start_outputs(res):
        gate_msg_u, sid, first, last, gate_vis, chat_tuple = res
        chat_vis, welcome_u = chat_tuple
        return gate_msg_u, sid, first, last, gate_vis, chat_vis, welcome_u

    start_btn.click(
        lambda f, l: _start_outputs(on_start(f, l)),
        [first_tb, last_tb],
        [gate_msg, st_session_id, st_first_name, st_last_name, gate_view, chat_view, welcome_md]
    )

    # Ask
    def on_ask(session_id, q):
        q = (q or "").strip()
        if not session_id:
            return gr.Textbox.update(value="Please start a session first."), "", "", False
        if not q:
            return gr.Textbox.update(value="Please enter a question."), "", "", False
        try:
            a = ask(q)
        except Exception as e:
            a = f"[runtime error] {e}"
        return gr.Textbox.update(value=a), q, a, True

    ask_btn.click(
        on_ask,
        [st_session_id, inp],
        [out, st_last_q, st_last_a, st_can_fb]
    )
    inp.submit(
        on_ask,
        [st_session_id, inp],
        [out, st_last_q, st_last_a, st_can_fb]
    )

    # Feedback
    def on_feedback(session_id, can_fb, first, last, last_q, last_a, rating):
        if not can_fb or not last_q or not last_a:
            return gr.Markdown.update(value="No recent answer to rate."), False
        ok, msg = save_feedback(session_id, first, last, last_q, last_a, rating)
        return gr.Markdown.update(value=msg), False

    fb_up.click(
        lambda sid, can, fn, ln, q, a: on_feedback(sid, can, fn, ln, q, a, +1),
        [st_session_id, st_can_fb, st_first_name, st_last_name, st_last_q, st_last_a],
        [fb_status, st_can_fb]
    )
    fb_down.click(
        lambda sid, can, fn, ln, q, a: on_feedback(sid, can, fn, ln, q, a, -1),
        [st_session_id, st_can_fb, st_first_name, st_last_name, st_last_q, st_last_a],
        [fb_status, st_can_fb]
    )

    # Export this session’s feedback
    def on_my_export(session_id):
        if not session_id:
            return gr.File.update(visible=False), gr.Markdown.update(value="No active session.")
        path = export_feedback_csv_for_session(session_id)
        return gr.File.update(value=path, visible=True), gr.Markdown.update(value="")

    my_export_btn.click(
        on_my_export,
        [st_session_id],
        [my_export_file, fb_status]
    )

    # Admin unlock
    def on_admin_unlock(token):
        ok = bool(ADMIN_TOKEN) and (token.strip() == ADMIN_TOKEN)
        if ok:
            return True, gr.Markdown.update(value="Admin unlocked ✅"), gr.update(visible=True)
        msg = "Invalid token or ADMIN_TOKEN not set."
        return False, gr.Markdown.update(value=msg), gr.update(visible=False)

    admin_unlock.click(
        on_admin_unlock,
        [admin_token_in],
        [st_admin_ok, admin_status, admin_controls]
    )

    # Admin exports
    def on_export_all(admin_ok):
        if not admin_ok:
            return gr.File.update(visible=False), gr.Markdown.update(value="Admin is locked.")
        path = export_all_feedback_csv()
        return gr.File.update(value=path, visible=True), gr.Markdown.update(value="")

    def on_export_sessions(admin_ok):
        if not admin_ok:
            return gr.File.update(visible=False), gr.Markdown.update(value="Admin is locked.")
        path = export_sessions_csv()
        return gr.File.update(value=path, visible=True), gr.Markdown.update(value="")

    all_export_btn.click(
        on_export_all,
        [st_admin_ok],
        [all_export_file, admin_status]
    )
    sessions_export_btn.click(
        on_export_sessions,
        [st_admin_ok],
        [sessions_export_file, admin_status]
    )

# Launch (HF Spaces will call `demo` automatically if not run as __main__)
if __name__ == "__main__":
    demo.launch(server_name="0.0.0.0", server_port=int(os.getenv("PORT", "7860")))