import streamlit as st
from flask import Flask, request, jsonify
from transformers import pipeline
from keybert import KeyBERT

# Initialize Flask app
app = Flask(__name__)

# Function to extract text from SRT-formatted text
def extract_text_from_srt_text(srt_text):
    lines = srt_text.strip().split("\n\n")
    texts = [subtitle.split("\n")[2] for subtitle in lines if subtitle.strip()]
    return " ".join(texts)

# Function to generate summary from text
def generate_summary(text, summary_length):
    summarizer = pipeline("summarization")
    summary = summarizer(text, max_length=summary_length, min_length=30, do_sample=False)
    summary_text = summary[0]["summary_text"]
    return summary_text

# Function to extract top 4 topics from text
def extract_top_topics(text, n_top_topics):
    model = KeyBERT('distilbert-base-nli-mean-tokens')
    keywords = model.extract_keywords(text, keyphrase_ngram_range=(1, 3), stop_words='english', use_maxsum=True, nr_candidates=20, top_n=n_top_topics)
    return [topic for topic, _ in keywords]

# Streamlit app
st.title("SRT Summarization")

# Text area for user to input SRT-formatted text
srt_text_input = st.text_area("Paste SRT-formatted text here:")

# Button to trigger summarization
if st.button("Summarize"):
    # Check if text area is not empty
    if srt_text_input.strip():
        # Extract text from SRT-formatted text
        text_to_summarize = extract_text_from_srt_text(srt_text_input)
        # Generate summary
        summary = generate_summary(text_to_summarize, 150)
        # Extract top 4 topics
        top_topics = extract_top_topics(text_to_summarize, 4)
        # Display summary and top 4 topics
        st.subheader("Summary:")
        st.write(summary)
        st.subheader("Top 4 Keywords:")
        for topic in top_topics:
            st.write(f"- {topic}")
    else:
        st.warning("Please enter some SRT-formatted text.")

# Define endpoint for REST API
@app.route("/summarize", methods=["POST"])
def summarize():
    data = request.json
    if "srt_text" not in data:
        return jsonify({"error": "Missing 'srt_text' parameter"}), 400
    srt_text = data["srt_text"]
    text_to_summarize = extract_text_from_srt_text(srt_text)
    summary = generate_summary(text_to_summarize, 150)
    top_topics = extract_top_topics(text_to_summarize, 4)
    return jsonify({"summary": summary, "top_topics": top_topics})

if __name__ == "__main__":
    app.run()