import streamlit as st
import re
from langchain.memory import ConversationBufferMemory
from langchain.chains import ConversationalRetrievalChain
from langchain_openai import ChatOpenAI
from langchain_community.vectorstores import Pinecone as LangchainPinecone
from TOT import TreeofThoughts
from utils import sample_suggestions

class Chatbot:
    def __init__(self, vector_store: LangchainPinecone):
        self.memory = ConversationBufferMemory(memory_key="chat_history", return_messages=True, output_key="answer")
        self.qa_chain = ConversationalRetrievalChain.from_llm(
            llm=ChatOpenAI(model="gpt-4o", temperature=0.7, max_retries=2, max_tokens=16384),
            retriever=vector_store.as_retriever(
                search_type="similarity",
                search_kwargs={"k": 5, "filter": {"user_id": st.session_state.user_id}}
            ),
            memory=self.memory,
            return_source_documents=True,
            output_key="answer"
        )
        self.tot = TreeofThoughts(self, "BFS")

    def generate(self, prompt):
        """Generate a response using the OpenAI API."""
        try:
            response = self.qa_chain.invoke({"question": prompt})
            return response
        except Exception as e:
            st.error(f"Error generating response: {str(e)}")
            return "An error occurred while processing your request."

    def is_complex_query(self, query):
        """Determine if a query requires complex reasoning."""
        complex_keywords = [
            "plan", "strategy", "design", "solve", "optimize", "compare",
            "analyze", "evaluate", "recommend", "prioritize", "pros and cons",
            "advantage", "disadvantage", "trade-off", "decision", "choose",
            "best approach", "steps to", "how would you", "what if", "scenario"
        ]

        # Check for question complexity indicators
        complexity_indicators = [
            len(query.split()) > 20,  # Long queries are often complex
            query.count("?") > 1,     # Multiple questions
            any(keyword in query.lower() for keyword in complex_keywords)
        ]

        return any(complexity_indicators)

    def get_response(self, query: str) -> tuple:
        """Retrieve answer, sources, and suggested questions in a single LLM call."""
        try:
            with st.spinner("Searching documents and generating response..."):
                retrieved_docs = self.qa_chain.retriever.get_relevant_documents(query)
                context = "\n\n".join([doc.page_content for doc in retrieved_docs])
                prompt = f"""Strictly based on your context,Generate a detailed elaborated structured response to this query: "{query}"

                Structure your answer with these elements:

                1. Comprehensive Answer:
                - Provide a well-organized response using paragraphs and bullet points
                - Focus on factual accuracy and completeness
                - Avoid markdown formatting

                2. Follow with exactly 3 numbered follow-up questions using this exact format:

                Suggested Questions:
                1) [Concise question under 15 words]
                2) [Specific question based on context]
                3) [Relevant follow-up query]

                Follow these rules strictly:
                - Questions must be numbered with 1), 2), 3)
                - Questions should not repeat the original query
                - No section headers or titles
                - In case the context does not have any information, answer: "I don't know.."
                - Maintain exactly this output structure:
                [Full answer text...]

                Suggested Questions:
                1) ...
                2) ...
                3) ..."""

                response = self.generate(prompt)

                answer = response.get("answer", "No answer found.")
                source_docs = response.get("source_documents", [])

                # Extract source texts and metadata
                sources = set()
                source_texts = []
                for doc in source_docs:
                    if "source" in doc.metadata:
                        sources.add(doc.metadata["source"])
                        source_texts.append(doc.page_content)

                answer = answer.strip()
                suggested_questions = []

                # Pattern to find the suggested questions section
                pattern = r"\n?\n?Suggested Questions:"

                # Search for the pattern
                match = re.search(pattern, answer)
                if match:
                    # Split into answer and questions sections
                    answer_part = answer[:match.start()].strip()
                    questions_part = answer[match.end():].strip()

                    # Clean and extract individual questions
                    suggested_questions = [
                        q.strip().lstrip("0123456789).-*")
                        for q in questions_part.split("\n")
                        if q.strip()
                    ]

                    # Remove empty strings that might result from splitting
                    suggested_questions = [q for q in suggested_questions if q]

                    # Update clean_answer to exclude questions section
                    answer = answer_part
                else:
                    suggested_questions = []

                # Check if we need to use Tree of Thoughts for complex queries
                if st.session_state.toggle:
                    # context_complete = context+ f"{answer}"
                    # doc_response = self.qa_chain.invoke({"question": query})
                    # context = doc_response.get("answer", "")
                    # sources = {doc.metadata["source"] for doc in doc_response.get("source_documents", [])
                    #           if "source" in doc.metadata}

                    # Formulate a prompt that includes the document context
                    enhanced_query = f"""
                    Based on the following information from documents:

                    {answer}

                    Please address this question:
                    {query}
                    """
                    print("pqpq")
                    print(enhanced_query)
                    # Process with Tree of Thoughts
                    tot_response = self.tot.solve(enhanced_query, k=3, T=3, b=3, vth=0.5, context=answer)
                    answer = tot_response
                    print("complex")
                else:
                    print("simple")
                    
                if answer.startswith("I don't know") or answer.startswith("I don't have") or not sources:
                    return answer.strip(), (), (), sample_suggestions

                return answer.strip(), sources, source_texts, suggested_questions[:3]

        except Exception as e:
            return str(e), (), [], []