# =========================
# ONE CELL: SmartCare HMS (Student Edition)
# Patients + Doctors (with timing) + Appointments (search doctor by time + select) + Billing + Groq AI + Login
# =========================


import os
import sqlite3
from datetime import datetime, date

import gradio as gr
from groq import Groq

# -------------------- (OPTIONAL) SET GROQ KEY HERE --------------------
# Paste your key (recommended) OR set it in a separate cell:
# os.environ["GROQ_API_KEY"] = "PASTE_YOUR_KEY"
# ---------------------------------------------------------------------
# os.environ["GROQ_API_KEY"] = "PASTE_YOUR_GROQ_KEY_HERE"

# -------------------- DB --------------------
DB_PATH = "hms.db"

def get_conn():
    return sqlite3.connect(DB_PATH, check_same_thread=False)

def _now():
    return datetime.now().strftime("%Y-%m-%d %H:%M:%S")

def not_empty(s):
    return s is not None and str(s).strip() != ""

def is_int(x):
    try:
        int(str(x).strip())
        return True
    except:
        return False

def is_float(x):
    try:
        float(str(x).strip())
        return True
    except:
        return False

def parse_date_yyyy_mm_dd(s):
    try:
        return datetime.strptime(str(s).strip(), "%Y-%m-%d")
    except:
        return None

def parse_time_hh_mm(s):
    try:
        return datetime.strptime(str(s).strip(), "%H:%M")
    except:
        return None

def init_db():
    conn = get_conn()
    cur = conn.cursor()

    # patients
    cur.execute("""
    CREATE TABLE IF NOT EXISTS patients (
        id INTEGER PRIMARY KEY AUTOINCREMENT,
        name TEXT NOT NULL,
        age INTEGER NOT NULL,
        gender TEXT NOT NULL,
        phone TEXT NOT NULL,
        address TEXT,
        created_at TEXT NOT NULL
    )""")

    # doctors
    cur.execute("""
    CREATE TABLE IF NOT EXISTS doctors (
        id INTEGER PRIMARY KEY AUTOINCREMENT,
        name TEXT NOT NULL,
        specialization TEXT NOT NULL,
        phone TEXT NOT NULL,
        available_from TEXT NOT NULL DEFAULT '09:00',
        available_to   TEXT NOT NULL DEFAULT '17:00',
        created_at TEXT NOT NULL
    )""")

    # appointments
    cur.execute("""
    CREATE TABLE IF NOT EXISTS appointments (
        id INTEGER PRIMARY KEY AUTOINCREMENT,
        patient_id INTEGER NOT NULL,
        doctor_id INTEGER NOT NULL,
        appt_date TEXT NOT NULL,
        appt_time TEXT NOT NULL,
        reason TEXT,
        status TEXT NOT NULL DEFAULT 'BOOKED',
        created_at TEXT NOT NULL
    )""")

    # bills
    cur.execute("""
    CREATE TABLE IF NOT EXISTS bills (
        id INTEGER PRIMARY KEY AUTOINCREMENT,
        patient_id INTEGER NOT NULL,
        services TEXT NOT NULL,
        total_amount REAL NOT NULL,
        created_at TEXT NOT NULL
    )""")

    conn.commit()
    conn.close()

# -------------------- DELETE FUNCTIONS --------------------

def delete_patient(pid):
    if not is_int(pid):
        return "❌ Patient ID must be a number."
    pid = int(pid)

    conn = get_conn()
    cur = conn.cursor()

    # check exists
    cur.execute("SELECT id FROM patients WHERE id=?", (pid,))
    if cur.fetchone() is None:
        conn.close()
        return "⚠️ Patient not found."

    # also delete related appointments & bills
    cur.execute("DELETE FROM appointments WHERE patient_id=?", (pid,))
    cur.execute("DELETE FROM bills WHERE patient_id=?", (pid,))
    cur.execute("DELETE FROM patients WHERE id=?", (pid,))
    conn.commit()
    conn.close()

    return "🗑️ Patient Deleted Successfully"


