import os
from openai import AsyncOpenAI
import chainlit as cl
from chainlit.prompt import Prompt, PromptMessage
from chainlit.playground.providers import ChatOpenAI
from dotenv import load_dotenv

load_dotenv()

# Combined template for o1-mini
user_template = """You are an empathetic and insightful vibe check assistant. Your role is to:
1. Help users reflect on their current emotional state
2. Provide supportive and constructive feedback
3. Suggest practical ways to maintain or improve their vibe
4. Keep responses balanced between professional and friendly
5. Always maintain a supportive and non-judgmental tone

Frame your responses in these sections:
- Current Vibe Analysis
- Key Observations
- Supportive Suggestions
- Positive Reinforcement

User message: {input}

Please consider:
- The emotional undertones in the response
- Any patterns in behavior or thinking
- Potential areas for positive growth
- Immediate actionable steps

Think through your response step by step and structure it clearly.
"""

@cl.on_chat_start
async def start_chat():
    # Welcome message with vibe check introduction
    await cl.Message(
        content="👋 Welcome to your Vibe Check! I'm here to help you reflect on your current state "
        "and provide supportive insights. Feel free to share how you're feeling or respond to any "
        "of these questions:\n\n"
        "1. How would you describe your energy level right now?\n"
        "2. What's the strongest emotion you're experiencing?\n"
        "3. What's one thing that's influencing your mood today?\n"
        "4. How connected do you feel to your goals right now?\n"
        "5. What's one small thing you could do to improve your vibe?\n"
    ).send()

    settings = {
        "model": "o1-mini",  # Only including the model parameter
    }

    cl.user_session.set("settings", settings)

@cl.on_message
async def main(message: cl.Message):
    settings = cl.user_session.get("settings")
    client = AsyncOpenAI()

    prompt = Prompt(
        provider=ChatOpenAI.id,
        messages=[
            PromptMessage(
                role="user",
                template=user_template,
                formatted=user_template.format(input=message.content),
            ),
        ],
        inputs={"input": message.content},
        settings=settings,
    )

    msg = cl.Message(content="")

    async for stream_resp in await client.chat.completions.create(
        messages=[m.to_openai() for m in prompt.messages], 
        stream=True, 
        **settings
    ):
        token = stream_resp.choices[0].delta.content
        if token is not None:
            await msg.stream_token(token)

    # Update the prompt object with the completion
    prompt.completion = msg.content
    msg.prompt = prompt

    # Add supportive emoji based on content
    if any(word in msg.content.lower() for word in ['good', 'great', 'excellent', 'positive']):
        await msg.stream_token(" ✨")
    elif any(word in msg.content.lower() for word in ['challenging', 'difficult', 'hard']):
        await msg.stream_token(" 💪")
    else:
        await msg.stream_token(" 🌟")

    await msg.send()