from huggingface_hub import HfApi
import tempfile
import streamlit as st
import pandas as pd
import io, os
import subprocess
import pdfplumber
from lxml import etree
from bs4 import BeautifulSoup
from PyPDF2 import PdfReader
from langchain_community.vectorstores import FAISS
from langchain.embeddings.base import Embeddings
from langchain_openai import ChatOpenAI
from langchain.agents import initialize_agent, AgentType
from langchain.agents import Tool
from langchain.memory import ConversationBufferMemory
from langchain.text_splitter import CharacterTextSplitter
from dotenv import load_dotenv
import google.generativeai as genai
from typing import List
from langchain_core.language_models import BaseLanguageModel
from langchain_core.runnables import Runnable
import google.generativeai as genai
from datetime import datetime

load_dotenv()

def load_environment():
    # Ensure HF_TOKEN is available
    if "HUGGINGFACEHUB_API_TOKEN" not in os.environ and "HF_TOKEN" in os.environ:
        os.environ["HUGGINGFACEHUB_API_TOKEN"] = os.environ["HF_TOKEN"] 
    if "GOOGLE_API_KEY" not in os.environ:
        raise ValueError("GOOGLE_API_KEY not found in environment variables.")
    genai.configure(api_key=st.secrets["GOOGLE_API_KEY"])

from keybert import KeyBERT
from sentence_transformers import CrossEncoder
from sentence_transformers import SentenceTransformer

class GeminiLLM(Runnable):
    def __init__(self, model_name="models/gemini-1.5-pro-latest", api_key=None):
        self.api_key = api_key or os.environ["GOOGLE_API_KEY"]
        if not self.api_key:
            raise ValueError("GOOGLE_API_KEY not found.")
        genai.configure(api_key=self.api_key)
        self.model = genai.GenerativeModel(model_name)
        
    def _call(self, prompt: str, stop=None) -> str:
        response = self.model.generate_content(prompt)
        return response.text

    @property
    def _llm_type(self) -> str:
        return "custom_gemini"
    
    def invoke(self, input, config=None):
        response = self.model.generate_content(input)
        return response.text.strip()

class GeminiEmbeddings(Embeddings):
    def __init__(self, model_name="models/embedding-001", api_key=None):
        api_key = "AIzaSyBIfGJRoet_wzzYXIiWXxStkIigEOzSR2o"
        if not api_key:
            raise ValueError("GOOGLE_API_KEY not found in environment variables.")
        os.environ["GOOGLE_API_KEY"] = api_key
        genai.configure(api_key=api_key)
        self.model_name = model_name

    def embed_documents(self, texts: List[str]) -> List[List[float]]:
        return [
            genai.embed_content(
                model=self.model_name,
                content=text,
                task_type="retrieval_document"
            )["embedding"]
            for text in texts
        ]

    def embed_query(self, text: str) -> List[float]:
        return genai.embed_content(
            model=self.model_name,
            content=text,
            task_type="retrieval_query"
        )["embedding"]


vectorstore_global = None

# Initialize feedback list
if "feedback_log" not in st.session_state:
    st.session_state.feedback_log = []


def preload_modtran_document():
    global vectorstore_global
    embeddings = GeminiEmbeddings()
    st.session_state.vectorstore = FAISS.load_local("modtran_vectorstore", embeddings, allow_dangerous_deserialization=True)
    set_global_vectorstore(st.session_state.vectorstore)
    st.session_state.chat_ready = True

def convert_pdf_to_xml(pdf_file, xml_path):
    os.makedirs("temp", exist_ok=True)
    pdf_path = os.path.join("temp", pdf_file.name)
    with open(pdf_path, 'wb') as f:
        f.write(pdf_file.getbuffer())
    subprocess.run(["pdftohtml", "-xml", pdf_path, xml_path], check=True)
    return xml_path

def extract_text_from_xml(xml_path, document_name):
    tree = etree.parse(xml_path)
    text_chunks = []
    for page in tree.xpath("//page"):
        page_num = int(page.get("number", 0))
        texts = [text.text for text in page.xpath('.//text') if text.text]
        combined_text = '\n'.join(texts)
        text_chunks.append({"text": combined_text, "page": page_num, "document": document_name})
    return text_chunks

def extract_text_from_pdf(pdf_file, document_name):
    text_chunks = []
    with pdfplumber.open(pdf_file) as pdf:
        for i, page in enumerate(pdf.pages):
            text = page.extract_text()
            if text:
                text_chunks.append({"text": text, "page": i + 1, "document": document_name})
    return text_chunks

