app.py

import os
import json
import random
import streamlit as st
from typing import Any, Dict, List, Optional, Tuple

# ----------------------------
# Paths
# ----------------------------
APP_DIR = os.path.dirname(os.path.abspath(__file__))
SCENARIOS_DIR = os.path.join(APP_DIR, "scenarios")
IMAGES_DIR = os.path.join(APP_DIR, "images")

os.makedirs(SCENARIOS_DIR, exist_ok=True)
os.makedirs(IMAGES_DIR, exist_ok=True)

IMAGE_EXTENSIONS = (".png", ".jpg", ".jpeg", ".gif", ".webp")

st.set_page_config(page_title="Interactive Training", page_icon="🧠", layout="centered")

ALIGNMENT_POINTS = {"high": 3, "medium": 2, "low": 1}

# ----------------------------
# Scenario JSON Validation (minimal but safe)
# ----------------------------
def normalize_alignment(x: Any) -> str:
    if not isinstance(x, str):
        return "unknown"
    x = x.strip().lower()
    return x if x in ("high", "medium", "low") else "unknown"

def validate_scenario_json(d: Dict[str, Any]) -> List[str]:
    errors = []
    if not isinstance(d, dict):
        return ["Scenario JSON must be a JSON object."]

    required = ["scenario_id", "title", "context", "customer", "objective", "response_options", "ideal_response"]
    for k in required:
        if k not in d:
            errors.append(f"Missing required key: '{k}'")

    if "scenario_id" in d and not isinstance(d["scenario_id"], str):
        errors.append("scenario_id must be a string.")
    if "title" in d and not isinstance(d["title"], str):
        errors.append("title must be a string.")

    ro = d.get("response_options")
    if ro is not None:
        if not isinstance(ro, list) or len(ro) < 2:
            errors.append("response_options must be an array with at least 2 options.")
        else:
            ids = set()
            for i, opt in enumerate(ro):
                if not isinstance(opt, dict):
                    errors.append(f"response_options[{i}] must be an object.")
                    continue
                oid = opt.get("option_id")
                txt = opt.get("text")
                ba = normalize_alignment(opt.get("brand_alignment"))
                if not isinstance(oid, str) or not oid.strip():
                    errors.append(f"response_options[{i}] missing valid option_id.")
                else:
                    if oid in ids:
                        errors.append(f"Duplicate option_id: {oid}")
                    ids.add(oid)
                if not isinstance(txt, str) or not txt.strip():
                    errors.append(f"response_options[{i}] missing valid text.")
                if ba == "unknown":
                    errors.append(f"response_options[{i}].brand_alignment must be high/medium/low.")

    ideal = d.get("ideal_response")
    if isinstance(ideal, dict):
        if not isinstance(ideal.get("option_id"), str) or not ideal.get("option_id"):
            errors.append("ideal_response.option_id must be a string.")
        if not isinstance(ideal.get("explanation"), str) or not ideal.get("explanation"):
            errors.append("ideal_response.explanation must be a string.")
    else:
        errors.append("ideal_response must be an object.")

    return errors

# ----------------------------
# Load scenarios from folder
# ----------------------------
def load_scenarios_from_folder() -> List[Dict[str, Any]]:
    scenarios: List[Dict[str, Any]] = []
    for fname in sorted(os.listdir(SCENARIOS_DIR)):
        if not fname.lower().endswith(".json"):
            continue
        fpath = os.path.join(SCENARIOS_DIR, fname)
        try:
            with open(fpath, "r", encoding="utf-8") as f:
                data = json.load(f)
            # Soft validate; skip broken files but keep app alive
            if isinstance(data, dict) and data.get("scenario_id"):
                scenarios.append(data)
        except Exception:
            # ignore broken files
            pass

    # Ensure stable ordering: by scenario_id then title
    scenarios.sort(key=lambda s: (str(s.get("scenario_id", "")), str(s.get("title", ""))))
    return scenarios

@st.cache_data(show_spinner=False)
def get_scenarios_cached() -> List[Dict[str, Any]]:
    return load_scenarios_from_folder()

