import os, io, re, json, base64, requests, numpy as np
import streamlit as st
from pypdf import PdfReader
import matplotlib.pyplot as plt

# -----------------------------
# Config
# -----------------------------
st.set_page_config(page_title="PDF Summarizer + Audio + QA", page_icon="📄", layout="wide")

HF_TOKEN = os.environ.get("HF_TOKEN", st.secrets.get("HF_TOKEN", ""))
HEADERS_JSON = {
    "Authorization": f"Bearer {HF_TOKEN}" if HF_TOKEN else "",
    "Content-Type": "application/json",
    "Accept": "application/json",
}

SUMMARIZER_MODEL = "facebook/bart-large-cnn"
TTS_MODEL = "facebook/mms-tts-eng"
EMB_MODEL = "sentence-transformers/all-MiniLM-L6-v2"
QA_MODEL = "deepset/roberta-base-squad2"

# -----------------------------
# API helpers
# -----------------------------
def hf_infer_json(model_id: str, payload: dict, router=False, accept=None):
    if router:
        url = f"https://router.huggingface.co/hf-inference/models/{model_id}"
    else:
        url = f"https://api-inference.huggingface.co/models/{model_id}"
    headers = HEADERS_JSON.copy()
    if accept:
        headers["Accept"] = accept
    r = requests.post(url, headers=headers, data=json.dumps(payload), timeout=120)
    r.raise_for_status()
    try:
        return r.json()
    except requests.exceptions.JSONDecodeError:
        return r.content

def split_into_chunks(text: str, max_chars: int = 1800, overlap: int = 200):
    text = re.sub(r"\s+", " ", text).strip()
    chunks = []
    i = 0
    while i < len(text):
        chunk = text[i:i+max_chars]
        last_dot = chunk.rfind(". ")
        if last_dot > 400:
            chunk = chunk[:last_dot+1]
            i += last_dot + 1 - overlap
        else:
            i += max_chars - overlap
        chunks.append(chunk.strip())
    return [c for c in chunks if c]

def embed_texts(texts):
    url = f"https://router.huggingface.co/hf-inference/models/{EMB_MODEL}/pipeline/feature-extraction"
    headers = {
        "Authorization": f"Bearer {HF_TOKEN}" if HF_TOKEN else "",
        "Content-Type": "application/json",
        "Accept": "application/json",
    }
    r = requests.post(url, headers=headers, data=json.dumps({"inputs": texts}), timeout=120)
    r.raise_for_status()
    arr = np.array(r.json(), dtype=np.float32)
    if arr.ndim == 2:
        return arr.mean(axis=0, keepdims=True)
    if arr.ndim == 3:
        pooled = [a.mean(axis=0) for a in arr]
        return np.vstack(pooled)
    return np.array(arr)

def cosine_sim(a, b):
    a = a / (np.linalg.norm(a, axis=-1, keepdims=True) + 1e-8)
    b = b / (np.linalg.norm(b, axis=-1, keepdims=True) + 1e-8)
    return a @ b.T

def summarize_long_text(text: str, per_chunk_max_len=220, final_max_len=250):
    chunks = split_into_chunks(text, max_chars=1800, overlap=200)
    mini_summaries = []
    for c in chunks:
        out = hf_infer_json(
            SUMMARIZER_MODEL,
            {"inputs": c, "parameters": {"max_length": per_chunk_max_len, "min_length": 60, "do_sample": False}},
            router=False
        )
        if isinstance(out, list) and len(out) and "summary_text" in out[0]:
            mini_summaries.append(out[0]["summary_text"])
        else:
            mini_summaries.append(c[:1000])
    joined = " ".join(mini_summaries)
    final = hf_infer_json(
        SUMMARIZER_MODEL,
        {"inputs": joined, "parameters": {"max_length": final_max_len, "min_length": 80, "do_sample": False}},
        router=False
    )
    if isinstance(final, list) and len(final) and "summary_text" in final[0]:
        return final[0]["summary_text"], chunks
    return joined[:1200], chunks

def tts_wav_bytes(text: str) -> bytes:
    res = hf_infer_json(TTS_MODEL, {"inputs": text}, router=False, accept="audio/wav")
    if isinstance(res, (bytes, bytearray)):
        return res
    if isinstance(res, dict) and "audio" in res:
        try:
            return base64.b64decode(res["audio"])
        except:
            pass
    raise RuntimeError("TTS API did not return audio bytes.")

