import streamlit as st
import os
import re

os.environ["OPENAI_API_KEY"] = "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"

from langchain_openai import ChatOpenAI
from openai import OpenAI
import tempfile

client = OpenAI()

def simple_split(text, chunk_size=1000):
    """Pure Python splitter"""
    sentences = re.split(r'[.!?]\s+', text)
    chunks = []
    current_chunk = ""
    
    for sentence in sentences:
        if len(current_chunk + sentence) < chunk_size:
            current_chunk += sentence + ". "
        else:
            if current_chunk:
                chunks.append(current_chunk.strip())
            current_chunk = sentence + ". "
    
    if current_chunk:
        chunks.append(current_chunk.strip())
    
    return chunks

def dynamic_rag(query, document_content):
    """Dynamic RAG - no external deps"""
    chunks = simple_split(document_content)
    
    # Simple similarity (keyword matching)
    best_chunks = []
    query_words = set(query.lower().split())
    
    for chunk in chunks:
        chunk_words = set(chunk.lower().split())
        score = len(query_words.intersection(chunk_words))
        best_chunks.append((score, chunk))
    
    best_chunks.sort(reverse=True, key=lambda x: x[0])
    context = "\n".join([chunk for score, chunk in best_chunks[:3]])
    
    prompt = f"""Use ONLY this context from document:

{context}

Question: {query}

Answer using context only:"""
    
    response = client.chat.completions.create(
        model="gpt-4o-mini",
        messages=[{"role": "user", "content": prompt}],
        temperature=0
    )
    return response.choices[0].message.content

st.title("🧠 Dynamic RAG Chatbot")
st.markdown("**Paste text or upload → Ask ANY question!**")

# Input options
col1, col2 = st.columns(2)

with col1:
    uploaded_file = st.file_uploader("📤 Upload TXT", type='txt')
    
with col2:
    pasted_text = st.text_area("📝 Or paste text here", height=150)

document_content = ""

if uploaded_file is not None:
    content = uploaded_file.read().decode('utf-8')
    document_content = content
    st.success("✅ TXT loaded!")
elif pasted_text:
    document_content = pasted_text
    st.success("✅ Text loaded!")

if document_content:
    st.session_state.document_content = document_content
    st.success("🚀 Chatbot ready! Ask about your text.")

if 'document_content' in st.session_state:
    if "messages" not in st.session_state:
        st.session_state.messages = []
    
    # Chat history
    for message in st.session_state.messages:
        with st.chat_message(message["role"]):
            st.markdown(message["content"])
    
    # Chat input
    if query := st.chat_input("💬 Ask about your document..."):
        st.session_state.messages.append({"role": "user", "content": query})
        with st.chat_message("user"):
            st.markdown(query)
        
        with st.chat_message("assistant"):
            with st.spinner("🔍 Searching document..."):
                response = dynamic_rag(query, st.session_state.document_content)
                st.markdown(response)
        
        st.session_state.messages.append({"role": "assistant", "content": response})

    # Clear
    if st.button("🗑️ Clear Chat"):
        st.session_state.messages = []
        st.rerun()

else:
    st.info("👆 **Paste text or upload TXT to start chatting!**")
    st.markdown("""
    **Test example:**
    ```
    Skills: Python, DSA, AI/ML
    Projects: RAG Chatbot (live demo)
    LeetCode: 300 problems solved
    ```
    Ask: "What projects?" → Perfect answer!
    """)

st.sidebar.markdown("### 🛠️ Pure Python RAG")
st.markdown("• Custom text splitter")
st.markdown("• Keyword similarity")
st.markdown("• OpenAI GPT-4o-mini")
st.markdown("• Dynamic input")