def delete_doctor(did):
    if not is_int(did):
        return "❌ Doctor ID must be a number."
    did = int(did)

    conn = get_conn()
    cur = conn.cursor()

    cur.execute("SELECT id FROM doctors WHERE id=?", (did,))
    if cur.fetchone() is None:
        conn.close()
        return "⚠️ Doctor not found."

    # delete related appointments
    cur.execute("DELETE FROM appointments WHERE doctor_id=?", (did,))
    cur.execute("DELETE FROM doctors WHERE id=?", (did,))
    conn.commit()
    conn.close()

    return "🗑️ Doctor Deleted Successfully"


def delete_appointment(aid):
    if not is_int(aid): return "❌ Appointment ID must be a number."
    aid = int(aid)

    conn = get_conn()
    cur = conn.cursor()
    cur.execute("DELETE FROM appointments WHERE id=?", (aid,))
    conn.commit()
    deleted = cur.rowcount
    conn.close()

    return "🗑️ Appointment Deleted" if deleted else "⚠️ Appointment not found."


def delete_bill(bid):
    if not is_int(bid): return "❌ Bill ID must be a number."
    bid = int(bid)

    conn = get_conn()
    cur = conn.cursor()
    cur.execute("DELETE FROM bills WHERE id=?", (bid,))
    conn.commit()
    deleted = cur.rowcount
    conn.close()

    return "🗑️ Bill Deleted" if deleted else "⚠️ Bill not found."

init_db()

# -------------------- Patients --------------------
def add_patient(name, age, gender, phone, address):
    if not not_empty(name): return "❌ Name is required."
    if not is_int(age): return "❌ Age must be a number."
    age = int(age)
    if age <= 0 or age > 120: return "❌ Age must be between 1 and 120."
    if gender not in ["Male", "Female", "Other"]: return "❌ Select gender."
    if not not_empty(phone): return "❌ Phone required."

    conn = get_conn()
    cur = conn.cursor()
    cur.execute("""INSERT INTO patients (name, age, gender, phone, address, created_at)
                   VALUES (?, ?, ?, ?, ?, ?)""",
                (name.strip(), age, gender, phone.strip(), (address or "").strip(), _now()))
    conn.commit()
    conn.close()
    return "✅ Patient Added 🎉"

def list_patients():
    conn = get_conn()
    cur = conn.cursor()
    cur.execute("SELECT id, name, age, gender, phone, address, created_at FROM patients ORDER BY id DESC")
    rows = cur.fetchall()
    conn.close()
    return rows

# -------------------- Doctors (WITH TIMING) --------------------
def add_doctor(name, specialization, phone, available_from, available_to):
    if not not_empty(name): return "❌ Doctor name required."
    if not not_empty(specialization): return "❌ Specialization required."
    if not not_empty(phone): return "❌ Phone required."

    t1 = parse_time_hh_mm(available_from)
    t2 = parse_time_hh_mm(available_to)
    if t1 is None or t2 is None:
        return "❌ Timing must be HH:MM (e.g., 09:00 to 17:00)."
    if t1 >= t2:
        return "❌ 'From' time must be earlier than 'To' time."

    conn = get_conn()
    cur = conn.cursor()
    cur.execute("""INSERT INTO doctors (name, specialization, phone, available_from, available_to, created_at)
                   VALUES (?, ?, ?, ?, ?, ?)""",
                (name.strip(), specialization.strip(), phone.strip(),
                 available_from.strip(), available_to.strip(), _now()))
    conn.commit()
    conn.close()
    return "✅ Doctor Added 🩺"

def list_doctors():
    conn = get_conn()
    cur = conn.cursor()
    cur.execute("SELECT id, name, specialization, phone, available_from, available_to, created_at FROM doctors ORDER BY id DESC")
    rows = cur.fetchall()
    conn.close()
    return rows

def search_doctors_by_time(time_str):
    """
    Returns:
      - table rows of doctors available at time_str
      - dropdown choices [(label, value), ...]
      - status msg
    """
    t = parse_time_hh_mm(time_str)
    if t is None:
        return [], gr.update(choices=[], value=None), "❌ Time format must be HH:MM (e.g., 14:30)."

    # HH:MM lexicographic compare works if zero-padded
    ts = time_str.strip()

    conn = get_conn()
    cur = conn.cursor()
    cur.execute("""
        SELECT id, name, specialization, phone, available_from, available_to
        FROM doctors
        WHERE available_from <= ? AND available_to >= ?
        ORDER BY id DESC
    """, (ts, ts))
    rows = cur.fetchall()
    conn.close()

    choices = []
    for r in rows:
        did, name, spec, phone, f, to = r
        label = f"{did} - {name} ({spec}) [{f}-{to}]"
        choices.append((label, str(did)))

    msg = f"✅ {len(rows)} doctor(s) found for {ts}." if rows else f"⚠️ No doctor available at {ts}."
    return rows, gr.update(choices=choices, value=(choices[0][1] if choices else None)), msg