def get_uploaded_text(uploaded_files):
    raw_text = []
    for uploaded_file in uploaded_files:
        document_name = uploaded_file.name
        if document_name.endswith(".pdf"):
            text_chunks = extract_text_from_pdf(uploaded_file, document_name)
            raw_text.extend(text_chunks)
        elif uploaded_file.name.endswith((".html", ".htm")):
            soup = BeautifulSoup(uploaded_file.getvalue(), 'lxml')
            raw_text.append({"text": soup.get_text(), "page": None, "document": document_name})
        elif uploaded_file.name.endswith((".txt")):
            content = uploaded_file.getvalue().decode("utf-8")
            raw_text.append({"text": content, "page": None, "document": document_name})
    return raw_text

def get_text_chunks(raw_text):
    splitter = CharacterTextSplitter(separator='\n', chunk_size=500, chunk_overlap=100)
    final_chunks = []
    for chunk in raw_text:
        for split_text in splitter.split_text(chunk["text"]):
            final_chunks.append({"text": split_text, "page": chunk["page"], "document": chunk["document"]})
    return final_chunks

def get_vectorstore(text_chunks):
    if not text_chunks:
        raise ValueError("text_chunks is empty. Cannot initialize FAISS vectorstore.")

    embeddings = GeminiEmbeddings()
    texts = [chunk["text"] for chunk in text_chunks]
    metadatas = [{"page": chunk["page"], "document": chunk["document"]} for chunk in text_chunks]
    
    return FAISS.from_texts(texts, embedding=embeddings, metadatas=metadatas)

def set_global_vectorstore(vectorstore):
    global vectorstore_global
    vectorstore_global = vectorstore

kw_model = None

def get_kw_model():
    global kw_model
    if kw_model is None:
        # Load sentence transformer with HF token explicitly
        model = SentenceTransformer(
            'sentence-transformers/all-MiniLM-L6-v2',
            use_auth_token=os.environ.get("HF_TOKEN")
        )
        kw_model = KeyBERT(model=model)
    return kw_model

def self_reasoning(query, context):
    print("🧪 self_reasoning received context of length:", len(context))
    llm = GeminiLLM()
    reasoning_prompt = f"""
    You are an AI assistant that analyzes the context provided to answer the user's query comprehensively and clearly. 
    Answer in a concise, factual way using the terminology from the context. Avoid extra explanation unless explicitly asked.
    YOU MUST mention the page number. 
    ### Example 1:
    **Question:** What is the purpose of the MODTRAN GUI?
    **Context:**
    [Page 10 of the docuemnt] The MODTRAN GUI helps users set parameters and visualize the model's output.
    **Answer:** The MODTRAN GUI assists users in parameter setup and output visualization. You can find the answer at Page 10 of the document provided.

    ### Example 2:
    **Question:** How do you run MODTRAN on Linux? Answer with page number. 
    **Context:**
    [Page 15 of the docuemnt] On Linux systems, MODTRAN can be run using the `mod6c` binary via terminal.
    **Answer:** Use the `mod6c` binary via terminal. (Page 15 of the document)

    ### Now answer:
    **Question:** {query}
    **Context:**
    {context}

    **Answer:**
    """
    try:
        result = llm._call(reasoning_prompt)
        print("✅ Gemini returned a result.")
        return result
    except Exception as e:
        print("❌ Error in self_reasoning:", e)
        return f"⚠️ Gemini failed: {e}"

def faiss_search_with_keywords(query):
    global vectorstore_global
    if vectorstore_global is None:
        raise ValueError("FAISS vectorstore is not initialized.")
    kw_model = get_kw_model()
    keywords = kw_model.extract_keywords(query, keyphrase_ngram_range=(1,2), stop_words='english', top_n=5)
    refined_query = " ".join([keyword[0] for keyword in keywords])
    retriever = vectorstore_global.as_retriever(search_kwargs={"k": 13})
    docs = retriever.get_relevant_documents(refined_query)
    context= '\n\n'.join([f"[Page {doc.metadata.get('page', 'Unknown')}] {doc.page_content}" for doc in docs])
    return self_reasoning(query, context)

def get_reranker():
    global reranker
    if reranker is None:
        reranker = CrossEncoder('cross-encoder/ms-marco-MiniLM-L-6-v2')
    return reranker
    
def faiss_search_with_reasoning(query):
    global vectorstore_global
    if vectorstore_global is None:
        raise ValueError("FAISS vectorstore is not initialized.")
    reranker = get_reranker()
    retriever = vectorstore_global.as_retriever(search_kwargs={"k": 13})
    docs = retriever.get_relevant_documents(query)
    pairs = [(query, doc.page_content) for doc in docs]
    scores = reranker.predict(pairs)
    reranked_docs = sorted(zip(scores, docs), key=lambda x: x[0], reverse=True)
    top_docs = [doc for _, doc in reranked_docs[:5]]
    context = '\n\n'.join([f"[Page {doc.metadata.get('page', 'Unknown')}] {doc.page_content.strip()}" for doc in top_docs])
    return self_reasoning(query, context)

