from pydantic import BaseModel, Field
from langchain_aws.chat_models.bedrock import ChatBedrock
from langchain import hub
import logging
from typing import Literal
from langchain_core.output_parsers import StrOutputParser
from langchain_core.tools import tool
from langchain.agents import AgentExecutor, create_tool_calling_agent
from aind_data_access_api.document_db import MetadataDbClient
from pprint import pprint

logging.basicConfig(filename='agentic_graph.log', level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s', filemode="w")

GLOBAL_VARIABLE = 90

MODEL_ID = "anthropic.claude-3-sonnet-20240229-v1:0"
LLM = ChatBedrock(
    model_id= MODEL_ID,
    model_kwargs= {
        "temperature": 0
    }
)

#determining if entire database needs to be surveyed
class RouteQuery(BaseModel):
    """Route a user query to the most relevant datasource."""

    reasoning: str = Field(
        description="Give a justification for the chosen method",
    )

    datasource: Literal["vectorstore", "direct_database"] = Field(
        description="Given a user question choose to route it to the direct database or its vectorstore.",
    )

structured_llm_router = LLM.with_structured_output(RouteQuery)
router_prompt = hub.pull("eden19/query_rerouter")
datasource_router = router_prompt | structured_llm_router


# Queries that require surveying the entire database (like count based questions)
# credentials = DocumentDbSSHCredentials()
# credentials.database = "metadata_vector_index"
# credentials.collection = "curated_assets"
API_GATEWAY_HOST = "api.allenneuraldynamics-test.org"
DATABASE = "metadata_vector_index"
COLLECTION = "curated_assets"

docdb_api_client = MetadataDbClient(
   host=API_GATEWAY_HOST,
   database=DATABASE,
   collection=COLLECTION,
)

@tool
def aggregation_retrieval(agg_pipeline: list) -> list:
    """Given a MongoDB query and list of projections, this function retrieves and returns the 
    relevant information in the documents. 
    Use a project stage as the first stage to minimize the size of the queries before proceeding with the remaining steps.
    The input to $map must be an array not a string, avoid using it in the $project stage.

    Parameters
    ----------
    agg_pipeline
        MongoDB aggregation pipeline

    Returns
    -------
    list
        List of retrieved documents
    """

    result = docdb_api_client.aggregate_docdb_records(
        pipeline=agg_pipeline
    )
    return result
        
tools = [aggregation_retrieval]
db_prompt = hub.pull("eden19/entire_db_retrieval")

#llm_with_tools = LLM.bind_tools(tools)

db_surveyor_agent = create_tool_calling_agent(LLM, tools, db_prompt)
query_retriever = AgentExecutor(agent=db_surveyor_agent, tools=tools, return_intermediate_steps = True, verbose=False)
# pprint(query_retriever.invoke({'chat_history': [],"query": "What are the data asset names using mouse 675387", "agent_scratchpad":[]}))
# class retrieve_aggregation(BaseModel):
#     """List of results retrieved from mongodb database after applying pipeline generated by the model"""

#     mongo_db_query: str = Field(
#         description="Aggregation pipeline used to retrieve information from database"
#     )
#     mongo_db_results: list= Field(
#         description="List of information retrieved from the database after applying aggregation pipeline"
#     )


# aggregation_retriever = llm_with_tools.with_structured_output(retrieve_aggregation)

# aggregation_retriever_chain = db_prompt | aggregation_retriever
# print(aggregation_retriever_chain.invoke({'chat_history': [],"query": "how many documents are stored in the collection?", 'agent_scratchpad': []}))


# Processing query
class ProcessQuery(BaseModel):
    """Binary score to check whether query requires retrieval to be filtered with metadata information to achieve accurate results."""

    binary_score: str = Field(
        description="Query requires further filtering during retrieval process, 'yes' or 'no'"
    )

query_grader = LLM.with_structured_output(ProcessQuery)
query_grade_prompt = hub.pull("eden19/processquery")
query_grader = query_grade_prompt | query_grader
#print(query_grader.invoke({"query": "What is the genotype for mouse 675387?"}).binary_score)

# Generating appropriate filter
class FilterGenerator(BaseModel):
    """MongoDB filter to be applied before vector retrieval"""

    filter_query: dict = Field(description="MongoDB filter")
    #top_k: int = Field(description="Number of documents to retrieve from the database")

filter_prompt = hub.pull("eden19/filtergeneration")
filter_generator_llm = LLM.with_structured_output(FilterGenerator)

filter_generation_chain = filter_prompt | filter_generator_llm
#print(filter_generation_chain.invoke({"query": "What is the genotype for mouse 675387?"}).filter_query)

# Check if retrieved documents answer question
class RetrievalGrader(BaseModel):
    """Binary score to check whether retrieved documents are relevant to the question"""

    reasoning: str = Field(
        description="Give a reasoning as to what makes the document relevant for the chosen method",
    )

    binary_score: str = Field(
        description="Retrieved documents are relevant to the query, 'yes' or 'no'"
    )

    relevant_context: str = Field(
        description="Relevant pieces of context in document"
    )

retrieval_grader = LLM.with_structured_output(RetrievalGrader)
retrieval_grade_prompt = hub.pull("eden19/retrievalgrader")
doc_grader = retrieval_grade_prompt | retrieval_grader
# doc_grade = doc_grader.invoke({"query": question, "document": doc}).binary_score
# logging.info(f"Retrieved document matched query: {doc_grade}")

# Generating response to documents
answer_generation_prompt = hub.pull("eden19/answergeneration")
rag_chain = answer_generation_prompt | LLM | StrOutputParser()

db_answer_generation_prompt = hub.pull("eden19/db_answergeneration")
db_rag_chain = answer_generation_prompt | LLM | StrOutputParser()
# generation = rag_chain.invoke({"documents": doc, "query": question})
# logging.info(f"Final answer: {generation}")