from transformers import AutoTokenizer, AutoModelForCausalLM
from langchain_core.output_parsers import BaseOutputParser
from langchain.prompts import PromptTemplate
from langchain_core.runnables import RunnablePassthrough, RunnableMap
from langchain.schema import Generation
import torch

# Load your tokenizer and model
tokenizer = AutoTokenizer.from_pretrained("NLPGenius/GovGPT-llama3")
model = AutoModelForCausalLM.from_pretrained("NLPGenius/GovGPT-llama3", 
                                             torch_dtype=torch.bfloat16, 
                                             device_map="auto",)


import transformers
from transformers import pipeline
llm_chain = transformers.pipeline(
    model=model, tokenizer=tokenizer,
    return_full_text=True,  # langchain expects the full text
    task='text-generation',
    # we pass model parameters here too
    temperature=0.3,  # 'randomness' of outputs, 0.0 is the min and 1.0 the max
    do_sample=True,
    max_length=2000,  # mex number of tokens to generate in the output
    truncation=True,
    repetition_penalty=1.1  # without this output begins repeating
)

from langchain.llms import HuggingFacePipeline

llm = HuggingFacePipeline(pipeline=llm_chain)

from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.embeddings import HuggingFaceEmbeddings
from langchain.vectorstores import FAISS

# Define the chunking and embedding strategy
text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
embedding_model = HuggingFaceEmbeddings()  # Adjust if you have a specific embedding model

import os
import rarfile
from langchain_community.document_loaders import PyPDFLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.vectorstores import FAISS
from langchain.embeddings import HuggingFaceEmbeddings
from langchain.prompts import PromptTemplate
from langchain.chains import LLMChain
from langchain_core.runnables import RunnablePassthrough
from langchain_core.output_parsers import StrOutputParser

# Function to create FAISS index
def create_faiss_index(docs):
    texts = text_splitter.split_documents(docs)  # Ensure `text_splitter` is defined in your script
    vector_store = FAISS.from_documents(texts, embedding_model)
    return vector_store

# Function to extract files from .rar archive
def extract_rar_files(rar_path, extract_to="extracted_files"):
    with rarfile.RarFile(rar_path) as rf:
        rf.extractall(extract_to)
    return extract_to

# Function to fetch and process PDF files
def fetch_and_process_pdfs(folder_path):
    pdf_files = [
        os.path.join(folder_path, file) for file in os.listdir(folder_path) if file.endswith(".pdf")
    ]
    
    all_documents = []
    for pdf_file in pdf_files:
        loader = PyPDFLoader(pdf_file)
        documents = loader.load()
        all_documents.extend(documents)
    return all_documents

# Function to create RAG pipeline
def create_rag_pipeline(rar_path):
    # Extract .rar file
    extracted_folder = extract_rar_files(rar_path)

    # Fetch and process PDF files
    documents = fetch_and_process_pdfs(extracted_folder)

    # Create vector store
    vector_store = create_faiss_index(documents)
    retriever = vector_store.as_retriever()

    return retriever

# Usage example
rar_path = "Rules folder.rar"
retriever = create_rag_pipeline(rar_path)



# Define a custom output parser to extract only the answer
class AnswerOutputParser(BaseOutputParser):
    def parse(self, text: str) -> str:
        # Extract everything after "Answer:"
        if "Answer:" in text:
            return text.split("Answer:")[-1].strip()
        return "I don't know."

# Function to perform model inference
def model_inference(retriever, question):
    # Define prompt template
    prompt = PromptTemplate(
        input_variables=["context", "question"],
        template="""\
        You are a helpful assistant. Answer the query accurately by focusing solely on the most relevant parts of the provided context.

        1. Identify and use only the sections of the context directly related to the query, ignoring unrelated or extraneous information.
        2. If the query cannot be explicitly answered using the relevant parts of the context, respond with: "The answer is not found in the context provided."
        3. If you are unsure or do not know the answer, respond with: "I don't know."
        4. Do not add information, interpretations, or assumptions beyond what is explicitly stated in the context.
        5. Ensure the response is concise, avoids redundancy, and directly addresses the query.

         Context:
        {context}
        
        Query:
        {question}
        
        Answer:
        """
    )


    rag_chain = (
            {"context": retriever, "question": RunnablePassthrough()}
            | prompt
            | llm
            | AnswerOutputParser() 
        )
    response = rag_chain.invoke(question)
    print("Bot response: ",response)
    return rag_chain

#response = model_inference(retriever, "Why was a new agriculture policy needed in Khyber Pakhtunkhwa?")



from flask import Flask, request, jsonify
from flask_cors import CORS

# Initialize Flask app
app = Flask(__name__)
CORS(app, supports_credentials=True, allow_headers=["Content-Type"])

# Mock chatbot function (replace with your actual pipeline)
def chatbot_response(query):
    try:
        # Call your actual RAG pipeline here
        response = model_inference(retriever, query)
        return response
    except Exception as e:
        return f"Error: {str(e)}"

# Route for chatbot responses
@app.route('/', methods=['GET','POST'])
def chat():
    
    # Try to get JSON data
    if request.content_type == 'application/json':
        data = request.get_json(silent=True)
        user_query = data.get("query", "").strip() if data else ""
    else:
        # Fallback to form-data
        user_query = request.form.get("query", "").strip()
    
    if not user_query:
        return jsonify({"error": "Please provide a valid query."}), 400

    # Get chatbot response
    bot_reply = chatbot_response(user_query)
    return jsonify({"response": bot_reply})

# Run the Flask app
if __name__ == '__main__':
    app.run()