import streamlit as st
import pandas as pd
import os
import torch
from huggingface_hub import InferenceClient
from deep_translator import GoogleTranslator
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_community.vectorstores import FAISS
from langchain_core.documents import Document

# --- CONFIG ---
MODEL_NAME = "mistralai/Mixtral-8x7B-Instruct-v0.1"

# --- Language Support ---
language_map = {
    "English": "en",
    "Malayalam": "ml",
    "Kannada": "kn",
    "Urdu": "ur"
}

# --- Hugging Face Token ---
try:
    hf_token = st.secrets["unicorn"]
except:
    hf_token = os.getenv("unicorn")

client = InferenceClient(
    provider="together",
    api_key=hf_token,
)

# --- Crop Images ---
crop_images = {
    "FAB Cabbage": "https://via.placeholder.com/300?text=Cabbage",
    "FAB Mint": "https://via.placeholder.com/300?text=Mint",
    "FAB Spinach": "https://via.placeholder.com/300?text=Spinach",
    "FAB Mustard": "https://via.placeholder.com/300?text=Mustard",
    "FAB Coriander leaves": "https://via.placeholder.com/300?text=Coriander",
    "FAB Spring onion": "https://via.placeholder.com/300?text=Spring+Onion",
    "FAB Lettuce": "https://via.placeholder.com/300?text=Lettuce",
    "FAB Celery": "https://via.placeholder.com/300?text=Celery",
    "FAB Mushroom": "https://via.placeholder.com/300?text=Mushroom"
}

# --- Translator ---
def translate(text, lang):
    try:
        return GoogleTranslator(source='auto', target=lang).translate(text)
    except:
        return text

# --- Load Dataset ---
@st.cache_data(show_spinner="📅 Loading crop dataset...")
def load_csv(crop):
    folder = os.path.join(os.path.dirname(__file__), "datasets")
    for filename in os.listdir(folder):
        if crop.replace(" ", "").lower() in filename.replace(" ", "").lower():
            return pd.read_csv(os.path.join(folder, filename))
    return None

# --- Ontology Reasoning ---
def apply_ontology(question, crop):
    if "insect" in question or "pest" in question:
        return f"{crop} is often affected by pests like aphids."
    return ""

# --- Knowledge Graph Retrieval ---
def from_knowledge_graph(question, crop):
    if "aphids" in question:
        return "Neem oil is effective against aphids."
    if "fungus" in question:
        return "Try a mix of baking soda and water."
    return ""

# --- FAISS Store ---
@st.cache_resource(show_spinner="🔎 Building vector store...")
def build_store(df):
    docs = [Document(page_content=f"Q: {row['Question']}\nA: {row['Answer']}") for _, row in df.iterrows()]
    device = "cuda" if torch.cuda.is_available() else "cpu"
    emb = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2", model_kwargs={"device": device})
    return FAISS.from_documents(docs, embedding=emb)

# --- LLM Answer with Step-by-Step Instructions ---
def get_llm_answer(context, question, crop):
    user_background = "I am building a crop question-answering chatbot."
    response_style = "Respond in a clear, step-by-step bullet point format. Only answer about the selected crop. Be concise and skip unrelated questions."

    messages = [
        {"role": "system", "content": f"You are a helpful, honest, and intelligent assistant.\n\nThe user has provided the following instructions:\n1. What the assistant should know about the user:\n{user_background}\n2. How the assistant should respond:\n{response_style}\n\nAlways follow these instructions.\n\nCrop: {crop}\n\nContext:\n{context}"},
        {"role": "user", "content": f"Question: {question}\nAnswer:"}
    ]

    try:
        completion = client.chat.completions.create(
            model=MODEL_NAME,
            messages=messages,
        )
        return f"[{crop}] {completion.choices[0].message.content.strip()}"
    except Exception as e:
        return f"LLM Error: {e}"

# --- Question Handler ---
def handle_question(df, vectorstore, crop, question):
    question_lower = question.lower()
    selected_crop_clean = crop.lower().replace("fab", "").strip()

    for other_crop in crop_images:
        other_clean = other_crop.lower().replace("fab", "").strip()
        if other_clean in question_lower and other_clean != selected_crop_clean:
            return f"This question is unrelated to {crop} and will not be addressed.\nPlease ask specifically about {crop} for tailored guidance.", "🚫 Unrelated"

    logic_reason = apply_ontology(question_lower, crop)
    kg_reason = from_knowledge_graph(question_lower, crop)

    retriever = vectorstore.as_retriever(search_type="similarity", k=1)
    doc = retriever.get_relevant_documents(question)[0]
    context = doc.page_content

    if logic_reason or kg_reason:
        context += f"\n\n{logic_reason}\n{kg_reason}"

    return get_llm_answer(context, question, crop), "✅ Combined (Hybrid + Ontology + KG)"

# --- Streamlit App ---
def main():
    st.set_page_config("Crop QA Chatbot", layout="wide")
    st.title("🌾 Crop QA with Hybrid Search + Ontology + Knowledge Graph")

    lang = st.selectbox("🌐 Language", list(language_map.keys()))
    lang_code = language_map[lang]
    translated_crop_names = {translate(c, lang_code): c for c in crop_images}

    cols = st.columns(3)
    for i, (label, orig) in enumerate(translated_crop_names.items()):
        with cols[i % 3]:
            st.image(crop_images[orig], caption=label, use_container_width=True)
            if st.button(label):
                st.session_state["selected_crop"] = orig
                st.session_state["lang_code"] = lang_code

    if "selected_crop" in st.session_state:
        crop = st.session_state["selected_crop"]
        lang_code = st.session_state["lang_code"]
        st.subheader(f"📝 Ask about: {translate(crop, lang_code)}")

        df = load_csv(crop)
        if df is None:
            st.error("❌ Dataset not found.")
            return

        vectorstore = build_store(df)

        user_q = st.text_input(translate("❓ Ask your question", lang_code))
        if user_q:
            with st.spinner("🧐 Thinking..."):
                translated_q = translate(user_q, "en")
                answer, source = handle_question(df, vectorstore, crop, translated_q)
                st.markdown(f"{translate(source, lang_code)}")
                st.write(translate(answer, lang_code))
  
if __name__ == "__main__":
    main()