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

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

# Prefer environment variable (Spaces sets secrets as env vars), *then* try st.secrets safely.
HF_TOKEN = os.environ.get("HF_TOKEN", "")
if not HF_TOKEN:
    try:
        # Access st.secrets inside try/except so we don't crash when no secrets file exists.
        HF_TOKEN = st.secrets.get("HF_TOKEN", "") if hasattr(st, "secrets") else ""
    except Exception:
        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"

# -----------------------------
# Helper: Hugging Face inference
# -----------------------------
def hf_infer_json(model_id: str, payload: dict, router=False, accept=None, timeout=120):
    """
    Send request to Hugging Face Hosted Inference API.
    If `router=True` we'll use the router base path (useful for some pipelines).
    If backend returns binary (audio), this returns raw bytes.
    """
    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

    try:
        r = requests.post(url, headers=headers, data=json.dumps(payload), timeout=timeout)
        r.raise_for_status()
    except requests.exceptions.RequestException as e:
        # Bubble up a useful message
        raise RuntimeError(f"Hugging Face request failed: {e}")

    # Try to decode JSON; if fails, return bytes/content
    try:
        return r.json()
    except ValueError:
        return r.content

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

def clean_text(s: str) -> str:
    return re.sub(r"\s+", " ", s).strip()

def split_into_chunks(text: str, max_chars: int = 1800, overlap: int = 200):
    text = clean_text(text)
    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]

# -----------------------------
# Embeddings + similarity
# -----------------------------
def embed_texts(texts):
    """
    Calls the feature-extraction pipeline on the router endpoint.
    Returns numpy array shape (n_texts, dim)
    """
    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",
    }
    try:
        r = requests.post(url, headers=headers, data=json.dumps({"inputs": texts}), timeout=120)
        r.raise_for_status()
    except requests.exceptions.RequestException as e:
        raise RuntimeError(f"Embedding request failed: {e}")

    arr = np.array(r.json(), dtype=np.float32)

    # Cases:
    # - arr.ndim == 1 -> single vector (dim,) -> reshape to (1,dim)
    # - arr.ndim == 2 -> batch of vectors (n, dim) -> return as-is
    # - arr.ndim == 3 -> model returned token-level vectors per item: mean-pool per item -> (n, dim)
    if arr.ndim == 1:
        return arr.reshape(1, -1)
    if arr.ndim == 2:
        return arr
    if arr.ndim == 3:
        pooled = np.array([a.mean(axis=0) for a in arr])
        return pooled
    # Fallback
    return arr.reshape(arr.shape[0], -1)

def cosine_sim(a, b):
    """
    a: (m, d), b: (n, d) -> returns (m, n)
    """
    a_n = a / (np.linalg.norm(a, axis=-1, keepdims=True) + 1e-8)
    b_n = b / (np.linalg.norm(b, axis=-1, keepdims=True) + 1e-8)
    return a_n @ b_n.T

# -----------------------------
# Summarization
# -----------------------------
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:
        try:
            out = hf_infer_json(
                SUMMARIZER_MODEL,
                {"inputs": c, "parameters": {"max_length": per_chunk_max_len, "min_length": 60, "do_sample": False}},
                router=False,
            )
        except Exception as e:
            # if API fails, include the chunk (truncated) as fallback
            mini_summaries.append(c[:1000])
            continue

        # Hosted inference often returns a list of dicts with 'summary_text'
        if isinstance(out, list) and len(out) and isinstance(out[0], dict) and "summary_text" in out[0]:
            mini_summaries.append(out[0]["summary_text"])
        elif isinstance(out, dict) and "summary_text" in out:
            mini_summaries.append(out["summary_text"])
        else:
            mini_summaries.append(c[:1000])

    joined = " ".join(mini_summaries)
    try:
        final = hf_infer_json(
            SUMMARIZER_MODEL,
            {"inputs": joined, "parameters": {"max_length": final_max_len, "min_length": 80, "do_sample": False}},
            router=False,
        )
    except Exception:
        return joined[:1200], chunks

    if isinstance(final, list) and len(final) and isinstance(final[0], dict) and "summary_text" in final[0]:
        return final[0]["summary_text"], chunks
    if isinstance(final, dict) and "summary_text" in final:
        return final["summary_text"], chunks

    return joined[:1200], chunks

# -----------------------------
# TTS
# -----------------------------
def tts_wav_bytes(text: str) -> bytes:
    try:
        res = hf_infer_json(TTS_MODEL, {"inputs": text}, router=False, accept="audio/wav", timeout=180)
    except Exception as e:
        raise RuntimeError(f"TTS request failed: {e}")

    if isinstance(res, (bytes, bytearray)):
        return res
    if isinstance(res, dict) and "audio" in res:
        try:
            return base64.b64decode(res["audio"])
        except Exception:
            pass
    raise RuntimeError("TTS API did not return audio bytes.")

# -----------------------------
# Visualization helper
# -----------------------------
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("No HF_TOKEN found. Add HF_TOKEN in Space Settings → Secrets (recommended). The app will still run but HF API calls will fail without a token.")

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

# session state
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..."):
                try:
                    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)
                except Exception as e:
                    st.error(f"Summarization failed: {e}")

    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:
                # embed once/cache
                if st.session_state.chunk_vecs is None:
                    st.session_state.chunk_vecs = embed_texts(st.session_state.chunks)
                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.")