def refresh_scenarios_cache() -> None:
    get_scenarios_cached.clear()  # type: ignore

def get_scenario_by_id(scenarios: List[Dict[str, Any]], sid: str) -> Optional[Dict[str, Any]]:
    for s in scenarios:
        if s.get("scenario_id") == sid:
            return s
    return None

def get_option(scenario: Dict[str, Any], option_id: str) -> Optional[Dict[str, Any]]:
    for o in scenario.get("response_options", []):
        if o.get("option_id") == option_id:
            return o
    return None

# ----------------------------
# Images
# ----------------------------
def get_scenario_image_path(scenario_id: str) -> Optional[str]:
    """Return path to scenario image if it exists (e.g. images/SCEN-HH-042.png)."""
    if not scenario_id:
        return None
    base = os.path.join(IMAGES_DIR, scenario_id)
    for ext in IMAGE_EXTENSIONS:
        path = base + ext
        if os.path.isfile(path):
            return path
    return None

def guess_ext_from_filename(name: str) -> str:
    ext = os.path.splitext(name)[1].lower()
    if ext in IMAGE_EXTENSIONS:
        return ext
    return ".png"  # fallback

# ----------------------------
# Session State
# ----------------------------
def reset_run(keep_progress: bool = True) -> None:
    st.session_state.selected_option_id = None
    st.session_state.submitted = False
    st.session_state.last_feedback = None
    if not keep_progress:
        st.session_state.total_score = 0
        st.session_state.attempts = 0
        st.session_state.correct = 0
        st.session_state.completed_ids = set()

def init_state(default_first_id: Optional[str]) -> None:
    if "current_scenario_id" not in st.session_state:
        st.session_state.current_scenario_id = default_first_id
    if "selected_option_id" not in st.session_state:
        st.session_state.selected_option_id = None
    if "submitted" not in st.session_state:
        st.session_state.submitted = False
    if "last_feedback" not in st.session_state:
        st.session_state.last_feedback = None

    # Progress metrics
    if "total_score" not in st.session_state:
        st.session_state.total_score = 0
    if "attempts" not in st.session_state:
        st.session_state.attempts = 0
    if "correct" not in st.session_state:
        st.session_state.correct = 0
    if "completed_ids" not in st.session_state:
        st.session_state.completed_ids = set()

# ----------------------------
# Main
# ----------------------------
scenarios = get_scenarios_cached()
first_id = scenarios[0]["scenario_id"] if scenarios else None
init_state(first_id)

st.title("🧠 Interactive Training App")

if not scenarios:
    st.warning("No scenarios found. Add JSON files into the 'scenarios/' folder or upload one in the sidebar.")

