# app.py
import os
import time
import json
import uuid
import sqlite3
from datetime import datetime
import requests
import streamlit as st
from fpdf import FPDF
import re

# -----------------------
# Config
# -----------------------
APP_TITLE = "L&D Studio AI"
DB_PATH = "history.db"
# prefer repo asset path; the app will try multiple locations
DEFAULT_ASSET_REL_PATH = "assets/logo.png"
UPLOADED_LOCAL_PATH = '/mnt/data/A_logo_for_"L&D_Studio_AI"_is_displayed_in_the_ima.png'
GROQ_API_KEY = os.getenv("GROQ_API_KEY", "")  # set this in environment / platform secrets
GROQ_ENDPOINT = "https://api.groq.com/openai/v1/chat/completions"
GROQ_MODELS_TO_TRY = ["llama-3.1-8b-instant"]

# -----------------------
# DB helpers (history)
# -----------------------
def init_db():
    conn = sqlite3.connect(DB_PATH)
    c = conn.cursor()
    c.execute('''
        CREATE TABLE IF NOT EXISTS sessions (
            id TEXT PRIMARY KEY,
            timestamp TEXT,
            user TEXT,
            name TEXT,
            industry TEXT,
            payload TEXT,
            exported INTEGER
        )
    ''')
    conn.commit()
    conn.close()

def save_session(name, industry, payload, user="demo", exported=0):
    conn = sqlite3.connect(DB_PATH)
    c = conn.cursor()
    sid = str(uuid.uuid4())
    # table has 7 columns -> provide exported value
    c.execute('INSERT INTO sessions VALUES (?,?,?,?,?,?,?)',
              (sid, datetime.utcnow().isoformat(), user, name, industry, json.dumps(payload), exported))
    conn.commit()
    conn.close()
    return sid

def fetch_history(filter_name=None, limit=100):
    conn = sqlite3.connect(DB_PATH)
    c = conn.cursor()
    if filter_name:
        c.execute("SELECT id,timestamp,user,name,industry FROM sessions WHERE name LIKE ? ORDER BY timestamp DESC LIMIT ?",
                  (f"%{filter_name}%", limit))
    else:
        c.execute("SELECT id,timestamp,user,name,industry FROM sessions ORDER BY timestamp DESC LIMIT ?",
                  (limit,))
    rows = c.fetchall()
    conn.close()
    return rows

# -----------------------
# Groq LLM call with retries & JSON validation
# -----------------------
def call_groq_chat(prompt_messages, max_tokens=1200, temperature=0.2, max_retries=3):
    """
    prompt_messages: list of dicts like {"role": "system|user", "content": "..."}
    Tries several models in GROQ_MODELS_TO_TRY if model-specific errors occur.
    Returns the assistant content string.
    Raises RuntimeError on final failure with helpful message.
    """
    if not GROQ_API_KEY:
        raise EnvironmentError("GROQ_API_KEY not set. Add it to your environment or platform secrets as GROQ_API_KEY.")

    headers = {
        "Authorization": f"Bearer {GROQ_API_KEY}",
        "Content-Type": "application/json"
    }

    last_exception = None
    for model in GROQ_MODELS_TO_TRY:
        payload = {
            "model": model,
            "messages": prompt_messages,
            "max_tokens": max_tokens,
            "temperature": temperature
        }
        attempt = 0
        while attempt < max_retries:
            try:
                resp = requests.post(GROQ_ENDPOINT, headers=headers, json=payload, timeout=30)
                # Debug: if you need logs, uncomment:
                # print("DEBUG: POST", GROQ_ENDPOINT, "model:", model, "status:", resp.status_code)
                # print("DEBUG: response text:", resp.text[:1000])
                resp.raise_for_status()
                data = resp.json()
                # attempt to extract assistant content robustly
                if "choices" in data and len(data["choices"])>0:
                    choice = data["choices"][0]
                    # OpenAI-compatible: choice["message"]["content"]
                    if isinstance(choice.get("message"), dict) and "content" in choice["message"]:
                        content = choice["message"]["content"]
                        # content can be string or dict
                        if isinstance(content, str):
                            return content
                        elif isinstance(content, dict):
                            # common pattern with text key
                            if "text" in content:
                                return content["text"]
                            else:
                                return json.dumps(content)
                    # fallback to choice["text"]
                    if "text" in choice:
                        return choice["text"]
                # fallback top-level text
                if "text" in data:
                    return data["text"]
                # if nothing found, raise to attempt retry
                raise RuntimeError("LLM response did not contain expected 'choices' content.")
            except Exception as e:
                last_exception = e
                attempt += 1
                time.sleep(1 + attempt * 0.3)
        # try next model if available
    # all attempts failed
    raise RuntimeError(f"Groq call failed after trying models {GROQ_MODELS_TO_TRY}. Last error: {last_exception}")

