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

load_dotenv()  

apikey = os.getenv("MISTRAL_API_KEY")

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


embedding_model = HuggingFaceEmbeddings(model_name="BAAI/bge-base-en-v1.5")
faiss_index = FAISS.load_local("faiss_index", embedding_model, allow_dangerous_deserialization=True)
retriever = faiss_index.as_retriever(search_kwargs={"k": 3})

tools = [
    Tool(
        name="FAISSRetriever",
        func=lambda query: "\n\n".join([doc.page_content for doc in retriever.invoke(query)]),
        description="Fetches the top 3 relevant document chunks from a FAISS index containing academic paper content based on a query. Use for scholarly text analysis."
    ),
    Tool(
        name="WebSearch",
        func=DuckDuckGoSearchRun().run,
        description="Fetches up-to-date information from the web. Use for verifying claims or finding additional context."
    ),
]


prompt_template = PromptTemplate(
    input_variables=["user_query", "summary"],
    template="""You are a critique agent tasked with strictly evaluating a summary of academic papers based on a user query. Using the provided user query and summary, analyze the content for accuracy, completeness, and biases. Use the FAISSRetriever and WebSearch tools to cross-reference information and validate claims. Provide a detailed critique, including recommendations, identified biases, and a relevance rating (1–5, where 5 is highly relevant and accurate).

User Query: {user_query}
Summary: {summary}

Instructions:
1. Use the FAISSRetriever tool to fetch document chunks related to the user query for additional context from academic papers.
2. Use the WebSearch tool to verify claims in the summary or find recent developments.
3. Evaluate the summary for:
   - Accuracy: Are claims supported by evidence from the tools?
   - Completeness: Does it cover key aspects of the user query?
   - Biases: Identify methodological, dataset, author, or other biases (e.g., overemphasis on positive results, lack of diverse perspectives).
4. Provide recommendations to improve the summary (e.g., additional topics, clearer explanations).
5. Assign a relevance rating (1–5) based on how well the summary addresses the user query and its reliability.
6. Be strict in identifying biases and unsupported claims.
7. Output a structured response with sections: Critique, Recommendations, Biases, Relevance Rating.

Example:
- User Query: "Advancements in diffusion models for image generation"
- Summary: "Diffusion models outperform GANs in image quality."
- Output:
  Critique: The summary claims diffusion models outperform GANs but lacks evidence or metrics.
  Recommendations: Include specific metrics (e.g., FID scores) and compare computational costs.
  Biases: Potential bias toward diffusion models; ignores GANs’ strengths in training speed.
  Relevance Rating: 2/5 (lacks depth and evidence).

Generate the critique for the provided user query and summary.
"""
)

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