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.

    We will be very explicit and not allow any formatting except the required lines.
    """
    prompt_template_str = f"""
    You are simulating a conversation of exactly 15 messages between two people: {name1} and {name2}.
    {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}.

    Requirements:
    - Output exactly 15 lines, no more, no less.
    - Each line must be a single message in the format:
      {name1}: <message> or {name2}: <message>
    - Do not add any headings, numbers, sample outputs, or explanations.
    - Do not mention code, programming, or instructions.
    - Each message should be 1-2 short sentences, friendly, natural, reflecting the style: {persona_style}.
    - Use everyday language, can ask questions, show opinions.
    - Use emojis sparingly if it fits the style (no more than 1-2 total).
    - No repeated lines, each message should logically follow from the previous one.
    - Do not produce anything after the 15th message. No extra lines or text.

    Produce all 15 messages now:
    """
    return ChatPromptTemplate.from_template(prompt_template_str)

def create_summary_prompt(name1: str, name2: str, conversation: str):
    """Prompt for generating a title and summary."""
    summary_prompt_str = f"""
    Below is a completed 15-message 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, and do not add extra formatting beyond the two lines:
    - One line starting with "Title:"
    - One line starting with "Summary:"
    """
    return ChatPromptTemplate.from_template(summary_prompt_str)

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

    model_names = [
        "meta-llama/Llama-3.3-70B-Instruct",
        "meta-llama/Llama-3.1-405B-Instruct",
        "Qwen/Qwen2.5-72B-Instruct",
        "deepseek-ai/DeepSeek-V3",
        "deepseek-ai/DeepSeek-V2.5"
        
    ]
    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="").strip()

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

                # Summarize the conversation
                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()