import torch
from transformers import pipeline
import spacy
import textwrap
import pdfplumber
import docx
import streamlit as st
import subprocess
import concurrent.futures

# Ensure spaCy model is downloaded
def ensure_spacy_model():
    try:
        nlp = spacy.load("en_core_web_sm")
    except OSError:
        subprocess.run(["python", "-m", "spacy", "download", "en_core_web_sm"])
        nlp = spacy.load("en_core_web_sm")
    return nlp

# Load pre-trained summarization model
summarizer = pipeline("summarization", model="facebook/bart-large-cnn")

# Load spaCy model for Named Entity Recognition (NER)
nlp = ensure_spacy_model()

def chunk_text(text, max_tokens=512):
    """Splits long text into smaller chunks for summarization, ensuring no words are cut off mid-sentence."""
    sentences = text.split('. ')
    chunks = []
    current_chunk = ""
    for sentence in sentences:
        if len(current_chunk) + len(sentence) < max_tokens:
            current_chunk += sentence + '. '
        else:
            if current_chunk.strip():
                chunks.append(current_chunk.strip())
            current_chunk = sentence + '. '
    if current_chunk.strip():
        chunks.append(current_chunk.strip())
    return chunks

def summarize_chunk(chunk, min_length=50, max_length=200):
    """Summarizes a single chunk of text."""
    try:
        result = summarizer(chunk, max_length=max_length, min_length=min_length, do_sample=False)
        return result[0]['summary_text'] if result and isinstance(result, list) and 'summary_text' in result[0] else "[Error: Unexpected output from summarization model]"
    except Exception as e:
        return f"[Error: {str(e)}]"

def summarize_lecture(transcript, min_length=50, max_length=200):
    """Summarizes a long lecture transcript using parallel processing."""
    chunks = chunk_text(transcript)
    if not chunks:
        return "Error: No valid text found for summarization."
    
    summaries = []
    with concurrent.futures.ThreadPoolExecutor() as executor:
        summaries = list(executor.map(lambda chunk: summarize_chunk(chunk, min_length, max_length), chunks))
    
    return "\n".join(summaries)

def extract_key_points(text):
    """Extracts key points using NER while filtering out irrelevant entity types."""
    doc = nlp(text)
    relevant_labels = {"ORG", "PERSON", "GPE", "EVENT", "WORK_OF_ART", "CONCEPT"}  # Focus on meaningful entities
    key_points = {}
    for ent in doc.ents:
        if ent.label_ in relevant_labels and len(ent.text) > 2:
            key_points.setdefault(ent.label_, set()).add(ent.text)
    return key_points

def extract_text_from_pdf(pdf_file):
    """Extract text from a PDF file."""
    text = ""
    with pdfplumber.open(pdf_file) as pdf:
        for page in pdf.pages:
            extracted_text = page.extract_text()
            if extracted_text:
                text += extracted_text + "\n"
    return text.strip()

def extract_text_from_docx(docx_file):
    """Extract text from a DOCX file."""
    doc = docx.Document(docx_file)
    return "\n".join([para.text for para in doc.paragraphs if para.text]).strip()

# Streamlit UI (saved as app.py for Hugging Face deployment)
def main():
    st.title("Lecture Summarizer")
    st.write("Upload a lecture transcript (TXT, PDF, DOCX) to generate key points and a summary.")
    
    uploaded_file = st.file_uploader("Choose a file", type=["txt", "pdf", "docx"])
    if uploaded_file is not None:
        file_type = uploaded_file.name.split(".")[-1]
        
        if file_type == "pdf":
            transcript = extract_text_from_pdf(uploaded_file)
        elif file_type == "docx":
            transcript = extract_text_from_docx(uploaded_file)
        else:
            transcript = uploaded_file.read().decode("utf-8").strip()
        
        if not transcript:
            st.error("Error: No text found in the uploaded file.")
            return
        
        if st.button("Summarize Lecture"):
            with st.spinner("Generating summary..."):
                summary = summarize_lecture(transcript)
                key_points = extract_key_points(transcript) if summary and not summary.startswith("Error") else {}
            
            st.subheader("Lecture Summary")
            st.write(textwrap.fill(summary, width=80))
            
            st.subheader("Key Points")
            if key_points:
                for label, items in key_points.items():
                    st.write(f"**{label}**: {', '.join(items)}")
            else:
                st.write("No key points could be extracted.")

if __name__ == "__main__":
    main()