faiss_keyword_tool = Tool(
    name="FAISS Keyword Search",
    func=faiss_search_with_keywords,
    description="Searches FAISS with a keyword-based approach to retrieve context."
)

faiss_reasoning_tool = Tool(
    name="FAISS Reasoning Search",
    func=faiss_search_with_reasoning,
    description="Searches FAISS with detailed reasoning to retrieve context."
)

def initialize_chatbot_agent():
    llm = GeminiLLM() 
    memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True)
    tools = [faiss_keyword_tool, faiss_reasoning_tool]
    agent = initialize_agent(
        tools=tools,
        llm=llm,
        agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
        memory=memory,
        verbose=False,
        handle_parsing_errors=True
    )
    return agent

def handle_user_query(query):
    try:
        global vectorstore_global
        if vectorstore_global is None:
            raise ValueError("Vectorstore is not initialized.")

        print("🔍 Starting handle_user_query with:", query)

        if "how" in query.lower():
            print("🧠 Routing to: faiss_search_with_reasoning")
            context = faiss_search_with_reasoning(query)
        else:
            print("🧠 Routing to: faiss_search_with_keywords")
            context = faiss_search_with_keywords(query)

        print("📚 Context length:", len(context))
        print("✍️ Calling self_reasoning...")

        answer = self_reasoning(query, context)

        print("✅ Answer generated.")
        return answer

    except Exception as e:
        print("❌ Error in handle_user_query:", e)
        return f"⚠️ Error: {e}"


def save_feedback_to_huggingface():
    try:
        if not st.session_state.feedback_log:
            return  # No feedback yet

        feedback_df = pd.DataFrame(st.session_state.feedback_log)
        now = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
        filename = f"feedback_{now}.csv"

        with tempfile.TemporaryDirectory() as tmpdir:
            filepath = os.path.join(tmpdir, filename)
            feedback_df.to_csv(filepath, index=False)

            token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGINGFACEHUB_API_TOKEN")
            if not token:
                raise ValueError("HF token not found.")

            api = HfApi(token=token)
            api.upload_file(
                path_or_fileobj=filepath,
                path_in_repo=filename,
                repo_id="ZarinT/chatbot-feedback",
                repo_type="dataset"
            )

        st.session_state.feedback_log.clear()
        print("Feedback uploaded successfully.")

    except Exception as e:
        print("Error uploading feedback:", e)
    
def main():
    load_environment()

    if "chat_ready" not in st.session_state:
        st.session_state.chat_ready = False
    if "chat_history" not in st.session_state:
        st.session_state.chat_history = []
    if "vectorstore" not in st.session_state:
        st.session_state.vectorstore = None

    st.header("Chat with MODTRAN Documents 📄")

    # Preload the document once when app starts
    if not st.session_state.chat_ready:
        with st.spinner("Loading MODTRAN document..."):
            preload_modtran_document()
            st.session_state.agent = initialize_chatbot_agent()
            st.session_state.chat_ready = True
            st.success("MODTRAN User Manual loaded successfully!")

    user_question = st.text_input("Ask your question:", key="user_input")

    if st.session_state.chat_ready and user_question:
        with st.spinner("Generating answer..."):
            try:
                set_global_vectorstore(st.session_state.vectorstore)
                response = handle_user_query(user_question)
            except Exception as e:
                response = f"⚠️ Something went wrong: {e}"

        # ✅ Show the response
        st.markdown(f"**Answer:**\n\n{response}")
        st.write(response)

    st.session_state.chat_history.append({"user": user_question, "bot": response})

        # After bot responds, ask for feedback
        with st.form(key=f"feedback_form_{len(st.session_state.chat_history)}"):
            rating = st.radio(
                "Rate this response:",
                options=["1", "2", "3", "4", "5"],
                key=f"rating_{len(st.session_state.chat_history)}",
                horizontal=True
            )
            submitted = st.form_submit_button("Submit Rating")
            if submitted:
                feedback = {
                    "question": user_question,
                    "response": response,
                    "rating": rating,
                    "timestamp": datetime.datetime.now().isoformat()
                }
                st.session_state.feedback_log.append(feedback)

                # If 5 feedbacks collected, auto upload to Hugging Face
                if len(st.session_state.feedback_log) >= 5:
                    save_feedback_to_huggingface()

        # Clear input box automatically
        st.rerun()

if __name__ == "__main__":
    load_environment()
    main()