app.py

import os
import re
import json
import time
from typing import List, Dict, Any, Tuple, Optional

import gradio as gr
from pydantic import BaseModel, Field
from huggingface_hub import InferenceClient

# Title and constants
APP_TITLE = "PeakMind OS — High-Performance Regulation Lab"
DEFAULT_MODEL = os.getenv("MODEL_ID", "google/functiongemma-270m-it")

# Crisis copy and patterns
CRISIS_COPY = (
    "I’m really sorry you’re feeling this way. I can’t help with self-harm content.\n\n"
    "If you might be in danger or feel like you may act on these thoughts, please seek immediate help:\n"
    "- US/Canada: Call or text 988 (Suicide & Crisis Lifeline)\n"
    "- If you are in immediate danger: call 911 (or your local emergency number)\n"
    "- If outside the US: find local crisis lines here: https://www.opencounseling.com/suicide-hotlines\n\n"
    "If you’d like, tell me where you are (country), and I’ll point you to the right resources."
)

SELF_HARM_PATTERNS = [
    r"\bkill myself\b", r"\bsuicide\b", r"\bend my life\b", r"\bself[- ]?harm\b",
    r"\bwant to die\b", r"\bno reason to live\b"
]

DISCLAIMER = (
    "**Important:** This tool is for coaching, skills practice, and performance routines — **not therapy**.\n\n"
    "- It does not diagnose or treat mental health conditions.\n"
    "- If you’re in crisis or may harm yourself/others, use emergency resources immediately.\n"
    "- Avoid entering identifying personal details.\n"
)

# Self harm detection
def looks_like_self_harm(text: str) -> bool:
    t = (text or "").lower()
    return any(re.search(p, t) for p in SELF_HARM_PATTERNS)

# JSON extraction
def robust_json_loads(text: str) -> Optional[Dict[str, Any]]:
    if not text:
        return None
    match = re.search(r"\{.*\}", text, re.DOTALL)
    if not match:
        return None
    candidate = match.group(0)
    try:
        return json.loads(candidate)
    except Exception:
        return None

# Action Plan schema
class ActionPlan(BaseModel):
    headline: str = Field(description="One-sentence focus for the next 30–90 minutes.")
    next_best_move: str = Field(description="The single most useful next action.")
    micro_steps: List[str] = Field(description="2–4 tiny steps to start immediately.")
    regulation: List[str] = Field(description="1–3 regulation steps (breath/body/attention).")
    if_then_rule: str = Field(description="Implementation intention to prevent spirals.")
    boundary_line: str = Field(description="Optional: a one-liner for workplace boundaries.")
    check_in: str = Field(description="When and how to re-check progress.")

def plan_prompt(stress: int, context: str, goal: str, minutes: int, style: str) -> str:
    return f"""
You are a high-performance coach focused on emotional regulation and execution.
You are NOT a therapist. Do not provide clinical or medical advice.

Write a concise, practical plan in VALID JSON only that matches this exact schema:
{ActionPlan.model_json_schema()}

User inputs:
- stress (0-10): {stress}
- context: {context}
- goal: {goal}
- time available (minutes): {minutes}
- style preference: {style}

Rules:
- Keep it grounded, non-judgmental, action-oriented.
- No therapy language. No diagnosis. No trauma processing.
- Include at least one regulation step that takes <= 90 seconds.
- micro_steps should be tiny and immediately doable.
Return JSON only, no backticks, no extra keys.
""".strip()

# Inference client
def get_client() -> Optional[InferenceClient]:
    token = os.getenv("HF_TOKEN")
    try:
        return InferenceClient(model=DEFAULT_MODEL, token=token) if token else InferenceClient(model=DEFAULT_MODEL)
    except Exception:
        return None

