import re
import gradio as gr
from datasets import load_dataset
from huggingface_hub import InferenceClient
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_community.vectorstores import Chroma

# 1. Connect to Hugging Face Inference API (free hosted models)
# You MUST add your Hugging Face token in the Space "Secrets" panel as HF_TOKEN
client = InferenceClient("microsoft/phi-2", token=None)  # if hosted on your HF account, no token needed

def generate_text(prompt, max_tokens=150):
    """Call Hugging Face Inference API to generate text."""
    response = client.text_generation(prompt, max_new_tokens=max_tokens)
    return response

# 2. Load dataset for RAG
data = load_dataset("Amod/mental_health_counseling_conversations")
docs = [f"Q: {entry['Context']}\nA: {entry['Response']}" for entry in data['train']]
embedding_model = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
vectordb = Chroma.from_texts(texts=docs, embedding=embedding_model)

# 3. Traits dictionary
traits = {"Secure": 0, "Anxious": 0, "Avoidant": 0}
MAX_SCENES = 5

# Narrative instructions for LLM
narrator_instructions = (
    "You are a supportive narrative AI creating an interactive story for the player.\n"
    "Write in second person (\"you\" as Alex) and describe feelings and events vividly.\n"
    "The story should feel emotionally deep and realistic, suitable for an adult audience (age 25-35).\n"
    "Focus on how attachment styles influence Alex's feelings and decisions in a romantic relationship.\n"
    "In each scene, after narrating the situation, offer the player choices labeled A), B), (and possibly C)).\n"
    "Each choice should represent a different attachment response (secure, avoidant, anxious).\n"
    "Do NOT reveal the outcome of the choices yet; just present the options.\n"
)

# --- STORY FUNCTIONS ---
def parse_choices(text):
    """Split LLM output into narrative and choices."""
    choice_start = text.find("A)")
    if choice_start == -1:
        return text, []
    narrative = text[:choice_start].strip()
    choice_lines = text[choice_start:].splitlines()
    choices = [line.strip() for line in choice_lines if re.match(r"^[A-C]\)", line.strip())]
    return narrative, choices

def generate_initial_scene():
    intro = (
        "Alex is a 26-year-old graduate student who recently moved to the UK. "
        "They are in a romantic relationship, but lately Alex feels insecure. "
        "One morning, Alex wakes up with a knot in their stomach. "
        "Their partner didn’t reply to a message from last night. "
        "Thoughts race through Alex’s mind..."
    )
    text = generate_text(narrator_instructions + intro + "\nWhat does Alex do?")
    return parse_choices(text)

def next_step(selected_choice, story_so_far, scene_index, traits_dict):
    if not selected_choice:
        return story_so_far, [], scene_index, traits_dict

    choice_label = selected_choice[0]
    choice_text = selected_choice[3:].strip()
    scene_index += 1

    # Update traits
    text_lower = choice_text.lower()
    if any(word in text_lower for word in ["talk", "share", "open", "honest", "calm"]):
        traits_dict["Secure"] += 1
    if any(word in text_lower for word in ["ignore", "avoid", "distance", "later", "alone"]):
        traits_dict["Avoidant"] += 1
    if any(word in text_lower for word in ["worry", "jealous", "clingy", "panic", "angry", "upset"]):
        traits_dict["Anxious"] += 1

    # Retrieve advice
    results = vectordb.similarity_search(choice_text + " relationship anxiety", k=1)
    advice = results[0].page_content.split("A:")[1].strip() if results else ""

    # Prompt for next scene
    last_line = story_so_far.splitlines()[-1] if story_so_far else ""
    prompt = narrator_instructions
    prompt += f"Previously: {last_line}\n"
    prompt += f"The player chose option {choice_label}: {choice_text}.\n"
    if advice:
        prompt += f"[Helpful thought: {advice}]\n"
    prompt += "Continue the story and provide new choices.\n"

    text = generate_text(prompt)
    return *parse_choices(text), scene_index, traits_dict

def summarize_story(story, traits_dict):
    prompt = (
        f"Alex's story is ending. Traits: Secure={traits_dict['Secure']}, "
        f"Anxious={traits_dict['Anxious']}, Avoidant={traits_dict['Avoidant']}.\n"
        "Write a reflective ending about what Alex learned about relationships."
    )
    return generate_text(prompt, max_tokens=120)

# --- GRADIO UI ---
with gr.Blocks(css=".gradio-container {background-color:#f5f7fa;} #story_box {padding:15px; background:#fff; border-radius:10px;}") as demo:
    story_state = gr.State("")
    scene_state = gr.State(0)
    traits_state = gr.State(traits.copy())

    gr.HTML("<h1>💙 Attachment Style Interactive Story</h1><p>Explore how attachment styles shape relationships.</p>")

    # Background music
    gr.HTML("""
    <audio src='file=bg_music.mp3' autoplay loop hidden></audio>
    """)

    story_box = gr.Markdown("*(Loading story...)*", elem_id="story_box")
    choices_radio = gr.Radio(label="Choose:", choices=[], visible=False)
    next_btn = gr.Button("Next")

    def start_game():
        narrative, choices = generate_initial_scene()
        return narrative, choices, 1, {"Secure":0,"Anxious":0,"Avoidant":0}

    def play_step(choice, story, scene, traits_dict):
        if scene >= MAX_SCENES:
            return story + "\n\n" + summarize_story(story, traits_dict), [], scene, traits_dict
        narrative, choices, new_scene, traits_dict = next_step(choice, story, scene, traits_dict)
        story += "\n\n" + narrative
        if new_scene >= MAX_SCENES:
            story += "\n\n" + summarize_story(story, traits_dict)
            return story, [], new_scene, traits_dict
        return story, choices, new_scene, traits_dict

    demo.load(start_game, inputs=None, outputs=[story_box, choices_radio, scene_state, traits_state])
    next_btn.click(play_step, inputs=[choices_radio, story_state, scene_state, traits_state],
                   outputs=[story_box, choices_radio, scene_state, traits_state])

demo.launch()