# -----------------------
# PDF Exporter
# -----------------------
def export_pdf(activity_dict, filename="activity_pack.pdf"):
    pdf = FPDF()
    pdf.set_auto_page_break(auto=True, margin=12)
    pdf.add_page()
    pdf.set_font("Arial", size=14)
    pdf.cell(0, 8, f"{APP_TITLE} — {activity_dict.get('activity_name', 'Activity')}", ln=True)
    pdf.ln(4)
    pdf.set_font("Arial", size=11)
    pdf.multi_cell(0, 6, f"Objective: {activity_dict.get('objective','')}")
    pdf.ln(2)
    pdf.multi_cell(0, 6, f"Why this works:\n{activity_dict.get('why','')}")
    pdf.ln(2)
    pdf.multi_cell(0, 6, f"Scenario:\n{activity_dict.get('scenario','')}")
    pdf.ln(2)
    pdf.multi_cell(0, 6, "Character sketches:")
    for ch in activity_dict.get('character_sketches', []):
        pdf.multi_cell(0, 6, f"- {ch.get('role')}: {ch.get('description')}")
    pdf.ln(2)
    pdf.multi_cell(0, 6, "Trainer Notes:")
    for step in activity_dict.get('trainer_notes', []):
        pdf.multi_cell(0, 6, f"- {step}")
    pdf.ln(2)
    pdf.multi_cell(0, 6, "Debrief Questions:")
    for q in activity_dict.get('debrief_questions', []):
        pdf.multi_cell(0, 6, f"- {q}")
    pdf.ln(4)
    pdf.multi_cell(0, 6, "Success indicators:")
    for s in activity_dict.get('success_indicators', []):
        pdf.multi_cell(0, 6, f"- {s}")
    pdf.output(filename)
    return filename

# -----------------------
# UI helpers
# -----------------------
def card(title, subtitle):
    st.markdown(
        f"""<div style="border-radius:12px;padding:12px;margin:8px 0;border:1px solid #eee;background:#fff">
            <div style="font-size:15px;font-weight:600">{title}</div>
            <div style="color:#666;margin-top:6px">{subtitle}</div>
        </div>""", unsafe_allow_html=True)

def get_asset_file():
    # Try repository relative path first, then uploaded local path, then placeholder
    candidates = [
        DEFAULT_ASSET_REL_PATH,
        os.path.join(os.getcwd(), DEFAULT_ASSET_REL_PATH),
        UPLOADED_LOCAL_PATH
    ]
    for p in candidates:
        if p and os.path.exists(p):
            return p
    return None