# Generate plan
def generate_plan(stress: int, context: str, goal: str, minutes: int, style: str) -> Tuple[str, str]:
    user_text = f"{context}\nGoal: {goal}"
    if looks_like_self_harm(user_text):
        return CRISIS_COPY, ""
    client = get_client()
    prompt = plan_prompt(stress, context, goal, minutes, style)
    # fallback deterministic plan if no client
    if client is None:
        plan = ActionPlan(
            headline="Stabilize → Choose one lever → Execute small chunk",
            next_best_move="Write the smallest next deliverable and start a timer.",
            micro_steps=[
                "Write down the deliverable in one sentence.",
                "Open the document or tool you need.",
                "Do the first 5 minutes imperfectly."
            ],
            regulation=[
                "Physiological sigh x2 (two short inhales, long exhale).",
                "Unclench jaw/shoulders; exhale longer than inhale for 60 seconds."
            ],
            if_then_rule="If I start spinning, then I do 60 seconds of slow exhale and take the tiniest next step.",
            boundary_line="I can do X by (time). If priorities change, what should drop?",
            check_in="In 20 minutes: rate stress again (0–10) and decide: continue, simplify, or pause."
        )
        return plan.model_dump_json(indent=2), "Using fallback (no inference client available)."
    debug: List[str] = []
    for attempt in range(2):
        try:
            out = client.text_generation(
                prompt,
                max_new_tokens=450,
                temperature=0.4,
                top_p=0.9,
                repetition_penalty=1.15,
            )
            data = robust_json_loads(out)
            if data:
                plan = ActionPlan(**data)
                return plan.model_dump_json(indent=2), f"Generated with model {DEFAULT_MODEL}."
            else:
                debug.append("JSON parse failed; retrying with repair instruction.")
                prompt = prompt + "\n\nREPAIR: Your last output was not valid JSON. Output ONLY valid JSON for the schema."
        except Exception as e:
            debug.append(f"Inference error: {e}")
            break
    return (
        "I hit an inference/formatting issue. Try again, or set a different MODEL_ID.\n\n"
        "Tip: use a smaller instruction model or ensure HF_TOKEN is set as a Space Secret.",
        "\n".join(debug)
    )

# Breathing timer helper
def breathing_timer(seconds: int, pattern: str) -> str:
    if seconds <= 0:
        return "Set a duration > 0."
    return (
        f"**Timer:** {seconds}s\n\n"
        f"**Pattern:** {pattern}\n\n"
        "Start now. Breathe gently—no straining.\n"
        "If you feel dizzy, return to normal breathing."
    )

# Reset text generator
def reset_text(rt: str, dur: int) -> str:
    if rt.startswith("Physiological"):
        pattern = "2 cycles: short inhale → top-up inhale → long exhale. Then slow breathing."
    elif rt.startswith("Box"):
        pattern = "Inhale 4s → Hold 4s → Exhale 4s → Hold 4s. Repeat."
    else:
        pattern = "Inhale gently 3–4s → Exhale 6–8s. Repeat."
    return breathing_timer(dur, pattern)

# Reframe generator
def simple_reframe(s: str, e: str, u: str, v: str) -> str:
    text = " ".join([s or "", e or "", u or "", v or ""])
    if looks_like_self_harm(text):
        return CRISIS_COPY
    return (
        f"**Reframe:** This is {e or 'a real feeling'}—not a command. "
        f"My job is to act from **{v or 'my values'}**, not from the urge to **{u or 'react'}**.\n\n"
        f"**One aligned action (10 minutes):** Identify the smallest professional next step and do it imperfectly.\n\n"
        f"**Boundary line (optional):** “To do this well, I need X. If we change priorities, what should drop?”"
    )

# After-action review
def aar(w: str, wi: str, c: str) -> str:
    text = " ".join([w or "", wi or "", c or ""])
    if looks_like_self_harm(text):
        return CRISIS_COPY
    lever = (c or "Pick one micro-skill to practice next time").strip()
    return (
        "### Your Loop\n"
        f"**Event:** {w or '—'}\n\n"
        f"**Strength:** {wi or '—'}\n\n"
        f"**Next-time lever:** {lever}\n\n"
        "**Practice rule:** make it *smaller than you think* and repeat it 3 times this week."
    )

# Log functions
def add_entry(entry: str, logs_json: str) -> Tuple[str, List[List[str]]]:
    try:
        logs = json.loads(logs_json) if logs_json else []
    except Exception:
        logs = []
    entry = entry.strip()
    if entry:
        logs.append(entry)
    display = [[e] for e in logs]
    return json.dumps(logs), display

def export_logs(logs_json: str) -> Tuple[str]:
    # Write logs JSON to a temporary file for download
    path = f"/tmp/peakmind_logs_{int(time.time())}.json"
    with open(path, "w") as f:
        f.write(logs_json or "[]")
    return path,

