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", "")
HEADERS_JSON = {
    "Authorization": f"Bearer {HF_TOKEN}" if HF_TOKEN else "",
    "Content-Type": "application/json",
    "Accept": "application/json",
}

SUMMARIZER_MODEL = "pszemraj/long-t5-tglobal-base-16384-book-summary"
TTS_MODELS = [
    "espnet/kan-bayashi_ljspeech_vits",
    "facebook/fastspeech2-en-ljspeech"
]
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 = 1500, 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 = HEADERS_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):
    chunks = split_into_chunks(text)
    mini_summaries = []
    for c in chunks:
        out = hf_infer_json(SUMMARIZER_MODEL, {"inputs": c}, 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[:800])
    return " ".join(mini_summaries), chunks

def tts_wav_bytes(text: str) -> bytes:
    for model in TTS_MODELS:
        try:
            res = hf_infer_json(model, {"inputs": text}, router=False, accept="audio/wav")
            if isinstance(res, (bytes, bytearray)):
                return res
            if isinstance(res, dict) and "audio" in res:
                return base64.b64decode(res["audio"])
        except Exception:
            continue
    raise RuntimeError("All TTS models failed.")

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("Free models via Hugging Face 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")
    with st.container():
        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(st.session_state.summary)

    with st.container():
        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 st.container():
        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.")