# Sidebar
with st.sidebar:
    st.header("Scenario Library")

    scenario_labels = [f"{s['scenario_id']} · {s['title']}" for s in scenarios]
    id_by_label = {f"{s['scenario_id']} · {s['title']}": s["scenario_id"] for s in scenarios}

    if scenarios and st.session_state.current_scenario_id is None:
        st.session_state.current_scenario_id = scenarios[0]["scenario_id"]

    if scenarios:
        current_label = next(
            (lbl for lbl in scenario_labels if id_by_label[lbl] == st.session_state.current_scenario_id),
            scenario_labels[0]
        )

        chosen_label = st.selectbox("Pick a scenario", scenario_labels, index=scenario_labels.index(current_label))
        st.session_state.current_scenario_id = id_by_label[chosen_label]

        colA, colB = st.columns(2)
        if colA.button("🎲 Random"):
            st.session_state.current_scenario_id = random.choice(scenarios)["scenario_id"]
            reset_run(keep_progress=True)
            st.rerun()

        if colB.button("➡️ Next"):
            ids = [s["scenario_id"] for s in scenarios]
            idx = ids.index(st.session_state.current_scenario_id)
            st.session_state.current_scenario_id = ids[(idx + 1) % len(ids)]
            reset_run(keep_progress=True)
            st.rerun()

    st.divider()
    st.subheader("Progress")
    st.metric("Total score", st.session_state.total_score)
    st.metric("Attempts", st.session_state.attempts)
    acc = (st.session_state.correct / st.session_state.attempts) * 100 if st.session_state.attempts else 0.0
    st.metric("Accuracy", f"{acc:.0f}%")
    st.metric("Completed", f"{len(st.session_state.completed_ids)}/{len(scenarios)}")
    st.caption("Points awarded by alignment")

    # ---- Push upload section "toward bottom" ----
    # Streamlit doesn't have true "bottom-left fixed" placement,
    # but placing it after everything else makes it appear at the bottom of the sidebar.
    st.divider()
    with st.expander("➕ Upload scenario", expanded=False):
        st.caption("Upload a scenario JSON and an optional image. Files will be saved to /scenarios and /images.")

        uploaded_json = st.file_uploader("Scenario JSON", type=["json"], key="upload_json")
        uploaded_img = st.file_uploader("Scenario image (optional)", type=["png", "jpg", "jpeg", "gif", "webp"], key="upload_img")

        overwrite = st.checkbox("Overwrite if scenario_id already exists", value=False)

        if st.button("Save uploaded scenario", type="primary"):
            if uploaded_json is None:
                st.error("Please upload a scenario JSON file.")
            else:
                try:
                    scenario_data = json.load(uploaded_json)
                except Exception as e:
                    st.error(f"Invalid JSON: {e}")
                    st.stop()

                errs = validate_scenario_json(scenario_data)
                if errs:
                    st.error("Scenario JSON has issues:")
                    for e in errs:
                        st.write(f"- {e}")
                    st.stop()

                scenario_id = scenario_data["scenario_id"].strip()
                json_path = os.path.join(SCENARIOS_DIR, f"{scenario_id}.json")

                if os.path.exists(json_path) and not overwrite:
                    st.error(f"Scenario '{scenario_id}' already exists. Enable overwrite to replace it.")
                    st.stop()

                # Save JSON
                try:
                    with open(json_path, "w", encoding="utf-8") as f:
                        json.dump(scenario_data, f, ensure_ascii=False, indent=2)
                except Exception as e:
                    st.error(f"Failed to save JSON: {e}")
                    st.stop()

                # Save image if provided
                if uploaded_img is not None:
                    ext = guess_ext_from_filename(uploaded_img.name)
                    img_path = os.path.join(IMAGES_DIR, f"{scenario_id}{ext}")
                    try:
                        with open(img_path, "wb") as f:
                            f.write(uploaded_img.getbuffer())
                    except Exception as e:
                        st.error(f"Failed to save image: {e}")
                        st.stop()

                # Refresh scenarios and jump to the new one
                refresh_scenarios_cache()
                st.session_state.current_scenario_id = scenario_id
                reset_run(keep_progress=True)
                st.success(f"Saved scenario '{scenario_id}'.")
                st.rerun()

# Main content: show scenario
if scenarios and st.session_state.current_scenario_id:
    scenario = get_scenario_by_id(get_scenarios_cached(), st.session_state.current_scenario_id)
else:
    scenario = None

if not scenario:
    st.info("Select a scenario from the sidebar, or upload one.")
    st.stop()

# Header
st.subheader(f"{scenario['title']}  ·  {scenario['scenario_id']}")

# Scenario image (if present)
img_path = get_scenario_image_path(scenario["scenario_id"])
if img_path:
    st.image(img_path, use_container_width=True)

# Context / Customer / Objective
ctx = scenario.get("context", {})
cust = scenario.get("customer", {})

with st.container(border=True):
    st.markdown("### Context")
    st.write(f"**Time:** {ctx.get('time_of_day', '—')}")
    st.write(f"**Store condition:** {ctx.get('store_condition', '—')}")
    st.write(f"**Role:** {ctx.get('employee_role', '—')}")

with st.container(border=True):
    st.markdown("### Customer")
    st.write(f"**Mood:** {cust.get('mood', '—')}")
    st.write(f"**Opening line:** {cust.get('opening_line', '—')}")

with st.container(border=True):
    st.markdown("### Objective")
    st.write(scenario.get("objective", "—"))