# -------------------- Appointments (SEARCH + SELECT DOCTOR + BOOK) --------------------
def book_appointment_admin(patient_id, selected_doctor_id, appt_date, appt_time, reason):
    # patient id
    if not is_int(patient_id):
        return "❌ Patient ID must be a number."

    # doctor dropdown value is string id
    if not not_empty(selected_doctor_id) or not is_int(selected_doctor_id):
        return "❌ Please select a doctor from the search list."

    pid = int(patient_id)
    did = int(selected_doctor_id)

    # validate date/time formats
    d = parse_date_yyyy_mm_dd(appt_date)
    if d is None:
        return "❌ Date format must be YYYY-MM-DD (e.g., 2026-02-19)."

    t = parse_time_hh_mm(appt_time)
    if t is None:
        return "❌ Time format must be HH:MM (e.g., 14:30)."

    if not not_empty(reason):
        return "❌ Reason is required (short text)."

    conn = get_conn()
    cur = conn.cursor()

    # patient exists?
    cur.execute("SELECT id FROM patients WHERE id=?", (pid,))
    if cur.fetchone() is None:
        conn.close()
        return "⚠️ Patient not found."

    # doctor exists + timing check
    cur.execute("SELECT available_from, available_to FROM doctors WHERE id=?", (did,))
    row = cur.fetchone()
    if row is None:
        conn.close()
        return "⚠️ Doctor not found."

    doc_from, doc_to = row[0], row[1]
    # ensure appointment time inside availability
    if not (doc_from <= appt_time.strip() <= doc_to):
        conn.close()
        return f"⚠️ Doctor available time is {doc_from} to {doc_to}. Choose a time within this range."

    # conflict check: same doctor/date/time and still booked
    cur.execute("""SELECT id FROM appointments
                   WHERE doctor_id=? AND appt_date=? AND appt_time=? AND status='BOOKED'""",
                (did, appt_date.strip(), appt_time.strip()))
    if cur.fetchone() is not None:
        conn.close()
        return "⚠️ Slot already booked for this doctor. Choose another time."

    cur.execute("""INSERT INTO appointments (patient_id, doctor_id, appt_date, appt_time, reason, status, created_at)
                   VALUES (?, ?, ?, ?, ?, 'BOOKED', ?)""",
                (pid, did, appt_date.strip(), appt_time.strip(), reason.strip(), _now()))
    conn.commit()
    conn.close()

    return "✅ Appointment Booked 📅"

def list_appointments():
    conn = get_conn()
    cur = conn.cursor()
    cur.execute("""
        SELECT a.id, a.patient_id, p.name, a.doctor_id, d.name, a.appt_date, a.appt_time, a.status, a.reason
        FROM appointments a
        JOIN patients p ON p.id = a.patient_id
        JOIN doctors d ON d.id = a.doctor_id
        ORDER BY a.id DESC
    """)
    rows = cur.fetchall()
    conn.close()
    return rows

def cancel_appointment(appt_id):
    if not is_int(appt_id): return "❌ Appointment ID must be a number."
    aid = int(appt_id)
    conn = get_conn()
    cur = conn.cursor()
    cur.execute("UPDATE appointments SET status='CANCELLED' WHERE id=?", (aid,))
    conn.commit()
    updated = cur.rowcount
    conn.close()
    return "✅ Cancelled 🧾" if updated else "⚠️ Appointment not found."

