import streamlit as st
from pi_shard import pi_shard, get_pi_digits
from gpt_utils import analyze_chunk
from pi_utils import random_pi_fact, generate_pi_graph
from pi_vector_utils import get_embedding, pi_rotation, pi_modulated_similarity
import fitz
import docx

st.set_page_config(page_title="Play with Pi", layout="wide")
st.title("🎲 Play with Pi - π-Based Chunking Engine")

st.sidebar.header("🔧 Controls")
openai_key = st.sidebar.text_input("OpenAI API Key", type="password")
uploaded_file = st.file_uploader("Upload a document", type=["txt", "pdf", "docx"])

if uploaded_file:
    # Handle uploaded file types
    if uploaded_file.name.endswith(".txt"):
        text = uploaded_file.read().decode("utf-8")
    elif uploaded_file.name.endswith(".pdf"):
        doc = fitz.open(stream=uploaded_file.read(), filetype="pdf")
        text = " ".join([page.get_text() for page in doc])
    elif uploaded_file.name.endswith(".docx"):
        doc = docx.Document(uploaded_file)
        text = "\n".join([para.text for para in doc.paragraphs])

    st.subheader("📄 Original Document")
    st.text_area("Document Preview", text[:1000] + "...", height=150)

    # Create π-based chunks
    chunks = pi_shard(text)
    st.subheader(f"🔍 π-Shards (Total: {len(chunks)})")
    selected = st.selectbox("Select Chunk", range(len(chunks)))
    st.code(chunks[selected], language="markdown")

    # GPT Analysis of Selected Chunk
    if openai_key:
        st.markdown("#### ✨ GPT Analysis")
        if st.button("Analyze Selected Chunk"):
            with st.spinner("Thinking like π..."):
                result = analyze_chunk(chunks[selected], openai_key)
                st.success("Done!")
                st.markdown(result)

    # Question Answering Section
    st.markdown("#### 🤔 Ask a Question about the Document")
    user_query = st.text_area("Enter your question:", "")

    if openai_key and st.button("🚀 Submit"):
        if user_query:
            st.info("Generating embeddings and rotating using π...")
            pi_digits = get_pi_digits(len(chunks))
            query_vec = get_embedding(user_query, openai_key)

            scores = []
            for i, chunk in enumerate(chunks):
                chunk_vec = get_embedding(chunk, openai_key)
                rotated = pi_rotation(chunk_vec, pi_digits[i])
                sim = pi_modulated_similarity(query_vec, rotated, pi_digits[i])
                scores.append((i, sim))

            scores.sort(key=lambda x: x[1], reverse=True)
            top_index = scores[0][0]

            st.success(f"✅ Best π-Chunk Match (Chunk #{top_index})")
            st.code(chunks[top_index])

            # Analyze matched chunk with GPT
            st.markdown("#### 📚 GPT Response to Query")
            with st.spinner("Analyzing the matched chunk..."):
                answer = analyze_chunk(chunks[top_index], openai_key)
                st.markdown(answer)

# Sidebar - Pi facts and visualization
st.sidebar.subheader("🎲 Pi Fact")
st.sidebar.info(random_pi_fact())

if st.sidebar.button("🌀 Show π-Graph"):
    fig = generate_pi_graph()
    st.pyplot(fig)