import os
import glob
import uuid
from langchain.schema import Document
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.document_loaders import TextLoader
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.chat_history import BaseChatMessageHistory
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain.chains import create_history_aware_retriever, create_retrieval_chain
from langchain.chains.combine_documents import create_stuff_documents_chain
from langchain_community.chat_message_histories import ChatMessageHistory
from langchain_core.chat_history import BaseChatMessageHistory
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.runnables.history import RunnableWithMessageHistory
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_chroma import Chroma
import time

from .doc_tracker import DocumentTracker
from .ollama_wrapper import OllamaWrapper


class RAG:
    def __init__(
        self,
        docs_dir,
        db_dir,
        tracker_file,
        ollama_url="127.0.0.1",
        embedding_model="nomic-embed-text",
        llm_model="llama3.2",
    ):
        self.docs_dir = docs_dir
        self.db_dir = db_dir
        self.tracker = DocumentTracker(tracker_file)
        self.ollama = OllamaWrapper(
            ollama_url, embedding_model=embedding_model, llm_model=llm_model
        )
        self.session_id = time.time()

    def process_documents(self, files, text_splitter):
        """Process document files into chunks with tracking metadata"""
        all_chunks = []
        file_chunk_map = {}

        for file_path in files:
            # Load the document
            loader = TextLoader(file_path)
            documents = loader.load()

            # Add source metadata
            for doc in documents:
                doc.metadata["source"] = file_path
                doc.metadata["source_id"] = file_path  # For easier identification

            # Split the document
            chunks = text_splitter.split_documents(documents)

            # Generate and track chunk IDs
            chunk_ids = []
            for chunk in chunks:
                chunk_id = str(uuid.uuid4())
                chunk.metadata["chunk_id"] = chunk_id
                chunk_ids.append(chunk_id)

            # Store chunk mappings if tracker is provided
            if self.tracker:
                self.tracker.update_chunk_mappings(file_path, chunk_ids)

            file_chunk_map[file_path] = chunks
            all_chunks.extend(chunks)

        return all_chunks, file_chunk_map

    def create_vector_db(self, extensions=None, force_refresh=False):
        """
        Create or update a vector database, with complete handling of changes.

        Args:
            extensions (list[str], optional): List of file extensions to include.
                                           If None, defaults to common programming languages.
            force_refresh (bool): Whether to force a complete refresh of the vector database.
        """
        # Set default extensions for common programming languages if none provided
        if extensions is None:
            extensions = default_file_extensions()
        elif isinstance(extensions, str):
            extensions = [extensions]

        # Text splitter configuration
        text_splitter = RecursiveCharacterTextSplitter(
            chunk_size=1000, chunk_overlap=200
        )

        # Create embeddings
        embeddings = self.ollama.embeddings

        # Load or create vector store
        if os.path.exists(self.db_dir) and not force_refresh:
            print("Loading existing vector store")
            vectorstore = Chroma(
                persist_directory=self.db_dir, embedding_function=embeddings
            )
        else:
            print("Creating new vector store")
            vectorstore = Chroma(
                persist_directory=self.db_dir, embedding_function=embeddings
            )

        # Find all files that match the extensions
        all_files = set()
        for ext in extensions:
            files = glob.glob(os.path.join(self.docs_dir, f"**/*{ext}"), recursive=True)
            all_files.update(files)

        if not all_files:
            print("No documents found to process")
            return vectorstore

        # Process all files
        all_documents = []
        for file_path in all_files:
            try:
                with open(file_path, "r", encoding="utf-8") as f:
                    content = f.read()
                    doc = Document(
                        page_content=content,
                        metadata={
                            "source": os.path.abspath(file_path),
                            "source_id": os.path.abspath(file_path),
                        },
                    )
                    all_documents.append(doc)
                    print(f"Successfully loaded {file_path}")
            except Exception as e:
                print(f"Error loading file {file_path}: {str(e)}")

        if not all_documents:
            print("No documents could be loaded")
            return vectorstore

        # Split documents
        chunks = text_splitter.split_documents(all_documents)

        if force_refresh:
            # Create new vector store from scratch
            vectorstore = Chroma.from_documents(
                documents=chunks, embedding=embeddings, persist_directory=self.db_dir
            )
        else:
            # Update existing vector store
            # First, get all existing document IDs
            existing_docs = vectorstore._collection.get()
            existing_sources = {m["source"] for m in existing_docs["metadatas"]}

            # Find new documents
            new_chunks = [
                chunk
                for chunk in chunks
                if chunk.metadata["source"] not in existing_sources
            ]

            if new_chunks:
                print(f"Adding {len(new_chunks)} new chunks to vector store")
                vectorstore.add_documents(new_chunks)

        print(f"Vector store updated with {len(chunks)} total chunks")
        return vectorstore

    def setup_rag(self):
        """Set up the RAG pipeline"""
        # Create vector store
        vectorstore = Chroma(
            persist_directory=self.db_dir, embedding_function=self.ollama.embeddings
        )

        # Create retriever
        retriever = vectorstore.as_retriever(
            search_type="similarity", search_kwargs={"k": 4}
        )

        # Create chat history buffer
        self.chat_history = []

        ### Contextualize question ###
        contextualize_q_system_prompt = (
            "Given a chat history and the latest user question "
            "which might reference context in the chat history, "
            "formulate a standalone question which can be understood "
            "without the chat history. Do NOT answer the question, "
            "just reformulate it if needed and otherwise return it as is."
        )
        contextualize_q_prompt = ChatPromptTemplate.from_messages(
            [
                ("system", contextualize_q_system_prompt),
                MessagesPlaceholder("chat_history"),
                ("human", "{input}"),
            ]
        )
        history_aware_retriever = create_history_aware_retriever(
            self.ollama.llm, retriever, contextualize_q_prompt
        )

        ### Answer question ###
        system_prompt = (
            "You are an expert at analyzing code and documentation. "
            "Use the following pieces of context to answer the question at the end."
            "If you don't know the answer, just say that you don't know, "
            "don't try to make up an answer"
            "\n\n"
            "{context}\n\n"
            "Answer in a clear and concise manner. "
            "If you're referring to code, use markdown formatting."
        )
        qa_prompt = ChatPromptTemplate.from_messages(
            [
                ("system", system_prompt),
                MessagesPlaceholder("chat_history"),
                ("human", "{input}"),
            ]
        )
        question_answer_chain = create_stuff_documents_chain(self.ollama.llm, qa_prompt)

        rag_chain = create_retrieval_chain(
            history_aware_retriever, question_answer_chain
        )

        ### Statefully manage chat history ###
        self.store = {}

        def get_session_history(session_id: str) -> BaseChatMessageHistory:
            if session_id not in self.store:
                self.store[session_id] = ChatMessageHistory()
            return self.store[session_id]

        return RunnableWithMessageHistory(
            rag_chain,
            get_session_history,
            input_messages_key="input",
            history_messages_key="chat_history",
            output_messages_key="answer",
        )

    def query_rag(self, rag_chain, query, stream=False):
        """
        Query the RAG system

        Args:
            rag_chain: The RAG chain to use
            query: Query string
            stream: If True, stream the response
        """
        if stream:
            response = rag_chain.stream(
                {"input": query},
                config={"configurable": {"session_id": self.session_id}},
            )
            # Store in chat history after getting full response
            full_response = ""
            for chunk in response:
                # Extract only the LLM answer chunk
                if "answer" in chunk:
                    chunk_text = (
                        chunk["answer"].content
                        if hasattr(chunk["answer"], "content")
                        else str(chunk["answer"])
                    )
                    full_response += chunk_text
                    yield chunk_text
        else:
            response = rag_chain.invoke(
                {"input": query},
                config={"configurable": {"session_id": self.session_id}},
            )
            # Extract only the LLM answer from the response
            result = (
                response["answer"].content
                if hasattr(response["answer"], "content")
                else str(response["answer"])
            )
            return result


