import streamlit as st
import os
import shutil
import tempfile
import numpy as np
import soundfile as sf
from dotenv import load_dotenv
from faster_whisper import WhisperModel
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain_core.prompts import PromptTemplate
from langchain_huggingface import HuggingFaceEmbeddings
from langchain_community.vectorstores import Chroma
from langchain_community.document_loaders import (
    PyPDFLoader,
    TextLoader,
    UnstructuredPowerPointLoader
)
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.retrievers.document_compressors import LLMChainExtractor

try:
    # Newer structure (LangChain 0.3.27+)
    from langchain_community.retrievers.contextual_compression import ContextualCompressionRetriever
except ImportError:
    # Fallback for older builds
    from langchain.retrievers import ContextualCompressionRetriever

# Streamlit WebRTC for live voice
from streamlit_webrtc import webrtc_streamer, AudioProcessorBase, WebRtcMode

# ------------------------------
# 🌍 Load Environment
# ------------------------------
load_dotenv()
GOOGLE_API_KEY = os.getenv("GOOGLE_API_KEY")
if not GOOGLE_API_KEY:
    st.error("⚠️ GOOGLE_API_KEY not found! Add it to .env or repo secrets.")
    st.stop()
os.environ["GOOGLE_API_KEY"] = GOOGLE_API_KEY

# ------------------------------
# ⚙️ Config
# ------------------------------
CHROMA_DIR = "./chroma_db"
MAX_MEMORY_TURNS = 10

# ------------------------------
# 🧩 Load Documents Helper
# ------------------------------
def load_documents(uploaded_files):
    docs = []
    for file in uploaded_files:
        temp_path = os.path.join(tempfile.gettempdir(), file.name)
        with open(temp_path, "wb") as f:
            f.write(file.getbuffer())
        if file.name.endswith(".pdf"):
            loader = PyPDFLoader(temp_path)
        elif file.name.endswith(".txt"):
            loader = TextLoader(temp_path, encoding="utf-8")
        elif file.name.endswith(".pptx"):
            loader = UnstructuredPowerPointLoader(temp_path)
        else:
            continue
        docs.extend(loader.load())
    return docs

# ------------------------------
# 🧠 Embeddings
# ------------------------------
@st.cache_resource
def get_embeddings():
    return HuggingFaceEmbeddings(model_name="sentence-transformers/paraphrase-MiniLM-L3-v2")
embeddings = get_embeddings()

# ------------------------------
# 💾 Session State
# ------------------------------
if "chat_history" not in st.session_state:
    st.session_state.chat_history = []
if "retriever" not in st.session_state:
    st.session_state.retriever = None

# ------------------------------
# 🧭 Sidebar Controls
# ------------------------------
st.sidebar.header("⚙️ Controls")
mode = st.sidebar.radio("Choose Mode:", ["🚀 Fast Mode", "🎯 Accurate Mode"])
uploaded_files = st.sidebar.file_uploader(
    "📂 Upload PDF/TXT/PPTX files",
    type=["pdf", "txt", "pptx"],
    accept_multiple_files=True,
)

if st.sidebar.button("🧹 Clear Chat"):
    st.session_state.chat_history = []
    st.rerun()

if st.sidebar.button("🗑️ Clear Database"):
    if os.path.exists(CHROMA_DIR):
        shutil.rmtree(CHROMA_DIR)
        st.session_state.retriever = None
        st.session_state.chat_history = []
        st.sidebar.success("✅ Database and chat cleared!")
    else:
        st.sidebar.warning("⚠️ No database found.")

# ------------------------------
# 🧾 Document Processing
# ------------------------------
if uploaded_files:
    splitter = RecursiveCharacterTextSplitter(
        chunk_size=250 if mode == "🚀 Fast Mode" else 800,
        chunk_overlap=30 if mode == "🚀 Fast Mode" else 100,
    )
    documents = load_documents(uploaded_files)
    split_docs = splitter.split_documents(documents)

    vectorstore = Chroma.from_documents(split_docs, embeddings, persist_directory=CHROMA_DIR)
    vectorstore.persist()

    if mode == "🎯 Accurate Mode":
        llm_for_compression = ChatGoogleGenerativeAI(
            model="gemini-2.5-flash",
            google_api_key=GOOGLE_API_KEY,
            temperature=0.2
        )
        compressor = LLMChainExtractor.from_llm(llm_for_compression)
        retriever = ContextualCompressionRetriever(
            base_compressor=compressor,
            base_retriever=vectorstore.as_retriever(search_kwargs={"k": 5}),
        )
    else:
        retriever = vectorstore.as_retriever(search_kwargs={"k": 3})

    st.session_state.retriever = retriever
    st.sidebar.success(f"✅ Docs indexed in *{mode}* mode!")