meta_cols = st.columns(2)
skills = scenario.get("skills_practiced", [])
difficulty = scenario.get("difficulty_level", "—")
meta_cols[0].write("**Skills practiced:** " + (", ".join(skills) if isinstance(skills, list) and skills else "—"))
meta_cols[1].write(f"**Difficulty:** {difficulty}")

st.divider()

# Options
ideal_id = scenario.get("ideal_response", {}).get("option_id")
ideal_expl = scenario.get("ideal_response", {}).get("explanation", "")

st.markdown("### Choose your response")

options = scenario.get("response_options", [])
labels = []
label_to_id = {}
seen_texts = set()

for o in options:
    oid = o.get("option_id")
    text = o.get("text", "")
    lbl = text if text not in seen_texts else f"{text} — Option {oid}"
    seen_texts.add(text)
    labels.append(lbl)
    label_to_id[lbl] = oid

selected_label = st.radio(
    "Select one:",
    labels,
    index=None if st.session_state.selected_option_id is None else next(
        (i for i, o in enumerate(options) if o.get("option_id") == st.session_state.selected_option_id),
        None
    ),
    disabled=st.session_state.submitted,
)

if selected_label:
    st.session_state.selected_option_id = label_to_id[selected_label]

btn_cols = st.columns([1, 1, 2])
submit = btn_cols[0].button(
    "Submit",
    type="primary",
    disabled=(st.session_state.selected_option_id is None or st.session_state.submitted),
)
retry = btn_cols[1].button("Try again", disabled=not st.session_state.submitted)

if retry:
    reset_run(keep_progress=True)
    st.rerun()

if submit:
    st.session_state.submitted = True
    st.session_state.attempts += 1

    chosen_opt = get_option(scenario, st.session_state.selected_option_id)
    if not chosen_opt:
        st.error("Chosen option not found.")
        st.stop()

    align = normalize_alignment(chosen_opt.get("brand_alignment"))
    points = ALIGNMENT_POINTS.get(align, 0)
    st.session_state.total_score += points

    correct = (st.session_state.selected_option_id == ideal_id)
    if correct:
        st.session_state.correct += 1
        st.success(f"✅ Ideal response! (+{points})")
    else:
        st.warning(f"⚠️ Not ideal. (+{points})")

    st.session_state.completed_ids.add(scenario["scenario_id"])

    st.markdown("### Feedback")
    if align == "high":
        st.write("Strong brand alignment: warm, clear, and ownership-driven.")
    elif align == "medium":
        st.write("Moderate brand alignment: clear, but could show more empathy/ownership or offer better alternatives.")
    elif align == "low":
        st.write("Low brand alignment: may feel vague, dismissive, or not customer-centered under pressure.")
    else:
        st.write("Brand alignment unknown.")

    ideal_opt = get_option(scenario, ideal_id) if ideal_id else None
    ideal_text = ideal_opt.get("text", "") if ideal_opt else f"Option {ideal_id}"
    st.info(f"✅ The ideal response was: **{ideal_text}**")
    if ideal_expl:
        st.markdown("### Why the ideal response works")
        st.write(ideal_expl)

    st.markdown("### Compare all options")
    for o in options:
        oid = o.get("option_id")
        is_ideal = (oid == ideal_id)
        is_chosen = (oid == st.session_state.selected_option_id)
        tag = " ⭐ Ideal" if is_ideal else ""
        tag2 = " 👉 Your choice" if is_chosen else ""
        with st.container(border=True):
            st.write(o.get("text", "") + tag + tag2)

    # Convenience navigation
    st.divider()
    nav_cols = st.columns(2)
    if nav_cols[0].button("🎲 Random scenario"):
        st.session_state.current_scenario_id = random.choice(get_scenarios_cached())["scenario_id"]
        reset_run(keep_progress=True)
        st.rerun()
    if nav_cols[1].button("➡️ Next scenario"):
        ids = [s["scenario_id"] for s in get_scenarios_cached()]
        idx = ids.index(st.session_state.current_scenario_id)
        st.session_state.current_scenario_id = ids[(idx + 1) % len(ids)]
        reset_run(keep_progress=True)
        st.rerun()