import gradio as gr
import json
import os
from datetime import date
from transformers import pipeline
from datasets import load_dataset, Dataset, DatasetDict
import huggingface_hub
import tempfile

PROFILE_FILE = "about_me.json"
DAILY_FILE = "daily_status.json"
HF_DATASET_REPO = "shee2205/bot-data"  # CHANGE THIS TO YOUR REPO!

def load_json(path, default):
    if os.path.exists(path):
        with open(path) as f:
            return json.load(f)
    return default

def save_json(path, data):
    with open(path, "w") as f:
        json.dump(data, f, indent=2)

# --- Hugging Face Datasets methods ---
def hf_login():
    # Reads from HF_TOKEN (Spaces env var)
    token = os.environ.get("HF_TOKEN")
    huggingface_hub.login(token)

def hf_load():
    try:
        ds = load_dataset(HF_DATASET_REPO, split="train")
        # Should be 1 row: {profile:..., daily:...}
        row = ds[0]
        return json.loads(row["profile"]), json.loads(row["daily"])
    except Exception as e:
        print("No HF data found, using empty. Reason:", e)
        return {}, {}

def hf_save(profile, daily):
    # Save current data as a single row
    # (delete all, upload new version)
    data = {
        "profile": [json.dumps(profile)],
        "daily": [json.dumps(daily)],
    }
    # Save locally to a tmp dir for upload
    with tempfile.TemporaryDirectory() as tmpdirname:
        path = os.path.join(tmpdirname, "data.jsonl")
        with open(path, "w") as f:
            for i in range(1):
                f.write(json.dumps({ "profile": data["profile"][i], "daily": data["daily"][i] }) + "\n")
        ds = Dataset.from_json(path)
        ds.push_to_hub(HF_DATASET_REPO, private=True, split="train", token=os.environ.get("HF_TOKEN"))

# --- Load from HF Dataset on startup ---
hf_login()
profile, daily = hf_load()
save_json(PROFILE_FILE, profile)
save_json(DAILY_FILE, daily)

def build_context(profile, daily):
    context = []
    for k, v in profile.items():
        context.append(f"{k.capitalize()}: {v}")
    for d in sorted(daily.keys(), reverse=True)[:7]:
        context.append(f"On {d}: {daily[d].get('log','')}")
    return "\n".join(context)

qa_pipeline = pipeline("question-answering", model="deepset/roberta-base-squad2")

def chatbot_qa(user_q):
    context = build_context(profile, daily)
    try:
        result = qa_pipeline(question=user_q, context=context)
        answer = result["answer"].strip()
        if not answer or answer.lower() in ["empty", ""]:
            return "Sheetal hasn't shared that yet!"
        return answer
    except Exception as e:
        return f"Error: {e}"

def save_profile_and_log(name, city, job, about, freeform):
    today = str(date.today())
    profile.update({
        "name": name,
        "city": city,
        "job": job,
        "about": about
    })
    save_json(PROFILE_FILE, profile)
    daily[today] = {"log": freeform}
    save_json(DAILY_FILE, daily)
    # Save to HF dataset as well
    hf_save(profile, daily)
    logs = ""
    for d in sorted(daily.keys(), reverse=True)[:5]:
        logs += f"**{d}**: {daily[d]['log']}\n"
    return (
        "✅ Profile & log updated (saved to cloud)!",
        logs
    )

def recent_logs():
    logs = ""
    for d in sorted(daily.keys(), reverse=True)[:5]:
        logs += f"**{d}**: {daily[d]['log']}\n"
    return logs

ADMIN_PASSWORD = "123"  # CHANGE THIS!

def check_password(password):
    if password == ADMIN_PASSWORD:
        return gr.update(visible=True), gr.update(value="", visible=False)
    else:
        return gr.update(visible=False), gr.update(value="❌ Access denied. Try again.", visible=True)

def show_profile():
    return json.dumps(profile, indent=2), json.dumps(daily, indent=2)

with gr.Blocks(title="Sheetal's Chatbot") as demo:
    gr.Markdown("# 🌸 Sheetal's Personal Chatbot")
    gr.Markdown("Ask anything about Sheetal!")

    with gr.Tab("📝 Admin (Sheetal)"):
        admin_password = gr.Textbox(label="Enter admin password", type="password", max_lines=1)
        unlock_panel = gr.Button("Unlock Admin Panel")
        password_status = gr.Textbox(label="Password status", value="", interactive=False, visible=False)

        with gr.Column(visible=False) as admin_panel:
            with gr.Row():
                name = gr.Textbox(label="Name", value=profile.get("name", "Sheetal"), max_lines=1)
                city = gr.Textbox(label="City", value=profile.get("city", ""), max_lines=1)
                job = gr.Textbox(label="Profession", value=profile.get("job", ""), max_lines=1)
            about = gr.Textbox(label="About You", value=profile.get("about", ""), lines=2, max_lines=3)
            today_log = gr.Textbox(label="What do you want to remember about today?", value=daily.get(str(date.today()), {}).get("log", ""), lines=2, max_lines=3)
            save_btn = gr.Button("💾 Save Profile & Today's Log")
            admin_status = gr.Markdown("")
            logs_output = gr.Markdown(recent_logs())
            save_btn.click(
                fn=save_profile_and_log,
                inputs=[name, city, job, about, today_log],
                outputs=[admin_status, logs_output]
            )
            # Debug: show current data
            show_btn = gr.Button("Show Current Data")
            profile_out = gr.Textbox(label="Profile JSON")
            daily_out = gr.Textbox(label="Daily JSON")
            show_btn.click(fn=show_profile, outputs=[profile_out, daily_out])

        unlock_panel.click(
            fn=check_password,
            inputs=[admin_password],
            outputs=[admin_panel, password_status]
        )

    with gr.Tab("💬 Ask About Sheetal"):
        user_q = gr.Textbox(label="Type your question here:")
        ask_btn = gr.Button("Ask")
        answer_box = gr.Textbox(label="Bot answer", interactive=False, lines=2, max_lines=4)
        ask_btn.click(
            fn=chatbot_qa,
            inputs=user_q,
            outputs=answer_box
        )

if __name__ == "__main__":
    demo.launch()