def default_file_extensions():
    """Return the default file extensions representing common plain text and code files"""
    return [
        # Python
        ".py",
        # C/C++
        ".c",
        ".cpp",
        ".cc",
        ".cxx",
        ".h",
        ".hpp",
        ".hxx",
        ".inc",
        ".inl",
        ".ipp",
        # Rust
        ".rs",
        # Java
        ".java",
        ".kt",
        ".kts",
        ".groovy",
        # Web
        ".html",
        ".htm",
        ".css",
        ".js",
        ".jsx",
        ".ts",
        ".tsx",
        ".vue",
        ".svelte",
        # Fortran
        ".f",
        ".for",
        ".f90",
        ".f95",
        ".f03",
        ".f08",
        # Go
        ".go",
        ".mod",
        # Ruby
        ".rb",
        ".rbw",
        ".rake",
        ".gemspec",
        # PHP
        ".php",
        ".phtml",
        ".php3",
        ".php4",
        ".php5",
        ".phps",
        # C#
        ".cs",
        ".csx",
        ".vb",
        # Swift
        ".swift",
        ".swiftmodule",
        # Shell/Scripts
        ".sh",
        ".bash",
        ".zsh",
        ".fish",
        ".ps1",
        ".bat",
        ".cmd",
        # Scala
        ".scala",
        ".sc",
        # Haskell
        ".hs",
        ".lhs",
        ".hsc",
        # Lua
        ".lua",
        ".luac",
        # R
        ".r",
        ".rmd",
        ".rds",
        # Perl
        ".pl",
        ".pm",
        ".t",
        # Documentation
        ".txt",
        ".md",
        ".rst",
        ".adoc",
        ".wiki",
        # Build/Config
        ".toml",
        ".yaml",
        ".yml",
        ".json",
        ".xml",
        ".ini",
        ".conf",
        ".cfg",
        # SQL
        ".sql",
        ".mysql",
        ".pgsql",
        ".sqlite",
        # Lisp family
        ".lisp",
        ".cl",
        ".el",
        ".clj",
        ".cljc",
        ".cljs",
        ".edn",
    ]