# -------------------- Billing --------------------
def create_bill(patient_id, services_csv, amounts_csv):
    if not is_int(patient_id): return "❌ Patient ID must be a number."
    pid = int(patient_id)

    if not not_empty(services_csv): return "❌ Services required."
    if not not_empty(amounts_csv): return "❌ Amounts required."

    services = [s.strip() for s in services_csv.split(",") if s.strip()]
    amounts_raw = [a.strip() for a in amounts_csv.split(",") if a.strip()]

    if len(services) == 0: return "❌ Add at least 1 service."
    if len(services) != len(amounts_raw):
        return "❌ Services count must match amounts count."

    amounts = []
    for a in amounts_raw:
        if not is_float(a): return f"❌ Amount '{a}' is not a number."
        val = float(a)
        if val < 0: return "❌ Amount cannot be negative."
        amounts.append(val)

    total = sum(amounts)

    conn = get_conn()
    cur = conn.cursor()
    cur.execute("SELECT id FROM patients WHERE id=?", (pid,))
    if cur.fetchone() is None:
        conn.close()
        return "⚠️ Patient not found."

    nice_services = " | ".join([f"{s} (Rs {amt:.0f})" for s, amt in zip(services, amounts)])

    cur.execute("""INSERT INTO bills (patient_id, services, total_amount, created_at)
                   VALUES (?, ?, ?, ?)""",
                (pid, nice_services, total, _now()))
    conn.commit()
    conn.close()
    return f"✅ Bill Created 💸 Total: Rs {total:.0f}"

def list_bills():
    conn = get_conn()
    cur = conn.cursor()
    cur.execute("""
        SELECT b.id, b.patient_id, p.name, b.services, b.total_amount, b.created_at
        FROM bills b
        JOIN patients p ON p.id = b.patient_id
        ORDER BY b.id DESC
    """
    )
    rows = cur.fetchall()
    conn.close()
    return rows

# -------------------- Reports --------------------
def reports_summary():
    return {
        "Total Patients": len(list_patients()),
        "Total Doctors": len(list_doctors()),
        "Total Appointments": len(list_appointments()),
        "Total Bills": len(list_bills()),
    }

# -------------------- Groq (SAFE + EXACT STYLE) --------------------
def groq_assistant(task_type, input_text):
    if not not_empty(input_text):
        return "⚠️ Please enter some text first."

    if not os.environ.get("GROQ_API_KEY"):
        return "⚠️ GROQ_API_KEY not set. Paste your key: os.environ['GROQ_API_KEY']='...'."

    system_note = (
        "You are a helpful hospital assistant for an undergraduate demo app. "
        "Do NOT provide medical diagnosis. Provide general educational info only. "
        "Always add: 'This is not medical advice. Please consult a qualified doctor.'"
    )

    prompt_map = {
        "Explain Diagnosis (Simple)": f"Explain this in simple, patient-friendly words:\n\n{input_text}",
        "Summarize Visit Notes": f"Summarize these visit notes into short bullet points:\n\n{input_text}",
        "Prescription Instructions": f"Rewrite these prescription instructions in clear steps:\n\n{input_text}",
        "Generate Discharge Advice": f"Write generic discharge advice (educational) for:\n\n{input_text}",
    }
    user_prompt = prompt_map.get(task_type, input_text)

    try:
        client = Groq(api_key=os.environ.get("GROQ_API_KEY"))
        chat_completion = client.chat.completions.create(
          messages=[
            {"role": "system", "content": system_note},
            {"role": "user", "content": user_prompt},
          ],
          model="llama-3.3-70b-versatile",
        )
        return chat_completion.choices[0].message.content
    except Exception as e:
        return f"❌ Groq API Error: {e}"

# -------------------- Login (Demo) --------------------
PINS = {"Admin": "1234", "Receptionist": "1111", "Doctor": "2222"}

def login(role, pin):
    if role not in PINS:
        return False, "❌ Invalid role."
    if str(pin).strip() != PINS[role]:
        return False, "❌ Wrong PIN."
    return True, f"✅ Logged in as {role} 🎉"

# -------------------- UI (GenZ + HCI) --------------------
CSS = """
:root { --radius: 18px; }
.gradio-container { font-family: ui-sans-serif, system-ui, -apple-system, Segoe UI, Roboto, Arial; }
h1,h2,h3 { letter-spacing: -0.5px; }
"""

