import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline
from peft import PeftModel
from langchain_community.llms import HuggingFacePipeline
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_community.vectorstores import FAISS
from langchain.text_splitter import CharacterTextSplitter
from langchain.docstore.document import Document
from langchain.chains import RetrievalQA

# Model and Tokenizer
model_name = "Meldashti/chatbot"
base_model = AutoModelForCausalLM.from_pretrained("unsloth/Llama-3.2-3B")
tokenizer = AutoTokenizer.from_pretrained("unsloth/Llama-3.2-3B")

# Merge PEFT weights with base model
model = PeftModel.from_pretrained(base_model, model_name)
model = model.merge_and_unload()

# Simplified pipeline with minimal parameters
generator = pipeline(
    "text-generation", 
    model=model, 
    tokenizer=tokenizer,
    max_new_tokens=50,  # Very low to test responsiveness
    do_sample=False
)

# LLM wrapper
llm = HuggingFacePipeline(pipeline=generator)

# Embeddings
embeddings = HuggingFaceEmbeddings(model_name="paraphrase-MiniLM-L3-v2")

# Sample documents (minimal)
documents = [
    Document(page_content="Example document about food industry caps"),
    Document(page_content="Information about manufacturing processes")
]

# Text splitting
text_splitter = CharacterTextSplitter(chunk_size=100, chunk_overlap=20)
split_documents = text_splitter.split_documents(documents)

# Vector store
vector_store = FAISS.from_documents(split_documents, embeddings)
retriever = vector_store.as_retriever(search_kwargs={"k": 2})

# Retrieval QA Chain
rag_chain = RetrievalQA.from_chain_type(
    llm=llm, 
    chain_type="stuff", 
    retriever=retriever
)

# Chat function with extensive logging
def chat(message, history):
    print(f"Received message: {message}")
    try:
        # Add timeout mechanism
        import timeout_decorator

        @timeout_decorator.timeout(10)  # 10 seconds timeout
        def get_response():
            response = rag_chain.invoke({"query": message})
            return str(response['result'])

        return get_response()
    except Exception as e:
        print(f"Error generating response: {type(e)}, {e}")
        return f"An error occurred: {str(e)}"

# Gradio interface
demo = gr.ChatInterface(chat, type="messages", autofocus=False)

# Launch
if __name__ == "__main__":
    demo.launch(debug=True)