app.py

import streamlit as st
import json
import pandas as pd
import os
from uuid import uuid4

st.set_page_config(page_title="Dataset Builder", layout="wide")
st.title("📚 JSONL Dataset Editor")

TMP_DIR = "temp"
TMP_FILE = os.path.join(TMP_DIR, "session_dataset.jsonl")
os.makedirs(TMP_DIR, exist_ok=True)


# --- Helpers ---
def get_all_fields(data):
    keys = set()
    for d in data:
        keys.update(d.keys())
    return sorted(list(keys))


def save_to_file():
    with open(TMP_FILE, "w", encoding="utf-8") as f:
        for row in st.session_state.data:
            f.write(json.dumps(row, ensure_ascii=False) + "\n")


# --- Session Initialization ---
if "data" not in st.session_state:
    st.session_state.data = []
if "all_fields" not in st.session_state:
    st.session_state.all_fields = []
if "editor_key" not in st.session_state:
    st.session_state.editor_key = str(uuid4())
if "just_added" not in st.session_state:
    st.session_state.just_added = False

# --- Load from TMP file ---
if os.path.exists(TMP_FILE) and not st.session_state.data:
    with open(TMP_FILE, "r", encoding="utf-8") as f:
        st.session_state.data = [json.loads(line) for line in f]
    st.session_state.all_fields = get_all_fields(st.session_state.data)

# --- File Upload ---
uploaded = st.file_uploader("Upload JSONL", type=["jsonl"])
if uploaded:
    lines = uploaded.read().decode("utf-8").splitlines()
    st.session_state.data = [json.loads(l) for l in lines]
    st.session_state.all_fields = get_all_fields(st.session_state.data)
    save_to_file()
    st.session_state.editor_key = str(uuid4())
    st.rerun()

# --- Add New Fields fallback ---
if not st.session_state.all_fields:
    st.session_state.all_fields = ["context", "question", "answer"]

# --- Add Entry Form ---
with st.form("add_form"):
    st.markdown("### ➕ Add New Entry")
    new_entry = {}
    for field in st.session_state.all_fields:
        new_entry[field] = st.text_area(field, key=f"add_{field}")

    submit_add = st.form_submit_button("Add Entry")

if submit_add:
    st.session_state.data.append(new_entry)
    save_to_file()
    st.session_state.editor_key = str(uuid4())
    st.session_state.just_added = True
    st.rerun()

# --- Wait for rerun before rendering editor ---
if st.session_state.just_added:
    st.session_state.just_added = False
    st.rerun()

# --- Display Editor ---
st.markdown("### ✏️ Edit Records")
df = pd.DataFrame(st.session_state.data)
df = df.reindex(columns=st.session_state.all_fields)

for field in st.session_state.all_fields:
    df[field] = df[field].astype(str)

column_configs = {
    field: st.column_config.TextColumn(field, width="large")
    if field.lower() in ["context", "question", "answer"]
    else None
    for field in st.session_state.all_fields
}

edited_df = st.data_editor(
    df,
    key=st.session_state.editor_key,
    use_container_width=True,
    num_rows="dynamic",
    column_config=column_configs,
)

# --- Save edits ---
if edited_df is not None:
    new_data = edited_df.fillna("").to_dict(orient="records")
    if new_data != st.session_state.data:
        st.session_state.data = new_data
        save_to_file()
        st.toast("✅ Changes auto-saved!", icon="💾")

# --- Add New Field ---
with st.expander("➕ Add New Field"):
    new_field = st.text_input("New field name", key="new_field_name")
    if st.button("Add Field"):
        if new_field and new_field not in st.session_state.all_fields:
            st.session_state.all_fields.append(new_field)
            st.rerun()

# --- Export Dataset ---
st.markdown("### 📤 Export")
export_path = st.text_input("Save path", value="./exports/exported_dataset.jsonl")

col1, col2, col3 = st.columns(3)

with col1:
    if st.button("📁 Export JSONL"):
        os.makedirs(os.path.dirname(export_path), exist_ok=True)
        with open(export_path, "w", encoding="utf-8") as f:
            for row in st.session_state.data:
                f.write(json.dumps(row, ensure_ascii=False) + "\n")
        with open(export_path, "r", encoding="utf-8") as f:
            content = f.read()
        st.download_button("⬇️ Download JSONL", content, file_name=os.path.basename(export_path))
        st.success("✅ Exported!")
        st.session_state.clear()
        if os.path.exists(TMP_FILE):
            os.remove(TMP_FILE)
        st.rerun()

with col2:
    if os.path.exists(TMP_FILE):
        with open(TMP_FILE, "r", encoding="utf-8") as f:
            tmp_data = f.read()
        st.download_button("⬇️ Temp File", tmp_data, file_name="session_dataset.jsonl")

with col3:
    if st.button("🧹 Clear Session"):
        if os.path.exists(TMP_FILE):
            os.remove(TMP_FILE)
        st.session_state.clear()
        st.rerun()