with gr.Blocks(css=CSS, title="SmartCare HMS (Student Edition)") as demo:
    auth_state = gr.State({"logged_in": False, "role": None})

    gr.Markdown("# 🏥 SmartCare HMS (Student Edition)\nGen-Z clean UI ✨ + HCI-friendly flows ✅")

    # LOGIN PANEL
    with gr.Row():
        role_in = gr.Dropdown(["Admin", "Receptionist", "Doctor"], value="Admin", label="Role")
        pin_in = gr.Textbox(label="PIN", placeholder="Enter PIN", type="password")
        btn_login = gr.Button("Login 🔐", variant="primary")
    login_msg = gr.Textbox(label="Status", interactive=False)

    app_wrap = gr.Column(visible=False)

    def do_login(role, pin, state):
        ok, msg = login(role, pin)
        state = dict(state)
        state["logged_in"] = ok
        state["role"] = role if ok else None
        return state, msg, gr.update(visible=ok)

    btn_login.click(do_login, inputs=[role_in, pin_in, auth_state], outputs=[auth_state, login_msg, app_wrap])

    # APP (Hidden until login)
    with app_wrap:
        with gr.Tab("✨ Dashboard"):
            btn_refresh = gr.Button("🔄 Refresh Stats", variant="primary")
            stats = gr.JSON(label="Hospital Snapshot")
            btn_refresh.click(fn=reports_summary, outputs=stats)

        with gr.Tab("🧑‍🤝‍🧑 Patients"):
            with gr.Row():
                p_name = gr.Textbox(label="Name", placeholder="e.g., Ali Khan")
                p_age = gr.Textbox(label="Age", placeholder="e.g., 21")
                p_gender = gr.Dropdown(["Male","Female","Other"], value="Male", label="Gender")
            with gr.Row():
                p_phone = gr.Textbox(label="Phone", placeholder="e.g., 03xx-xxxxxxx")
                p_addr = gr.Textbox(label="Address", placeholder="e.g., Lahore (optional)")
            btn_add_p = gr.Button("➕ Add Patient", variant="primary")
            p_msg = gr.Textbox(label="Status", interactive=False)
            btn_add_p.click(add_patient, [p_name, p_age, p_gender, p_phone, p_addr], p_msg)

            btn_list_p = gr.Button("📋 Show Patients")
            p_table = gr.Dataframe(interactive=False,
                headers=["ID","Name","Age","Gender","Phone","Address","Created At"])
            btn_list_p.click(list_patients, outputs=p_table)

            with gr.Row():
                del_p_id = gr.Textbox(label="Delete Patient ID", placeholder="e.g., 1")
                btn_del_p = gr.Button("🗑️ Delete Patient", variant="stop")
            del_p_msg = gr.Textbox(label="Delete Status", interactive=False)

            btn_del_p.click(delete_patient, del_p_id, del_p_msg)

        with gr.Tab("🩺 Doctors"):
            gr.Markdown("### Add Doctor (with availability timing)")
            with gr.Row():
                d_name = gr.Textbox(label="Doctor Name", placeholder="e.g., Dr. Sara")
                d_spec = gr.Textbox(label="Specialization", placeholder="e.g., Cardiologist")
            with gr.Row():
                d_phone = gr.Textbox(label="Phone", placeholder="e.g., 03xx-xxxxxxx")
                d_from = gr.Textbox(label="Available From (HH:MM)", value="09:00")
                d_to   = gr.Textbox(label="Available To (HH:MM)", value="17:00")

            btn_add_d = gr.Button("➕ Add Doctor", variant="primary")
            d_msg = gr.Textbox(label="Status", interactive=False)
            btn_add_d.click(add_doctor, [d_name, d_spec, d_phone, d_from, d_to], d_msg)

            btn_list_d = gr.Button("📋 Show Doctors")
            d_table = gr.Dataframe(interactive=False,
                headers=["ID","Name","Specialization","Phone","From","To","Created At"])
            btn_list_d.click(list_doctors, outputs=d_table)

            with gr.Row():
                del_d_id = gr.Textbox(label="Delete Doctor ID", placeholder="e.g., 1")
                btn_del_d = gr.Button("🗑️ Delete Doctor", variant="stop")
            del_d_msg = gr.Textbox(label="Delete Status", interactive=False)

            btn_del_d.click(delete_doctor, del_d_id, del_d_msg)

        with gr.Tab("📅 Appointments"):
            gr.Markdown("### 1) Search Doctor by Time → 2) Select Doctor → 3) Add Patient ID + Date + Time + Reason → Book")

            # search by time
            with gr.Row():
                search_time = gr.Textbox(label="Search Time (HH:MM)", placeholder="e.g., 14:30")
                btn_search_doc = gr.Button("🔎 Search Doctors", variant="primary")
            doc_search_msg = gr.Textbox(label="Search Status", interactive=False)

            doc_table = gr.Dataframe(
                interactive=False,
                headers=["ID","Name","Specialization","Phone","From","To"]
            )

            doc_select = gr.Dropdown(
                choices=[],
                label="Select Doctor from Results"
            )

            btn_search_doc.click(
                fn=search_doctors_by_time,
                inputs=[search_time],
                outputs=[doc_table, doc_select, doc_search_msg]
            )

            # booking
            with gr.Row():
                ap_pid = gr.Textbox(label="Patient ID", placeholder="e.g., 1")
                ap_date = gr.Textbox(label="Date (YYYY-MM-DD)", value=str(date.today()))
                ap_time = gr.Textbox(label="Time (HH:MM)", placeholder="e.g., 14:30")
            ap_reason = gr.Textbox(label="Reason (required)", placeholder="e.g., fever follow-up / routine checkup")

            btn_book = gr.Button("✅ Book Appointment", variant="primary")
            ap_msg = gr.Textbox(label="Booking Status", interactive=False)

            btn_book.click(
                fn=book_appointment_admin,
                inputs=[ap_pid, doc_select, ap_date, ap_time, ap_reason],
                outputs=ap_msg
            )

            # list appointments
            btn_list_ap = gr.Button("📋 Show Appointments")
            ap_table = gr.Dataframe(interactive=False,
                headers=["Appt ID","Patient ID","Patient Name","Doctor ID","Doctor Name","Date","Time","Status","Reason"])
            btn_list_ap.click(list_appointments, outputs=ap_table)

            # cancel
            with gr.Row():
                ap_cancel_id = gr.Textbox(label="Cancel Appointment ID", placeholder="e.g., 3")
                btn_cancel = gr.Button("🧾 Cancel", variant="stop")
            cancel_msg = gr.Textbox(label="Cancel Status", interactive=False)
            btn_cancel.click(cancel_appointment, ap_cancel_id, cancel_msg)

            with gr.Row():
                del_ap_id = gr.Textbox(label="Delete Appointment ID", placeholder="e.g., 2")
                btn_del_ap = gr.Button("🗑️ Delete Appointment", variant="stop")
            del_ap_msg = gr.Textbox(label="Delete Status", interactive=False)

            btn_del_ap.click(delete_appointment, del_ap_id, del_ap_msg)

        with gr.Tab("💸 Billing"):
            b_pid = gr.Textbox(label="Patient ID", placeholder="e.g., 1")
            b_services = gr.Textbox(label="Services (comma-separated)", placeholder="Consultation, X-Ray, Medicine")
            b_amounts = gr.Textbox(label="Amounts (comma-separated)", placeholder="1000, 1500, 800")
            btn_bill = gr.Button("🧾 Generate Bill", variant="primary")
            b_msg = gr.Textbox(label="Status", interactive=False)
            btn_bill.click(create_bill, [b_pid, b_services, b_amounts], b_msg)

            btn_list_b = gr.Button("📋 Show Bills")
            b_table = gr.Dataframe(interactive=False,
                headers=["Bill ID","Patient ID","Patient Name","Services","Total","Created At"])
            btn_list_b.click(list_bills, outputs=b_table)

            with gr.Row():
                del_b_id = gr.Textbox(label="Delete Bill ID", placeholder="e.g., 1")
                btn_del_b = gr.Button("🗑️ Delete Bill", variant="stop")
            del_b_msg = gr.Textbox(label="Delete Status", interactive=False)

            btn_del_b.click(delete_bill, del_b_id, del_b_msg)

        with gr.Tab("🤖 AI Assistant"):
            task = gr.Dropdown(
                ["Explain Diagnosis (Simple)", "Summarize Visit Notes", "Prescription Instructions", "Generate Discharge Advice"],
                value="Summarize Visit Notes",
                label="Task"
            )
            ai_input = gr.Textbox(label="Input Text", lines=6, placeholder="Paste notes here…")
            btn_ai = gr.Button("⚡ Generate", variant="primary")
            ai_output = gr.Textbox(label="AI Output", lines=10)
            btn_ai.click(groq_assistant, [task, ai_input], ai_output)

demo.launch(share=True)