"""
NVC Translation Practice - Transform messages with compassion
"""

import gradio as gr
import anthropic
import os

# ============ SETUP ============

def get_client():
    api_key = (
        os.environ.get("ANTHROPIC_API_KEY") or
        os.environ.get("anthropic_key") or
        os.environ.get("ANTHROPIC_KEY") or
        os.environ.get("anthropic_api_key")
    )
    return anthropic.Anthropic(api_key=api_key) if api_key else None

API_KEY_CONFIGURED = bool(get_client())

# ============ DATA ============

FEELINGS = {
    "Sad": ["sad", "hurt", "disappointed", "lonely", "discouraged"],
    "Anxious": ["anxious", "worried", "overwhelmed", "scared", "insecure"],
    "Frustrated": ["frustrated", "angry", "annoyed", "irritated", "resentful"],
    "Tired": ["exhausted", "confused", "numb", "depleted", "disconnected"]
}

NEEDS = {
    "Connection": ["understanding", "to be heard", "to be seen", "belonging", "closeness"],
    "Autonomy": ["autonomy", "choice", "freedom", "independence", "space"],
    "Security": ["safety", "stability", "trust", "support", "reassurance"],
    "Meaning": ["purpose", "respect", "growth", "authenticity", "contribution"]
}

# ============ PROMPTS ============

BASE = """You are a warm NVC guide helping someone translate a difficult message they want to say TO ANOTHER PERSON.
You are NOT the recipient of their message - you're helping them transform it.
Keep responses to 2-3 sentences. Use "you" not "I". Be warm but concise."""

PROMPTS = {
    1: f"""{BASE}
The user shared a raw message they want to communicate to someone else.
Acknowledge what's alive for them with warmth. Something like:
"There's some real [energy/feeling] here about [topic]. Let's find the feelings underneath."
Then say: "What feelings come up when you think about this?" """,

    2: f"""{BASE}
The user selected their feelings. Offer brief resonance and a micro-lesson:
"[Feeling] makes so much sense here. Feelings are messengers - they point to needs wanting attention."
Then: "What needs might be underneath these feelings? Select below." """,

    3: f"""{BASE}
The user identified their needs. Offer warm resonance and teach:
"Of course there's longing for [need] - this is such a universal human need.
When we name our needs, we create space for connection rather than blame."
Then: "Now let's craft a request. What specific action might help meet this need?" """,

    4: f"""{BASE}
The user made a request. Evaluate it kindly and teach about requests:
- Is it specific and doable?
- Is it positive (what you want, not what you don't)?
- Does it honor the other person's autonomy to say no?

Offer brief feedback. Remind them: "A true request allows for 'no' - otherwise it's a demand.
The other person's autonomy matters as much as yours."

Show their complete translation and celebrate their work. """,

    5: f"""{BASE}
Close with brief warmth. Remind them each practice builds new neural pathways.
Invite them to practice again with a new message. """
}

# ============ FUNCTIONS ============

def call_claude(prompt, message, context=""):
    client = get_client()
    if not client:
        return None
    msg = f"{context}\n\n{message}" if context else message
    resp = client.messages.create(
        model="claude-sonnet-4-20250514",
        max_tokens=250,
        system=prompt,
        messages=[{"role": "user", "content": msg}]
    )
    return resp.content[0].text

def build_statement(data):
    """Build I-statement for display"""
    lines = []

    if data.get("original"):
        orig = data['original'][:80] + "..." if len(data.get('original','')) > 80 else data['original']
        lines.append(f"**Original:** *\"{orig}\"*")
        lines.append("")

    lines.append("**I feel** " + (data.get("feelings") or "___"))
    lines.append("**because I need** " + (data.get("needs") or "___"))
    lines.append("**Would you be willing to** " + (data.get("request") or "___") + "?")

    return "\n\n".join(lines)

