import os
import streamlit as st
import torch
from langchain.chains import LLMChain
from langchain.prompts import ChatPromptTemplate
from langchain_huggingface import HuggingFaceEndpoint

def create_conversation_prompt(name1: str, name2: str, persona_style: str):
    """
    Create a prompt that instructs the model to produce exactly 15 messages
    of conversation, alternating between name1 and name2, starting with name1.
    """
    prompt_template_str = f"""
    You are to simulate a conversation of exactly 15 messages total between two people: {name1} and {name2}.
    The conversation should reflect the style: {persona_style}.
    {name1} speaks first (message 1), then {name2} (message 2), then {name1} (message 3), and so forth, 
    alternating until all 15 messages are complete (the 15th message is by {name1}).

    Rules:
    - Output only the 15 messages of the conversation, nothing else.
    - Each message must be in the format:
      {name1}: <one or two short sentences>
      {name2}: <one or two short sentences>
    - Do not add any headings, such as '# Conversation' or '---'.
    - Do not repeat the conversation.
    - Do not produce anything after the 15th message.
    - Use everyday language, can ask questions, show opinions.
    - Use emojis sparingly, if at all (no more than 1-2 in the entire conversation).
    - No repetition of the same message. Each message is unique and responds logically to the previous one.
    - The conversation ends exactly after the 15th message by {name1} with no extra lines or content.

    Now, produce all 15 messages strictly following the above instructions:
    """
    return ChatPromptTemplate.from_template(prompt_template_str)

def create_summary_prompt(name1: str, name2: str, conversation: str):
    """Create a prompt specifically for generating a title and summary of the conversation."""
    summary_prompt_str = f"""
    The following is a completed conversation between {name1} and {name2}:

    {conversation}

    Please provide:
    Title: <A short descriptive title of the conversation>
    Summary: <A few short sentences highlighting the main points, tone, and conclusion>

    Do not continue the conversation, do not repeat it, just provide a title and summary.
    """
    return ChatPromptTemplate.from_template(summary_prompt_str)

def main():
    st.title("LLM Conversation Simulation")

    model_names = [
        "meta-llama/Llama-3.3-70B-Instruct",
        "mistralai/Mistral-7B-v0.1",
        "lmsys/vicuna-13b-v1.5",
        "tiiuae/falcon-180B",
        "EleutherAI/gpt-neox-20b",
        "dice-research/lola_v1"
    ]
    selected_model = st.selectbox("Select a model:", model_names)

    name1 = st.text_input("Enter the first user's name:", value="Alice")
    name2 = st.text_input("Enter the second user's name:", value="Bob")
    persona_style = st.text_area("Enter the persona style characteristics:", 
                                 value="friendly, curious, and a bit sarcastic")

    if st.button("Start Conversation Simulation"):
        st.write("**Loading model...**")
        print("Loading model...")

        with st.spinner("Starting simulation..."):
            endpoint_url = f"https://api-inference.huggingface.co/models/{selected_model}"

            try:
                llm = HuggingFaceEndpoint(
                    endpoint_url=endpoint_url,
                    huggingfacehub_api_token=os.environ.get("HUGGINGFACEHUB_API_TOKEN"),
                    task="text-generation",
                    temperature=0.7,
                    max_new_tokens=512
                )
                st.write("**Model loaded successfully!**")
                print("Model loaded successfully!")
            except Exception as e:
                st.error(f"Error initializing HuggingFaceEndpoint: {e}")
                print(f"Error initializing HuggingFaceEndpoint: {e}")
                return

            conversation_prompt = create_conversation_prompt(name1, name2, persona_style)
            conversation_chain = LLMChain(llm=llm, prompt=conversation_prompt)

            st.write("**Generating the full 15-message conversation...**")
            print("Generating the full 15-message conversation...")

            try:
                # Generate all 15 messages in one go
                conversation = conversation_chain.run(chat_history="", input="Produce the full conversation now.")
                conversation = conversation.strip()

                st.subheader("Final Conversation:")
                st.text(conversation)
                print("Conversation Generation Complete.\n")
                print("Full Conversation:\n", conversation)

                # Now summarize
                summary_prompt = create_summary_prompt(name1, name2, conversation)
                summary_chain = LLMChain(llm=llm, prompt=summary_prompt)

                st.subheader("Summary and Title:")
                st.write("**Summarizing the conversation...**")
                print("Summarizing the conversation...")

                summary = summary_chain.run(chat_history="", input="")
                st.write(summary)
                print("Summary:\n", summary)

            except Exception as e:
                st.error(f"Error generating conversation: {e}")
                print(f"Error generating conversation: {e}")

if __name__ == "__main__":
    main()