# Build UI
with gr.Blocks(title=APP_TITLE) as demo:
    gr.Markdown(f"# {APP_TITLE}\n\n{DISCLAIMER}")
    with gr.Row():
        with gr.Column(scale=1):
            gr.Markdown("## Quick Inputs")
            stress = gr.Slider(0, 10, value=5, step=1, label="Current stress (0–10)")
            minutes = gr.Slider(5, 180, value=30, step=5, label="Time available (minutes)")
            style = gr.Dropdown(
                ["Direct & no-BS", "Warm & encouraging", "Analytical & structured"],
                value="Analytical & structured",
                label="Coaching style"
            )
            model_note = gr.Markdown(
                f"**Model:** `{DEFAULT_MODEL}`  \n"
                "You can set `MODEL_ID` as a Space Variable and `HF_TOKEN` as a Space Secret."
            )
        with gr.Column(scale=2):
            with gr.Tab("Pressure → Clarity Plan"):
                context = gr.Textbox(lines=3, label="Context (keep non-identifying)",
                                     placeholder="Example: Big client call in 2 hours; I'm scattered and irritated.")
                goal = gr.Textbox(lines=2, label="Goal",
                                  placeholder="Example: Enter the call calm, clear, and decisive.")
                run_btn = gr.Button("Generate Plan")
                plan_out = gr.Code(label="Action Plan (JSON)", language="json")
                status = gr.Markdown()
                run_btn.click(
                    fn=generate_plan,
                    inputs=[stress, context, goal, minutes, style],
                    outputs=[plan_out, status]
                )
            with gr.Tab("60–120s Reset"):
                reset_type = gr.Radio(
                    ["Physiological sigh (fast reset)", "Box breathing (4-4-4-4)", "Slow exhale (downshift)"],
                    value="Physiological sigh (fast reset)",
                    label="Choose a reset"
                )
                duration = gr.Slider(30, 180, value=90, step=15, label="Duration (seconds)")
                reset_out = gr.Markdown()
                reset_btn = gr.Button("Start Reset")
                reset_btn.click(reset_text, inputs=[reset_type, duration], outputs=reset_out)
            with gr.Tab("Cognitive Reframe (Work Mode)"):
                sit = gr.Textbox(lines=2, label="Situation", placeholder="What happened?")
                emo = gr.Textbox(lines=1, label="Emotion", placeholder="Name it: anxious, angry, ashamed, etc.")
                urge = gr.Textbox(lines=1, label="Urge", placeholder="What do you feel like doing?")
                val = gr.Textbox(lines=1, label="Value", placeholder="What matters here? professionalism, courage, honesty…")
                refr_btn = gr.Button("Generate Reframe (no therapy)")
                refr_out = gr.Markdown()
                refr_btn.click(simple_reframe, inputs=[sit, emo, urge, val], outputs=refr_out)
            with gr.Tab("After-Action Review"):
                what = gr.Textbox(lines=2, label="What happened?")
                win = gr.Textbox(lines=2, label="What did you do well?")
                change = gr.Textbox(lines=2, label="What will you do differently next time?")
                aar_btn = gr.Button("Summarize into 1 weekly lever")
                aar_out = gr.Markdown()
                aar_btn.click(aar, inputs=[what, win, change], outputs=aar_out)
            with gr.Tab("Private Log + Export"):
                log_state = gr.State(value="[]")
                new_entry = gr.Textbox(lines=2, label="New entry", placeholder="Type your journal entry here.")
                add_btn = gr.Button("Add to log")
                log_list = gr.Dataframe(headers=["Entries"], datatype=["str"], interactive=False, label="In-session log")
                logs_json_display = gr.Textbox(label="Copyable JSON", interactive=False)
                export_btn = gr.Button("Export log as JSON file")
                download_file = gr.File(label="Download logs", interactive=False)
                # Add entry callback
                add_btn.click(
                    fn=add_entry,
                    inputs=[new_entry, log_state],
                    outputs=[log_state, log_list]
                )
                # Update JSON display
                add_btn.click(
                    lambda logs_json: logs_json,
                    inputs=log_state,
                    outputs=logs_json_display
                )
                # Export callback
                export_btn.click(
                    fn=export_logs,
                    inputs=[log_state],
                    outputs=[download_file]
                )
    gr.Markdown(
        "---\n"
        "### About\n"
        "Built as a coaching demo focused on performance routines and emotional regulation.\n\n"
        "**Privacy:** This Space is designed to avoid data retention. Don’t enter identifying info."
    )

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