# ------------------------------
# 💬 Main Chat Interface
# ------------------------------
st.title("📚 AskDocs – Smart Multimodal Chatbot")

for msg in st.session_state.chat_history:
    with st.chat_message(msg["role"]):
        st.markdown(msg["content"])

# ------------------------------
# 🎙️ Live Voice Input
# ------------------------------
st.subheader("🎤 Speak Your Question")

class WhisperAudioProcessor(AudioProcessorBase):
    def __init__(self):
        self.frames = []

    def recv_audio(self, frame):
        self.frames.append(frame.to_ndarray())
        return frame

webrtc_ctx = webrtc_streamer(
    key="whisper",
    mode=WebRtcMode.SENDRECV,
    audio_processor_factory=WhisperAudioProcessor,
    media_stream_constraints={"audio": True, "video": False},
    async_processing=True
)

voice_question = None
if webrtc_ctx.audio_processor and webrtc_ctx.audio_processor.frames:
    frames = webrtc_ctx.audio_processor.frames
    audio_data = np.concatenate(frames, axis=0)
    tmp_path = os.path.join(tempfile.gettempdir(), "voice.wav")
    sf.write(tmp_path, audio_data, 48000)

    with st.spinner("🎧 Transcribing voice..."):
        model = WhisperModel("base", device="cpu", compute_type="int8")
        segments, _ = model.transcribe(tmp_path)
        voice_question = " ".join([seg.text for seg in segments])
        st.success(f"🗣️ You said: {voice_question}")

# ------------------------------
# 💬 Text Chat Input
# ------------------------------
st.markdown("---")
user_input = st.chat_input("Type your question here...")
question = voice_question or user_input

# ------------------------------
# 🤖 Response Generation
# ------------------------------
if question:
    st.session_state.chat_history.append({"role": "user", "content": question})
    with st.chat_message("user"):
        st.markdown(question)

    if not st.session_state.retriever:
        with st.chat_message("assistant"):
            st.error("⚠️ Please upload documents first!")
    else:
        retrieved_docs = st.session_state.retriever.invoke(question)
        context_text = "\n\n".join(doc.page_content for doc in retrieved_docs)
        history_text = "\n".join([
            f"{t['role']}: {t['content']}"
            for t in st.session_state.chat_history[-MAX_MEMORY_TURNS * 2:]
        ])

        llm = ChatGoogleGenerativeAI(
            model="gemini-2.5-flash",
            google_api_key=GOOGLE_API_KEY,
            temperature=0.2
        )

        template = """
        You are a smart and reliable assistant. Use the chat history and retrieved context to answer.
        If unsure, say so.

        Chat History:
        {history}

        Retrieved Context:
        {context}

        Question:
        {question}

        Answer Instructions:
        - Provide answers in this order:
          1. *LLM Answer:* Direct and concise.
          2. *Explanation:* Detailed reasoning and context.
          3. *Summary:* Short key takeaway.
        - Cite sources when relevant.
        Answer:
        """
        prompt = PromptTemplate.from_template(template)
        final_prompt = prompt.invoke({
            "history": history_text,
            "context": context_text,
            "question": question
        })

        with st.chat_message("assistant"):
            with st.spinner("🤖 Thinking..."):
                response = llm.invoke(final_prompt).content
                st.write_stream(iter([response]))

        st.session_state.chat_history.append({"role": "assistant", "content": response})

# ------------------------------
# 🪄 Footer
# ------------------------------
st.divider()
st.caption("⚡ Powered by Gemini + LangChain + HuggingFace + Streamlit + Whisper + WebRTC")