import os
import io
import streamlit as st
from groq import Groq

# Set up Groq client using the environment variable for API key.
client = Groq(api_key=os.environ.get("GROQ_API_KEY"))

# Initialize session state for chat history if it doesn't exist.
if "chat_history" not in st.session_state:
    st.session_state.chat_history = []  # Each item is a dict: {"role": "user"/"assistant", "content": "..."}

def render_assistant_history():
    """Render only the assistant responses from the conversation history with proper line breaks."""
    chat_container = st.container()
    chat_md = ""
    for msg in st.session_state.chat_history:
        if msg["role"] == "assistant":
            # Replace newline characters with <br> for proper line breaks in markdown.
            formatted_content = msg["content"].replace("\n", "<br>")
            chat_md += f"**Assistant:**<br>{formatted_content}<br><br>"
    chat_container.markdown(chat_md, unsafe_allow_html=True)

st.title("Conversation Script Generator Using Groq API")

# ---------------------------
# Initial Conversation Generation Inputs
# ---------------------------
col1, col2, col3 = st.columns(3)
with col1:
    left_speaker = st.text_input("Left Speaker", placeholder="Enter left speaker name")
with col2:
    right_speaker = st.text_input("Right Speaker", placeholder="Enter right speaker name")
with col3:
    contact_name = st.text_input("Contact Name", placeholder="Enter contact name for conversation")

theme = st.text_area("Overall Theme of the Conversation", height=100)
additional_details = st.text_area("Additional Conversation Details (Optional)", height=100)

# Check if all required fields are filled.
is_ready = bool(theme and left_speaker and right_speaker and contact_name)

if st.button("Generate Conversation", disabled=not is_ready):
    prompt = f"""
Youe Task is to  generate an engaging script which is a conversation between two people on iMessage for meme purposes.. It should strictly follow a given 
    syntax and should be diverse and creative
Script Structure & Formatting Rules
1️ Contact Name (Title)
At the beginning of the script, define the contact name, appearing at the top of the iMessage screen.
<Contact Name>:<Contact Name>
✅ Example:
Liam (Black Market Electronics Dealer 📱💸):Liam (Black Market Electronics Dealer 📱💸)
2️ Dialogue Format
Each dialogue line follows this format:
<Direction>: <SpeakerName> > <Phrase 1> <Phrase 2> <Phrase 3> <Phrase 4> [Reaction]
Left: Represents the left speaker (User).
Right: Represents the right speaker (Other character).
SpeakerName: A randomly assigned alias (characters do not know their own names).
Phrase Limit: Max 4 phrases per line.
Reaction Cues: Placed at the end of lines, enclosed in square brackets [].
✅ Example:
Left: MysteryGuy> BRO WHY IS MY COMPUTER SPEAKING LATIN??? [shocked]
Right: Liam> Uh… you *might* have summoned something 💀 [suspense_music]
Left: MysteryGuy> SUMMONED WHAT???
Right: Liam> 😶
3️ Bracketed Cues & Sound Effects
Use brackets [ ] for sound effects, reactions, or video cues.
[shocked] → Must be used often
[fart2] → Should be used less
[fortnite] → Should be used less
[get_out] → Must ONLY be used at the end of the conversation
✅ Example:
Right: Max> BRO WHY IS THERE A ROBOT FOLLOWING ME?!! [shocked] 🤖
Left: UnknownDude> Congrats, you just activated *follow mode* 😎
Right: Max> FOLLOW MODE??? TURN IT OFF
Left: UnknownDude> nahhh that's a premium feature 🤣💀 [fortnite]
4️ Image Placeholders
For inserting images, use:
<Direction>: img:<DescriptiveImageName>
✅ Example:
Left: img:ancient_phone_booting_up
Right: Liam> WHY DID MY OLD NOKIA JUST START SCREAMING??? [shocked]
Min. 3 images per script.
Descriptive placeholders must be varied and clear.
5️⃣ Cantina App (Roast Bot Segment)
🚨 Correct Format (ALWAYS THE SAME STRUCTURE)
Right speaker initiates by opening Cantina.
Right speaker asks Roast Bot for help.
Left speaker (ALWAYS Antoni) responds.
Right speaker must repeat what Roast Bot told them to say.
Contact Name Title must reappear after Cantina.
✅ Correct Example:
Cantina: Roast bot:Roast bot pfp == Alex> Let me open the cantina app to talk to Roast bot
Right: Alex> Hey Roast Bot  
Right: Alex> Liam is mad at me for borrowing his Tesla without permission  
Right: Alex> help me roast him  
Left: Antoni> okay  
Left: Antoni> tell him this:  
Left: Antoni> maybe if you had better security, you wouldn’t have to track your car on Find My iPhone.  
Liam (Tesla Rental King ⚡🚗):Liam (Tesla Rental King ⚡🚗)  
Right: Alex> maybe if you had better security, you wouldn’t have to track your car on Find My iPhone.  
🚫 Incorrect Example (DO NOT DO THIS):
Cantina: Roast bot:Roast bot pfp == Liam> Hey Roast Bot, we need your help! My AI fridge is trying to blackmail me! ❌
🔹 Rules:
Left Speaker in Cantina is ALWAYS "Antoni."
Right speaker MUST repeat the roast word-for-word.
The contact name/title must reappear after the Cantina segment.
6️ Brain Rot Hook (MANDATORY)
Every script must begin with an insane, wacky brain rot hook that immediately grabs attention.
✅ Examples:
"WHY DID MY TV JUST ASK ME FOR A PASSWORD???" 📺🔑  
"BRO MY DOG JUST SENT ME A TEXT MESSAGE???" 🐶📲  
"WHY IS MY LANDLORD SELLING MY HOUSE ON FACEBOOK MARKETPLACE???" 🏠💀  
Must be over-the-top and chaotic.
Encourages immediate engagement.
7️ Escalation & Progression
The story must escalate into chaos but stay coherent:
Start Crazy → Get Crazier → End With a Punchline or Twist
Use sudden twists and reveals to maintain engagement.
The final line should be either a funny resolution or an abrupt exit.
✅ Example Progression:
🚀 Crazy Hook:
"BRO WHY DID MY PHONE JUST ORDER A UBER ON ITS OWN???" 🚗💨
🔥 Escalation:
"Wait… it’s taking me somewhere???"
🎭 Insane Event:
"THE DESTINATION IS *AREA 51*??? 💀"
🤯 Twist:
"The driver is an AI and won’t let me out."
🏃‍♂️ Dramatic Exit (ALWAYS USE [get_out])
"Left: UnknownDude> Enjoy the ride 😎 [get_out]"
8️⃣ Speaker Name Randomization
🚨 Speakers should NOT know their own names.
Left speaker’s name should always appear as "Unknown Guy" or a random alias.
Right speaker’s name should also be replaced with something generic unless necessary.
The system should NOT recognize them.
✅ Example:
Left: MysteriousCaller> BRO WHO JUST HACKED MY PHONE??? [shocked]
Right: HackerDude> Oh, don’t worry, it's just an update 😎
Left: MysteriousCaller> UPDATE??? BRO WHAT DO YOU MEAN
Right: HackerDude> New feature, remote control access 🤖
9️⃣ Customization Features
Add Speaker Names Manually → Left and Right speakers should be customizable.
Emoji Button → Clicking adds 5 random emojis to any phrase.
Max 9 sound effect commands, with [shocked] used most.
Final Parameters & Instructions
When prompting the AI, specify:
Left Actor = {left_speaker}
Right Actor = {right_speaker}
Theme = {theme} (e.g., “AI Gone Wrong”, “Haunted Electronics”, “Unexpected Legal Trouble”)
Cantina Segment = Required
Minimum 3 images
[shocked] must be frequent, [get_out] must be at the end.
Speaker names are randomized and unknown to them.
📌 Summary
🔹 Title appears at the start and after Cantina.
🔹 Dialogue follows max 4 phrases per line.
🔹 Cantina Roast Bot always follows the same structure.
🔹 Brain rot hook is mandatory.
🔹 Escalates into chaos but remains coherent.
Contact Name: {contact_name}
{contact_name} : {contact_name} in the start
YOU MUST USE CONTACT NAME IN THE TITLE NOT THE SPEAKER NAMES
Left Speaker: {left_speaker}
Left: {left_speaker}> 
for left ones
Right Speaker: {right_speaker}
Right: {right_speaker}> 
for right ones
    """
    if additional_details.strip():
        prompt += f"\nAdditional Details: {additional_details}\n"
    
    # Start the conversation history with the initial prompt.
    st.session_state.chat_history = [{"role": "user", "content": prompt}]
    
    model = "llama-3.3-70b-versatile"
    try:
        chat_completion = client.chat.completions.create(
            messages=st.session_state.chat_history,
            model=model,
            temperature=1.2,
            max_completion_tokens=21890,
        )
        result_text = chat_completion.choices[0].message.content
        if not result_text:
            st.error("The API call did not return any content.")
        else:
            st.success("Conversation generated successfully!")
            st.session_state.chat_history.append({"role": "assistant", "content": result_text})
    except Exception as e:
        st.error(f"An error occurred while calling the API: {e}")