def extract_text_from_pdf(file) -> str:
    reader = PdfReader(file)
    pages = []
    for p in reader.pages:
        try:
            pages.append(p.extract_text() or "")
        except:
            pages.append("")
    return "\n".join(pages)

def make_word_freq_chart(text: str, top_k=20):
    text = text.lower()
    stop = set(("the a an and of to in is are for with on by as at this that from be was were it its it’s into or if not your you we they their our can may such more most other also than which".split()))
    tokens = re.findall(r"[a-zA-Z]{3,}", text)
    freq = {}
    for t in tokens:
        if t in stop:
            continue
        freq[t] = freq.get(t, 0) + 1
    items = sorted(freq.items(), key=lambda x: x[1], reverse=True)[:top_k]
    if not items:
        st.info("Not enough text to show a frequency chart.")
        return
    words, counts = zip(*items)
    fig = plt.figure()
    plt.bar(words, counts)
    plt.xticks(rotation=60, ha="right")
    plt.title("Top word frequencies")
    plt.tight_layout()
    st.pyplot(fig)

# -----------------------------
# UI
# -----------------------------
st.title("📄 PDF → Summary · 🔊 Audio · 📊 Chart · ❓ Q&A")
st.caption("Powered by Hugging Face Hosted Inference API (free models).")

if not HF_TOKEN:
    st.warning("Set HF_TOKEN in environment or in your Space secrets to use the Hosted Inference API.")

uploaded = st.file_uploader("Upload a PDF", type=["pdf"])

if "doc_text" not in st.session_state:
    st.session_state.doc_text = ""
    st.session_state.chunks = []
    st.session_state.chunk_vecs = None
    st.session_state.summary = ""

if uploaded:
    with st.spinner("Extracting text..."):
        text = extract_text_from_pdf(uploaded)
        st.session_state.doc_text = text
        st.success(f"Loaded {len(text)} characters.")

    st.write("### Actions")
    c1, c2, c3 = st.columns(3)

    with c1:
        if st.button("📝 Summarize"):
            with st.spinner("Summarizing..."):
                summary, chunks = summarize_long_text(st.session_state.doc_text)
                st.session_state.summary = summary
                st.session_state.chunks = chunks
            st.success("Summary ready.")
            st.write("#### Summary")
            st.write(summary)

    with c2:
        if st.button("🔊 Generate Audio (summary)"):
            target_text = st.session_state.summary or st.session_state.doc_text[:1200]
            with st.spinner("Generating audio..."):
                try:
                    wav = tts_wav_bytes(target_text)
                    st.audio(wav, format="audio/wav")
                    st.success("Audio ready.")
                except Exception as e:
                    st.error(f"TTS failed: {e}")

    with c3:
        if st.button("📊 Show Word-Frequency Chart"):
            with st.spinner("Building chart..."):
                make_word_freq_chart(st.session_state.doc_text)

    st.write("---")
    st.subheader("Ask questions about the PDF")
    question = st.text_input("Your question")
    if st.button("Answer"):
        if not st.session_state.chunks:
            st.session_state.chunks = split_into_chunks(st.session_state.doc_text)
        with st.spinner("Thinking..."):
            try:
                if st.session_state.chunk_vecs is None:
                    vecs = embed_texts(st.session_state.chunks)
                    st.session_state.chunk_vecs = vecs
                else:
                    vecs = st.session_state.chunk_vecs

                q_vec = embed_texts([question])
                sims = cosine_sim(q_vec, vecs).flatten()
                top_idx = np.argsort(sims)[::-1][:3]
                context = "\n".join([st.session_state.chunks[i] for i in top_idx])

                qa_out = hf_infer_json(QA_MODEL, {"inputs": {"question": question, "context": context}}, router=False)
                if isinstance(qa_out, dict):
                    ans = qa_out.get("answer", "")
                    score = qa_out.get("score", 0.0)
                elif isinstance(qa_out, list) and len(qa_out) and isinstance(qa_out[0], dict):
                    ans = qa_out[0].get("answer", "")
                    score = qa_out[0].get("score", 0.0)
                else:
                    ans, score = "", 0.0

                st.write("**Answer:**", ans or "_(no confident answer)_")
                st.caption(f"Confidence: {score:.3f}")
                with st.expander("Context used"):
                    st.write(context)
            except Exception as e:
                st.error(f"QA failed: {e}")

else:
    st.info("Upload a PDF to get started.")