import json
from transformers import pipeline
import docx
from bs4 import BeautifulSoup
import streamlit as st

# Load the Hugging Face NER pipeline using a pre-trained model
nlp = pipeline("ner", model="dbmdz/bert-large-cased-finetuned-conll03-english")

def extract_text_from_file(uploaded_file):
    """Extract text from various file formats."""
    file_type = uploaded_file.name.split(".")[-1].lower()
    
    if file_type == "txt":
        text = uploaded_file.getvalue().decode("utf-8")
    
    elif file_type == "html":
        soup = BeautifulSoup(uploaded_file.getvalue(), "html.parser")
        text = soup.get_text()
    
    elif file_type == "json":
        data = json.load(uploaded_file)
        text = "\n".join(str(value) for value in data.values()) if isinstance(data, dict) else json.dumps(data)
    
    elif file_type == "docx":
        doc = docx.Document(uploaded_file)
        text = "\n".join([para.text for para in doc.paragraphs])
    
    else:
        raise ValueError("Unsupported file format")

    return text

def infer_relationships(chat_data):
    """Process large text inputs in chunks to avoid memory issues."""
    relationship_scores = []
    max_chars = 1000  # Keep within the model's safe limits

    # Split the text into chunks if it's too long
    for i in range(0, len(chat_data), max_chars):
        chunk = chat_data[i:i + max_chars]
        relationship_scores.append(analyze_relationships(chunk))

    return relationship_scores

def analyze_relationships(text_chunk):
    """Analyze relationships (NER) in text using the Hugging Face model."""
    # Run Named Entity Recognition
    entities = nlp(text_chunk)
    return {"entities": [(ent['word'], ent['entity']) for ent in entities]}

def main():
    st.title("Chat Relationship Analyzer")

    uploaded_file = st.file_uploader("Upload a file", type=["txt", "html", "json", "docx"])

    if uploaded_file is not None:
        try:
            chat_text = extract_text_from_file(uploaded_file)
            
            # If the text is too long, warn the user and process in chunks
            if len(chat_text) > 1000:
                st.warning("Text is too long. Processing in chunks.")

            relationship_scores = infer_relationships(chat_text)
            st.json(relationship_scores)

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

if __name__ == "__main__":
    main()