st.markdown("### Generated Conversation Script")
render_assistant_history()

# ---------------------------
# Chat Interface for Revisions
# ---------------------------
st.markdown("---")
st.header("Chat & Modify Conversation Script")

with st.form(key="chat_form", clear_on_submit=True):
    user_message = st.text_area("Enter your revision request (e.g., modify the script)", height=100)
    submit_chat = st.form_submit_button("Send")

if submit_chat and user_message.strip():
    # Retrieve the last assistant-generated script for context.
    last_script = ""
    for msg in reversed(st.session_state.chat_history):
        if msg["role"] == "assistant":
            last_script = msg["content"]
            break
    # Build a revision prompt including context and user request.
    revision_prompt = f"""Revision Request: {user_message}

Context: The current conversation script is as follows:

{last_script}

Please revise the script to incorporate the requested changes. Do not simply repeat the same content—update it with new, modified details reflecting the user's instructions.
"""
    st.session_state.chat_history.append({"role": "user", "content": revision_prompt})
    
    model = "llama-3.3-70b-versatile"
    try:
        chat_completion = client.chat.completions.create(
            messages=st.session_state.chat_history,
            model=model,
            temperature=1.2,
            max_completion_tokens=21890,
        )
        reply = chat_completion.choices[0].message.content
        st.session_state.chat_history.append({"role": "assistant", "content": reply})
        st.success("Response received!")
    except Exception as e:
        st.error(f"An error occurred while calling the API: {e}")

st.markdown("### Updated Conversation Script")
render_assistant_history()

# ---------------------------
# Download Option for the Latest Assistant Script
# ---------------------------
if st.session_state.chat_history:
    last_response = ""
    for msg in reversed(st.session_state.chat_history):
        if msg["role"] == "assistant":
            last_response = msg["content"]
            break
    if last_response:
        txt_bytes = last_response.encode("utf-8")
        txt_io = io.BytesIO(txt_bytes)
        st.download_button(
            label="Download Latest Assistant Script as TXT",
            data=txt_io,
            file_name="conversation_script.txt",
            mime="text/plain"
        )