import streamlit as st
from datasets import load_dataset
from langchain.llms import HuggingFaceEndpoint
from langchain.prompts import FewShotChatMessagePromptTemplate, ChatPromptTemplate
from langchain.schema.messages import SystemMessage

# Set page config at the very top, before anything else
st.set_page_config(page_title="DialogSum Few-Shot Summarizer", page_icon="🧠")

# Load few-shot examples from dialogsum
@st.cache_data
def load_examples(n=3):
    dataset = load_dataset("knkarthick/dialogsum", split="train[:20]")
    return [{"dialogue": row["dialogue"], "summary": row["summary"]} for row in dataset.select(range(n))]

examples = load_examples()

# Template for each example
example_prompt = ChatPromptTemplate.from_messages([
    ("human", "Summarize the following dialog:\n\n{dialogue}"),
    ("ai", "{summary}")
])

# Few-shot prompt template (no prefix/suffix here)
few_shot_prompt = FewShotChatMessagePromptTemplate(
    example_prompt=example_prompt,
    examples=examples
)

# Now add intro system message + user input separately
final_prompt = ChatPromptTemplate.from_messages([
    SystemMessage(content="The following are examples of dialogues and their summaries."),
    ("human", "Summarize the following dialog:\n\n{dialogue}")
])

# Streamlit UI setup
st.title("🧠 Few-Shot Dialog Summarizer")
st.markdown("Uses real examples from `dialogsum` to guide the summary output.")

user_input = st.text_area("✍️ Paste your dialogue here:", height=200)

if user_input:
    # Prepare messages for the final prompt (include few-shot examples here directly)
    formatted_prompt = few_shot_prompt.format_messages(dialogue=user_input)
    
    # Add formatted examples to the final prompt
    final_formatted_prompt = final_prompt.format_messages(dialogue=user_input)
    
    # Convert the list of messages into a single string
    prompt_string = "\n".join([msg.content for msg in final_formatted_prompt])  # Access content with .content

    # Get response from model with correct explicit parameters
    llm = HuggingFaceEndpoint(
        repo_id="google/pegasus-xsum",
        task="summarization",  # Correct task for text summarization
        temperature=0.3,  # Explicitly passing temperature here
        max_length=128  # Explicitly passing max_length here
    )
    
    # Run the LLM with the prompt string
    response = llm(prompt_string)

    # Output the summary
    st.subheader("📌 Summary:")
    st.write(response)