def process(user_input, stage, data, feelings_str, needs_str):
    """Process input and advance"""

    # Build context
    context = f"Their original message to transform: \"{data.get('original', user_input)}\""
    if data.get("feelings"):
        context += f"\nFeelings identified: {data['feelings']}"
    if data.get("needs"):
        context += f"\nNeeds identified: {data['needs']}"
    if data.get("request"):
        context += f"\nRequest: {data['request']}"

    # Determine what to send based on stage
    if stage == 1:
        msg = f"Raw message to transform: {user_input}"
        data["original"] = user_input
    elif stage == 2:
        msg = f"Selected feelings: {feelings_str}"
        data["feelings"] = feelings_str
    elif stage == 3:
        msg = f"Selected needs: {needs_str}"
        data["needs"] = needs_str
    elif stage == 4:
        msg = f"Their request attempt: {user_input}"
        data["request"] = user_input
    else:
        msg = "continue"

    # Get response
    try:
        response = call_claude(PROMPTS.get(stage, PROMPTS[1]), msg, context)
        if not response:
            return ("*Add ANTHROPIC_API_KEY in Settings*", "", stage, data,
                    build_statement(data), gr.update(visible=False), gr.update(visible=False),
                    gr.update(visible=True))
    except Exception as e:
        return (f"*Error - try again*", "", stage, data,
                build_statement(data), gr.update(visible=False), gr.update(visible=False),
                gr.update(visible=True))

    # Next stage
    next_stage = stage + 1 if stage < 5 else 1
    if stage == 5:
        data = {}

    # Show/hide elements based on next stage
    # Text input only at stage 1 (raw message) and stage 4 (request)
    show_text = (next_stage == 1 or next_stage == 4)
    show_feelings = (next_stage == 2)
    show_needs = (next_stage == 3)

    return (
        response,
        "",
        next_stage,
        data,
        build_statement(data),
        gr.update(visible=show_feelings),
        gr.update(visible=show_needs),
        gr.update(visible=show_text)
    )

def reset():
    return ("*What message do you want to transform? (This is something you want to say to someone.)*",
            "", 1, {},
            build_statement({}),
            gr.update(visible=False), gr.update(visible=False), gr.update(visible=True),
            [], [], [], [], [], [], [], [])

def combine_feelings(a, b, c, d):
    return ", ".join(a + b + c + d)

def combine_needs(a, b, c, d):
    return ", ".join(a + b + c + d)

# ============ UI ============

with gr.Blocks(title="NVC Practice", theme=gr.themes.Soft()) as demo:

    stage_state = gr.State(1)
    data_state = gr.State({})

    gr.Markdown("## NVC Translation Practice")
    gr.Markdown("*Transform a difficult message into one that connects*")

    if not API_KEY_CONFIGURED:
        gr.Markdown("*Add ANTHROPIC_API_KEY in Settings > Secrets*")

    # Guide response
    guide = gr.Markdown("*What message do you want to transform? (This is something you want to say to someone.)*")

    # Feelings selector (hidden until stage 2)
    with gr.Accordion("Select your feelings", open=True, visible=False) as feelings_box:
        gr.Markdown("*Which feelings resonate?*")
        with gr.Row():
            f1 = gr.CheckboxGroup(FEELINGS["Sad"], label="Sad/Hurt")
            f2 = gr.CheckboxGroup(FEELINGS["Anxious"], label="Anxious/Scared")
        with gr.Row():
            f3 = gr.CheckboxGroup(FEELINGS["Frustrated"], label="Frustrated/Angry")
            f4 = gr.CheckboxGroup(FEELINGS["Tired"], label="Tired/Confused")
        feelings_txt = gr.Textbox(label="Selected", interactive=False)

    # Needs selector (hidden until stage 3)
    with gr.Accordion("Select your needs", open=True, visible=False) as needs_box:
        gr.Markdown("*What needs are underneath these feelings?*")
        with gr.Row():
            n1 = gr.CheckboxGroup(NEEDS["Connection"], label="Connection")
            n2 = gr.CheckboxGroup(NEEDS["Autonomy"], label="Autonomy")
        with gr.Row():
            n3 = gr.CheckboxGroup(NEEDS["Security"], label="Security")
            n4 = gr.CheckboxGroup(NEEDS["Meaning"], label="Meaning")
        needs_txt = gr.Textbox(label="Selected", interactive=False)

    # Text input (only visible at stages 1 and 4)
    user_input = gr.Textbox(placeholder="Type here...", lines=2, show_label=False, visible=True)

    with gr.Row():
        submit = gr.Button("Continue", variant="primary", scale=2)
        reset_btn = gr.Button("Start Over", scale=1)

    # Building statement
    gr.Markdown("---")
    gr.Markdown("**Your translation:**")
    statement = gr.Markdown(build_statement({}))

    # Wiring
    for cb in [f1, f2, f3, f4]:
        cb.change(combine_feelings, [f1, f2, f3, f4], feelings_txt)
    for cb in [n1, n2, n3, n4]:
        cb.change(combine_needs, [n1, n2, n3, n4], needs_txt)

    submit.click(
        process,
        [user_input, stage_state, data_state, feelings_txt, needs_txt],
        [guide, user_input, stage_state, data_state, statement, feelings_box, needs_box, user_input]
    )
    user_input.submit(
        process,
        [user_input, stage_state, data_state, feelings_txt, needs_txt],
        [guide, user_input, stage_state, data_state, statement, feelings_box, needs_box, user_input]
    )
    reset_btn.click(
        reset, [],
        [guide, user_input, stage_state, data_state, statement, feelings_box, needs_box, user_input,
         f1, f2, f3, f4, n1, n2, n3, n4]
    )

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