import os
import pdfplumber
import streamlit as st
import time
from dotenv import load_dotenv, find_dotenv

from langchain.embeddings import HuggingFaceEmbeddings
from langchain_chroma import Chroma
from langchain_core.documents import Document
from langchain_text_splitters import RecursiveCharacterTextSplitter

from langchain_groq.chat_models import ChatGroq
from langchain.prompts import PromptTemplate
from langchain.chains import RetrievalQA

# ------------------ CONFIG ------------------

load_dotenv(find_dotenv())
PDF_PATH = os.path.join(os.path.dirname(__file__), "docs", "document.pdf")
CHROMA_DIR = "chroma_db"
EMBED_MODEL = "sentence-transformers/all-MiniLM-L6-v2"  # Modèle HF
GROQ_MODEL = "llama3-70b-8192"

# ------------------ UTILS ------------------

@st.cache_resource
def load_vectorstore():
    embedding = HuggingFaceEmbeddings(model_name=EMBED_MODEL)

    if not os.path.exists(CHROMA_DIR):
        with pdfplumber.open(PDF_PATH) as pdf:
            full_text = ""
            for page in pdf.pages:
                full_text += page.extract_text() + "\n"

        if not full_text.strip():
            st.error("❌ Le PDF est vide ou illisible.")
            st.stop()

        cleaned = " ".join(full_text.split())
        doc = Document(page_content=cleaned, metadata={"file": os.path.basename(PDF_PATH)})

        splitter = RecursiveCharacterTextSplitter(chunk_size=2000, chunk_overlap=100)
        chunks = splitter.split_documents([doc])
        return Chroma.from_documents(chunks, embedding=embedding, persist_directory=CHROMA_DIR)
    else:
        return Chroma(persist_directory=CHROMA_DIR, embedding_function=embedding)

def create_qa_chain(db):
    retriever = db.as_retriever(search_type="mmr", search_kwargs={"k": 5})
    llm = ChatGroq(model=GROQ_MODEL, api_key=os.getenv("GROQ_API_KEY"))

    prompt = PromptTemplate(
        template="""
        Tu es un assistant agricole expert francophone nommé Terr'Y, créé par Terr'O.
        Tu ne dois répondre qu'à partir du document fourni. Si la réponse n'est pas dans le document, dis-le honnêtement.
        Ta mission est de répondre **uniquement** à partir du contenu du document fourni.

        Si la réponse contient des éléments structurés (ex : primes, cultures, conditions, barèmes, dates, etc.), tu dois les présenter sous forme de **tableau clair**.

        Consignes :
        - Utilise un ton clair, professionnel et en français.
        - Réponds exclusivement à partir des documents.
        - Ne complète jamais avec des données absentes.
        - Si l’information n’est pas présente, indique-le simplement.
        - Tu t'appelles Terr'Y l'assistant agricole créé par Terr'O.

        Question : {question}
        Contenu : {context}
        """,
        input_variables=["question", "context"]
    )

    return RetrievalQA.from_chain_type(
        llm=llm,
        retriever=retriever,
        chain_type="stuff",
        chain_type_kwargs={"prompt": prompt}
    )

def typewriter(text, delay=0.03):
    container = st.empty()
    words = text.split(" ")
    typed = ""
    for word in words:
        typed += word + " "
        container.markdown(typed)
        time.sleep(delay)

# ------------------ INTERFACE STREAMLIT ------------------

st.set_page_config(page_title="Terr'Y - Assistant Collecte", page_icon="🌾", layout="wide")

col1, col2 = st.columns([1, 6])
with col1:
    st.image("images/terrY.png", width=100)
with col2:
    st.title("🌾 Terr'Y - Assistant Agricole")
    st.markdown("Pose une question sur le livret de campagne collecte 2024 📘")

question = st.text_input("Ta question ici :", placeholder="Ex. : Quelles sont les primes pour le blé tendre ?")

if question:
    with st.spinner("Terr'Y réfléchit..."):
        db = load_vectorstore()
        qa = create_qa_chain(db)
        response = qa.invoke(question)
        answer = response["result"]
        st.markdown("### 🧠 Réponse")
        typewriter(answer)