# -----------------------
# Main
# -----------------------
def main():
    init_db()
    st.set_page_config(page_title=APP_TITLE, layout="centered")
    st.sidebar.title(APP_TITLE)
    username = st.sidebar.text_input("User (prototype)", value="demo_user")
    st.sidebar.markdown("**Note**: This is prototype history stored locally. Add GROQ_API_KEY to Secrets for LLM calls.")

    asset_file = get_asset_file()
    if asset_file:
        st.image(asset_file, width=220)
    else:
        st.warning("Logo not found. Add assets/logo.png to the repo or set the uploaded local path.")
        st.image("https://via.placeholder.com/440x220.png?text=L%26D+Studio+AI", width=220)

    st.title(APP_TITLE)
    st.caption("AI-Powered Trainer Companion — mobile-first prototype")

    st.markdown("### Idea Generator")
    with st.form("idea_form"):
        industry = st.selectbox("Industry (context)", ["General / Cross-industry", "Pharmaceutical", "Healthcare", "Insurance", "Technology", "Education", "Other"])
        category = st.selectbox("Activity Category", ["Ice Breaking", "Activities During Workshop", "Learning Solution", "Energizer"])
        message = st.text_area("Message You Want to Drive (one-line)", placeholder="e.g., Active listening, feedback culture", height=70)
        audience = st.multiselect("Audience Profile", ["First-line managers", "Senior managers", "Sales reps", "New joiners", "Cross-functional teams", "Leadership team"])
        batch_size = st.text_input("Batch Size (e.g., 15–20)", value="15-20")
        duration = st.selectbox("Activity Duration", ["5-10 mins", "10-20 mins", "20-30 mins", "30-60 mins"])
        delivery = st.radio("Delivery Format", ["In-person", "Virtual"])
        training_stage = st.selectbox("Training Stage", ["Ice-breaker", "Mid-session energizer", "Concept introduction", "Deep-dive practice", "Reflection/closing"])
        scenario_context = st.text_area("Scenario Context (brief)", placeholder="Describe setting, culture, common challenges", height=80)
        desired_outcome = st.text_area("Desired Outcome (what should participants do differently?)", height=60)
        include_dialogues = st.checkbox("Include sample dialogues in output (roleplay only)", value=False)
        submitted = st.form_submit_button("Generate Activity Pack")

    if submitted:
        # mandatory check
        required = {
            "Message": message,
            "Audience": audience,
            "Batch size": batch_size,
            "Duration": duration,
            "Delivery": delivery,
            "Training stage": training_stage,
            "Scenario context": scenario_context,
            "Desired outcome": desired_outcome
        }
        for k, v in required.items():
            if not v or (isinstance(v, list) and len(v) == 0):
                st.error(f"Please complete: {k}")
                st.stop()

        # build messages for Groq chat
        system_prompt = {
            "role": "system",
            "content": (
                "You are an L&D assistant that generates trainer-ready activity packs. "
                "Output only a JSON object (no extra commentary) with keys: activity_name, objective, why, scenario, "
                "character_sketches (list of objects with role+description), trainer_notes (list), props (list), "
                "debrief_questions (3-5), kirkpatrick (dict), success_indicators (list), export_text, dialogues (optional)."
            )
        }
        user_content = (
            f"Industry: {industry}\nCategory: {category}\nMessage: {message}\nAudience: {', '.join(audience)}\n"
            f"Batch size: {batch_size}\nDuration: {duration}\nDelivery: {delivery}\nStage: {training_stage}\n"
            f"Scenario context: {scenario_context}\nDesired outcome: {desired_outcome}\nInclude dialogues: {include_dialogues}"
        )
        user_message = {"role": "user", "content": user_content}

        with st.spinner("Calling LLM (Groq)…"):
            try:
                llm_text = call_groq_chat([system_prompt, user_message], max_tokens=2000)
            except Exception as e:
                st.error(f"LLM call failed: {e}")
                st.stop()

        # attempt to parse JSON from the response (works even if LLM wraps text)
        parsed = None
        try:
            parsed = json.loads(llm_text)
        except Exception:
            m = re.search(r'\{(?:.|\n)*\}', str(llm_text))
            if m:
                try:
                    parsed = json.loads(m.group(0))
                except Exception:
                    parsed = None

        if not parsed:
            st.error("Failed to parse JSON from LLM output. Raw output is shown for debugging.")
            st.code(llm_text)
            st.stop()

        # show results
        st.markdown("### Activity Pack — Primary Recommendation")
        st.markdown(f"**{parsed.get('activity_name')}**")
        st.markdown(f"*Objective:* {parsed.get('objective')}")
        st.markdown(f"*Why this works:* {parsed.get('why')}")

        with st.expander("Scenario & Character Sketches"):
            st.write(parsed.get('scenario'))
            for cs in parsed.get('character_sketches', []):
                st.write(f"**{cs.get('role')}** — {cs.get('description')}")
            if include_dialogues and parsed.get('dialogues'):
                st.write("**Sample dialogues:**")
                st.write(parsed.get('dialogues'))

        with st.expander("Trainer Notes & Run-sheet"):
            for i, n in enumerate(parsed.get('trainer_notes', []), 1):
                st.write(f"{i}. {n}")

        with st.expander("Debrief Questions"):
            for q in parsed.get('debrief_questions', []):
                st.write(f"- {q}")

        with st.expander("Kirkpatrick & Success indicators"):
            st.json(parsed.get('kirkpatrick', {}))
            st.write("Success indicators:")
            for s in parsed.get('success_indicators', []):
                st.write(f"- {s}")

        # Save and export
        sid = save_session(parsed.get('activity_name', 'unnamed'), industry,
                           {"form": {"category": category, "message": message}, "result": parsed},
                           user=username)
        st.success("Saved to History (local).")

        if st.button("Export as PDF"):
            fname = f"{APP_TITLE.replace(' ', '_')}_{sid}.pdf"
            path = export_pdf(parsed, filename=fname)
            st.success(f"Exported PDF: {path}")
            # update exported flag (simple)
            conn = sqlite3.connect(DB_PATH)
            c = conn.cursor()
            c.execute("UPDATE sessions SET exported=1 WHERE id=?", (sid,))
            conn.commit()
            conn.close()

        if st.button("Download .txt (full pack)"):
            st.download_button("Download Pack as TXT", json.dumps(parsed, indent=2), file_name=f"{sid}.txt")

    st.markdown("---")
    st.subheader("History")
    q = st.text_input("Filter by activity name", "")
    rows = fetch_history(filter_name=q)
    if rows:
        for r in rows:
            st.write(f"- {r[3]}  ·  {r[1][:19]}  ·  {r[4]}  ·  id:{r[0]}")
    else:
        st.write("No history yet.")

    st.caption("Prototype note: secrets (GROQ_API_KEY) should be stored in environment or platform secrets, not in code.")

if __name__ == "__main__":
    main()