import os
from langchain.chat_models import ChatOpenAI
from langchain.chains import RetrievalQA
from langchain.memory import ConversationBufferMemory
from langchain.prompts import PromptTemplate
from langchain.agents import initialize_agent, AgentType
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.embeddings.huggingface import HuggingFaceEmbeddings
from langchain.vectorstores import FAISS
from langchain.tools import Tool
from langchain.tools import DuckDuckGoSearchRun
from langchain_core.documents import Document


from dotenv import load_dotenv
import os

load_dotenv()  

apikey = os.getenv("MISTRAL_API_KEY")

llm = ChatOpenAI(
    openai_api_key=apikey,
    openai_api_base="https://api.mistral.ai/v1",
    model="mistral-medium"
)

prompt_template = PromptTemplate(
    input_variables=["user_prompt"],
    template="""You are a retriever agent tasked with creating an efficient search small query to retrieve academic papers from arxiv relevant to a user’s request. Based on the user’s input prompt, generate a concise and precise search query (a string of keywords or phrases) that will be used by the function `retrieve_and_extract_papers(query: str, max_papers: int = 3) -> str` to fetch up to 3 relevant papers. The query should focus on key concepts, avoid ambiguity, and prioritize relevance to ensure the extracted text is suitable for summarization.

User Input Prompt: {user_prompt}

Instructions:
1. Identify the core concepts, topics, or questions in the user prompt.
2. Formulate a search query using relevant keywords or short phrases.
3. Exclude overly broad or irrelevant terms to improve precision.
4. Output only the search query as a string.

Example:
- User Prompt: "Recent advancements in large language models for natural language processing"
- Search Query: "large language models NLP advancements"

Generate the search query for the provided user prompt.
"""
)


search_tool = DuckDuckGoSearchRun()
tools = [
    Tool(
        name="WebSearch",
        func=search_tool.run,
        description="Useful for fetching up-to-date healthcare information from the web.Use it rarely"
    ),
]

retriever_agent = initialize_agent(
    tools=tools,
    agent_type=AgentType.CONVERSATIONAL_REACT_DESCRIPTION,
    llm=llm,
    verbose=True,
    prompt=prompt_template
)



import arxiv
import pdfplumber
import requests
import os
from typing import List
embedding_model = HuggingFaceEmbeddings(model_name="BAAI/bge-base-en-v1.5")
def retrieve_and_extract_papers(query: str, max_papers: int = 3) -> str:
    """
    Retrieves research papers from arXiv, downloads PDFs, extracts text,
    and returns a single string with papers separated by '---pprN: Title'.
    
    Args:
        query: Search query (e.g., "diffusion models 2024")
        max_papers: Number of papers to retrieve (default: 3)
    
    Returns:
        Single string with extracted text from PDFs, separated by '---pprN: Title'
    """
    # Initialize arXiv client
    client = arxiv.Client()
    search = arxiv.Search(
        query=query,
        max_results=max_papers,
        sort_by=arxiv.SortCriterion.Relevance
    )
    
    papers = list(client.results(search))
    if not papers:
        return "No papers found for the query."
    
    # Create temporary directory for PDFs
    temp_dir = "temp_papers"
    os.makedirs(temp_dir, exist_ok=True)
    
    # Process each paper and collect extracted text
    formatted_texts = []
    for i, paper in enumerate(papers, 1):
        try:
            # Download PDF
            pdf_url = paper.pdf_url
            response = requests.get(pdf_url)
            pdf_path = os.path.join(temp_dir, f"paper_{i}.pdf")
            
            with open(pdf_path, 'wb') as f:
                f.write(response.content)
            
            # Extract text from PDF
            text = ""
            with pdfplumber.open(pdf_path) as pdf:
                for page in pdf.pages:
                    page_text = page.extract_text()
                    if page_text:
                        text += page_text + "\n"
            
            # Append formatted text with separator including paper title
            separator = f"---ppr{i}: {paper.title}"
            formatted_texts.append(f"{separator}\n{text}")
            
            # Clean up: Remove the downloaded PDF
            os.remove(pdf_path)
            
        except Exception as e:
            print(f"Error processing paper {i} ({paper.title}): {e}")
            formatted_texts.append(f"---ppr{i}: {paper.title}\nError: Could not process paper.")
    
    # Clean up: Remove temporary directory if empty
    if os.path.exists(temp_dir) and not os.listdir(temp_dir):
        os.rmdir(temp_dir)
    
    # Join all texts into a single string
    return "\n".join(formatted_texts)
pdftext = retrieve_and_extract_papers(retriever_output)
text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=50)
document = Document(page_content=pdftext)
documents = text_splitter.split_documents([document])
faiss_index = FAISS.from_documents(documents, embedding_model)
faiss_index.save_local("faiss_index")
retriever = faiss_index.as_retriever(search_kwargs={"k": 3})